ZacharyZcR
cbd599024e
* Fuse CUDA expert MLP execution * Group CUDA expert transfers by device * Instrument grouped CUDA expert execution * Bound grouped CUDA decode scratch * Execute expert groups across GPUs in parallel * Release host backing for multi-GPU experts * Define quality-preserving memory policies * Overlap cold expert loading with resident compute * Adapt expert placement with session LFRU * Fuse q4 expert gate and up dispatch * Plan CPU work on physical cores * Batch grouped expert CUDA kernels * Separate VRAM and RAM expert placement * Add ragged multi-sequence decode forward * feat(runtime): add continuous decode scheduler * Route concurrent API requests through batch scheduler * Harden multiplex request lifecycle and framing * Cancel disconnected multiplex requests * Bind API port before starting the engine * fix automatic KV slot allocation * add native int4 Tensor Core grouped GEMM * add Tensor Core throughput benchmark * optimize packed int4 low-row kernels * add asynchronous CUDA staging streams * document validated six-GPU dense acceleration * tune six-GPU expert hot set * raise validated expert hot-set target * add CUDA MLA absorption core * fuse grouped expert gate and up projections * Warn for explicit lossy routing flags
559 lines
28 KiB
Python
Executable File
559 lines
28 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
colibrì — tiny engine, immense model.
|
||
Run GLM-5.2 (744B) locally on CPU with roughly 15-26 GB of RAM.
|
||
|
||
coli chat interactive chat (loads the model once)
|
||
coli serve OpenAI-compatible HTTP API (persistent engine)
|
||
coli run "prompt" one-shot generation
|
||
coli info model, RAM, disk, and configuration status
|
||
coli plan Disk / RAM / VRAM resource plan
|
||
coli doctor installation and execution-plan diagnostics
|
||
coli bench [task...] quality benchmarks (MMLU/HellaSwag/...)
|
||
coli convert convert GLM-5.2-FP8 to int4, one shard at a time
|
||
coli build build the engine
|
||
|
||
Configuration through environment variables or flags (also valid after the subcommand):
|
||
COLI_MODEL=<dir> model directory (default /home/vincenzo/glm52_i4)
|
||
--ram N RAM budget in GB (automatically sizes the expert cache)
|
||
--repin N adapt RAM/VRAM experts every N tokens
|
||
--topp P adaptive expert top-p --topk N fixed top-k
|
||
--ngen N maximum response tokens --cap N cache slots/layer
|
||
"""
|
||
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap
|
||
|
||
# Windows: forza output UTF-8 (console cp1252 tronca Unicode box-drawing/emoji)
|
||
if sys.platform == "win32":
|
||
for s in (sys.stdout, sys.stderr):
|
||
try: s.reconfigure(encoding="utf-8")
|
||
except (AttributeError, OSError): pass
|
||
|
||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
TOOLS = os.path.join(HERE, "tools")
|
||
GLM = os.path.join(HERE, "glm" + (".exe" if sys.platform == "win32" else ""))
|
||
DEF_MODEL = os.environ.get("COLI_MODEL", "/home/vincenzo/glm52_i4")
|
||
END = b"\x01\x01END\x01\x01\n"
|
||
READY = b"\x01\x01READY\x01\x01\n"
|
||
|
||
# ---------- palette & stile ----------
|
||
def _c(n): return f"\033[38;5;{n}m"
|
||
class C:
|
||
teal=_c(37); cyan=_c(80); mag=_c(170); org=_c(208); grn=_c(78); yel=_c(179)
|
||
dim="\033[2m"; b="\033[1m"; r="\033[0m"; gray=_c(242); dgray=_c(238)
|
||
@staticmethod
|
||
def off():
|
||
for k,v in vars(C).items():
|
||
if isinstance(v,str) and v.startswith("\033"): setattr(C,k,"")
|
||
TTY = sys.stdout.isatty() or os.environ.get("COLI_COLOR")=="1"
|
||
if not TTY: C.off()
|
||
|
||
# ---------- colibrì 8-bit (pixel art, 2 pixel verticali per carattere) ----------
|
||
SPRITE = [
|
||
"....MMM.........",
|
||
"...MMMMM..w.....",
|
||
"....MMMM.ww.....",
|
||
"OOOOTTeTCC......",
|
||
"....TTTTTCC.....",
|
||
".....TTTTCC.....",
|
||
"......TTCC......",
|
||
".......TC.......",
|
||
"........C.......",
|
||
"................",
|
||
]
|
||
PAL = {"M":170, "T":37, "C":80, "O":208, "e":231, "w":80, ".":None}
|
||
|
||
def sprite_lines():
|
||
if not TTY:
|
||
return [" (\\ ", " )·> ", " / \\ ", " ", " "]
|
||
out=[]
|
||
for y in range(0,len(SPRITE),2):
|
||
top, bot = SPRITE[y], SPRITE[y+1] if y+1<len(SPRITE) else "."*len(SPRITE[y])
|
||
row=""
|
||
for x in range(len(top)):
|
||
ct, cb = PAL.get(top[x]), PAL.get(bot[x])
|
||
if ct is None and cb is None: row+= "\033[0m "
|
||
elif ct is not None and cb is None: row+= f"\033[38;5;{ct}m\033[49m▀"
|
||
elif ct is None and cb is not None: row+= f"\033[38;5;{cb}m\033[49m▄"
|
||
else: row+= f"\033[38;5;{ct}m\033[48;5;{cb}m▀"
|
||
out.append(row+"\033[0m")
|
||
return out
|
||
|
||
def banner(sub=""):
|
||
sp=sprite_lines()
|
||
txt=[
|
||
f"{C.teal}{C.b}colibrì{C.r} {C.dim}v1.0{C.r}",
|
||
f"{C.dim}tiny engine, immense model{C.r}",
|
||
f"{C.gray}GLM-5.2 · 744B MoE · int4 · streaming CPU{C.r}",
|
||
f"{C.dgray}{sub}{C.r}" if sub else "",
|
||
"",
|
||
]
|
||
print()
|
||
for i,s in enumerate(sp):
|
||
t = txt[i] if i<len(txt) else ""
|
||
print(f" {s} {t}")
|
||
print(f" {C.dgray}{'─'*58}{C.r}")
|
||
|
||
def hline(w): return f"{C.dgray}{'─'*w}{C.r}"
|
||
|
||
# ---------- util ----------
|
||
def term_w(): return min(shutil.get_terminal_size((80,20)).columns, 100)
|
||
|
||
def need_model(model):
|
||
if not os.path.isdir(model):
|
||
sys.exit(f"{C.yel}model not found:{C.r} {model}\n set COLI_MODEL or use --model")
|
||
if not os.path.exists(os.path.join(model,"tokenizer.json")):
|
||
sys.exit(f"{C.yel}tokenizer.json is missing from {model}{C.r}")
|
||
if not os.path.exists(GLM):
|
||
sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build")
|
||
|
||
def cuda_binary():
|
||
if not os.path.exists(GLM) or sys.platform != "linux": return False
|
||
try:
|
||
linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
|
||
return any("libcudart" in line and "not found" not in line
|
||
for line in linked.stdout.splitlines())
|
||
except (OSError,subprocess.SubprocessError): return False
|
||
|
||
def resource_request(a, env):
|
||
ctx=a.ctx or int(env.get("CTX",4096))
|
||
ram=a.ram or float(env.get("RAM_GB",0))
|
||
vram=a.vram or float(env.get("CUDA_EXPERT_GB",0))
|
||
gpu=a.gpu
|
||
if gpu is None:
|
||
gpu=env.get("COLI_GPUS",env.get("COLI_GPU","auto"))
|
||
devices=None if gpu=="auto" else ([] if gpu=="none" else
|
||
[int(value) for value in gpu.split(",")])
|
||
return ram,ctx,devices,vram
|
||
|
||
def env_for(a):
|
||
e = dict(os.environ, SNAP=a.model)
|
||
e["COLI_POLICY"]=a.policy
|
||
if a.ram: e["RAM_GB"]=str(a.ram)
|
||
if a.ngen: e["NGEN"]=str(a.ngen)
|
||
if a.topp: e["TOPP"]=str(a.topp)
|
||
if a.topk: e["TOPK"]=str(a.topk)
|
||
if a.temp is not None: e["TEMP"]=str(a.temp) # 0 = greedy; default motore: 1.0 + nucleus 0.95
|
||
if a.repin: e["REPIN"]=str(a.repin)
|
||
if a.ctx: e["CTX"]=str(a.ctx)
|
||
if a.auto_tier:
|
||
from resource_plan import build_plan, environment_for_plan, format_bytes
|
||
if a.gpu is not None:
|
||
e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
|
||
if a.gpu=="none":
|
||
e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
|
||
else: e.pop("COLI_CUDA",None)
|
||
elif e.get("COLI_CUDA")=="0":
|
||
e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
|
||
e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
|
||
if a.vram and a.gpu!="none": e["CUDA_EXPERT_GB"]=str(a.vram)
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,e)
|
||
plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
|
||
except (OSError,ValueError,json.JSONDecodeError) as error:
|
||
sys.exit(f"{C.yel}invalid resource plan:{C.r} {error}")
|
||
has_cuda=cuda_binary()
|
||
e=environment_for_plan(plan,e,has_cuda)
|
||
rt=plan["tiers"]["ram"]; vt=plan["tiers"]["vram"]
|
||
gpu=f" · VRAM {format_bytes(vt['budget_bytes'])}" if has_cuda and vt["devices"] else " · CPU"
|
||
print(f" {C.dim}[PLAN] RAM {format_bytes(rt['budget_bytes'])} · cap {rt['cache_slots_per_layer']}/layer{gpu}{C.r}",file=sys.stderr)
|
||
return e
|
||
|
||
# ---------- rendering markdown in STREAMING per il terminale ----------
|
||
class MDStream:
|
||
"""Interpreta il markdown della risposta mentre arriva: i ``` diventano riquadri,
|
||
**x** grassetto vero, `x` colorato, # titoli, - puntini. I marker non si vedono mai.
|
||
Regge i chunk spezzati a meta' marker (hold-back) e l'output sporco (``` doppi)."""
|
||
def __init__(self, indent=" "):
|
||
self.ind=indent
|
||
self.cur="" # riga parziale non ancora emessa
|
||
self.code=False; self.lang=""
|
||
self.bold=False; self.icode=False
|
||
self.justclosed=False # l'ultima riga era una chiusura ```? (anti ``` doppi)
|
||
self.printed=0 # caratteri della riga corrente gia' emessi
|
||
def _fence(self, line):
|
||
lang=line.strip()[3:].strip().strip("`")
|
||
if not self.code:
|
||
if not lang and self.justclosed: return # ``` orfano dopo una chiusura: rumore, ignora
|
||
self.code=True; self.lang=lang
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u256d\u2500 {lang or 'code'}{C.r}\n")
|
||
elif lang: # ```lang mentre siamo GIA' in code: chiudi e riapri
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n{self.ind}{C.dgray}\u256d\u2500 {lang}{C.r}\n")
|
||
self.lang=lang
|
||
else:
|
||
self.code=False; self.justclosed=True
|
||
sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n")
|
||
def _inline(self, txt, out):
|
||
i=0
|
||
while i<len(txt):
|
||
ch=txt[i]
|
||
if ch=="`":
|
||
self.icode=not self.icode
|
||
out.append(C.org if self.icode else C.r); i+=1; continue
|
||
if ch=="*":
|
||
j=i
|
||
while j<len(txt) and txt[j]=="*": j+=1
|
||
if j-i>=2: # **/***: grassetto on/off, gli asterischi spariscono
|
||
self.bold=not self.bold
|
||
out.append(C.b if self.bold else C.r)
|
||
else: out.append("*") # * singolo: lascialo (moltiplicazioni ecc.)
|
||
i=j; continue
|
||
out.append(ch); i+=1
|
||
def _line(self, line, partial=False):
|
||
if not partial and line.lstrip().startswith("```"):
|
||
self._fence(line); self.printed=0; return
|
||
if line.strip(): self.justclosed=False
|
||
seg=line[self.printed:] # emetti solo la parte nuova della riga
|
||
out=[]
|
||
if self.code:
|
||
if self.printed==0: out.append(f"{self.ind}{C.dgray}\u2502{C.r} {C.cyan}")
|
||
out.append(seg)
|
||
else:
|
||
if self.printed==0:
|
||
out.append(self.ind)
|
||
st=seg.lstrip()
|
||
if st.startswith("#"): # titolo: via i #, grassetto teal
|
||
seg=st.lstrip("#").strip(); out.append(f"{C.teal}{C.b}"); self.bold=True
|
||
elif st.startswith(("- ","* ")): # lista: puntino vero
|
||
seg=st[2:]; out.append(f"{C.teal}\u2022{C.r} ")
|
||
self._inline(seg,out)
|
||
sys.stdout.write("".join(out)); sys.stdout.flush()
|
||
self.printed=len(line)
|
||
if not partial: # fine riga: reset stati inline (robusto ai marker orfani)
|
||
sys.stdout.write(C.r+"\n"); sys.stdout.flush()
|
||
self.bold=self.icode=False
|
||
self.printed=0
|
||
def feed(self, s):
|
||
self.cur+=s
|
||
while "\n" in self.cur:
|
||
line,self.cur=self.cur.split("\n",1)
|
||
self._line(line)
|
||
st=self.cur.lstrip() # riga parziale: possibile fence? aspetta il newline
|
||
if st and (st.startswith("```") or (len(st)<3 and "```".startswith(st))):
|
||
return
|
||
if st.startswith("#") and self.printed==0:
|
||
return # titolo: rendi la riga intera al newline
|
||
hold=0 # trattieni marker potenzialmente spezzati in coda
|
||
while hold<len(self.cur) and self.cur[-1-hold] in "*`": hold+=1
|
||
safe=self.cur[:len(self.cur)-hold] if hold else self.cur
|
||
if len(safe)>self.printed: self._line(safe, partial=True)
|
||
def close(self):
|
||
if self.cur: self._line(self.cur); self.cur=""
|
||
if self.code:
|
||
sys.stdout.write(f"\n{self.ind}{C.dgray}\u2570\u2500{C.r}"); self.code=False
|
||
sys.stdout.write(C.r); sys.stdout.flush()
|
||
|
||
class Spinner:
|
||
FRAMES=["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
|
||
def __init__(self,label,tick=None):
|
||
self.label=label; self.tick=tick; self.suffix=""
|
||
self.stop_evt=threading.Event(); self.t0=time.time(); self.th=None
|
||
def start(self):
|
||
if not TTY: return
|
||
def run():
|
||
i=0
|
||
while not self.stop_evt.is_set():
|
||
el=time.time()-self.t0
|
||
if self.tick and i%8==0: # ~1 Hz: legge il progresso dal log
|
||
try: self.suffix=self.tick() or self.suffix
|
||
except Exception: pass
|
||
suf=f" {C.dgray}· {self.suffix}{C.r}" if self.suffix else ""
|
||
sys.stdout.write(f"\r {C.teal}{self.FRAMES[i%10]}{C.r} {C.dim}{self.label} {el:.0f}s{C.r}{suf}\033[K")
|
||
sys.stdout.flush(); i+=1; time.sleep(0.12)
|
||
self.th=threading.Thread(target=run,daemon=True); self.th.start()
|
||
def stop(self):
|
||
self.stop_evt.set()
|
||
if self.th: self.th.join(timeout=0.4)
|
||
if TTY: sys.stdout.write("\r\033[K"); sys.stdout.flush()
|
||
|
||
def stream_turn(p, sentinel, on_bytes):
|
||
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT."""
|
||
pend=b""
|
||
while True:
|
||
b=p.stdout.read(1)
|
||
if b==b"": return None
|
||
pend+=b
|
||
if pend.endswith(sentinel):
|
||
rest=pend[:-len(sentinel)]
|
||
if rest: on_bytes(rest)
|
||
line=p.stdout.readline().decode("utf-8","replace").strip() # STAT tok tps hit rss
|
||
m=re.match(r"STAT (\S+) (\S+) (\S+) (\S+)", line)
|
||
return {"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
|
||
if len(pend)>len(sentinel):
|
||
out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
|
||
on_bytes(out)
|
||
|
||
# ---------- comandi ----------
|
||
def cmd_build(a):
|
||
banner("build")
|
||
sys.exit(subprocess.call(["make","-C",HERE,"glm"]))
|
||
|
||
def cmd_info(a):
|
||
banner("info")
|
||
cfgp=os.path.join(a.model,"config.json")
|
||
def row(k,v): print(f" {C.gray}{k:<10}{C.r} {v}")
|
||
if os.path.exists(cfgp):
|
||
c=json.load(open(cfgp))
|
||
row("model", a.model)
|
||
row("arch", f"hidden {c.get('hidden_size')} · {c.get('num_hidden_layers')} layer · "
|
||
f"{c.get('n_routed_experts')} expert/layer · top-{c.get('num_experts_per_tok')}")
|
||
sts=[x for x in os.listdir(a.model) if x.endswith('.safetensors')]
|
||
sz=sum(os.path.getsize(os.path.join(a.model,x)) for x in sts)
|
||
row("shards", f"{len(sts)} files · {sz/1e9:.0f} GB on disk")
|
||
else:
|
||
print(f" {C.yel}config.json is missing (incomplete conversion?){C.r}")
|
||
try:
|
||
mi=open('/proc/meminfo').read()
|
||
tot=int(re.search(r'MemTotal:\s+(\d+)',mi).group(1))/1e6
|
||
av=int(re.search(r'MemAvailable:\s+(\d+)',mi).group(1))/1e6
|
||
row("RAM", f"{tot:.0f} GB total · {av:.1f} GB available")
|
||
except Exception: pass
|
||
try:
|
||
fs = shutil.disk_usage(a.model if os.path.isdir(a.model) else HERE)
|
||
row("disk", f"{fs.free/1e9:.0f} GB free")
|
||
except OSError:
|
||
row("disk", "? GB (unavailable)")
|
||
row("engine", "ready ✓" if os.path.exists(GLM) else "not built (coli build)")
|
||
knobs=[]
|
||
if a.ram: knobs.append(f"ram {a.ram}GB")
|
||
if a.topp: knobs.append(f"topp {a.topp}")
|
||
if a.topk: knobs.append(f"topk {a.topk}")
|
||
if knobs: row("tuning", " · ".join(knobs))
|
||
print()
|
||
|
||
def cmd_plan(a):
|
||
from resource_plan import build_plan, format_plan
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,os.environ)
|
||
if ctx<1: raise ValueError("--ctx must be positive")
|
||
if a.vram<0: raise ValueError("--vram cannot be negative")
|
||
plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
|
||
except (OSError, ValueError, json.JSONDecodeError) as error:
|
||
sys.exit(f"{C.yel}cannot create resource plan:{C.r} {error}")
|
||
if a.json:
|
||
print(json.dumps(plan,indent=2))
|
||
return
|
||
banner("plan · Disk / RAM / VRAM")
|
||
print(textwrap.indent(format_plan(plan)," "))
|
||
print()
|
||
|
||
def cmd_doctor(a):
|
||
from doctor import exit_code, format_doctor, run_doctor
|
||
try:
|
||
ram,ctx,devices,vram=resource_request(a,os.environ)
|
||
if ctx<1: raise ValueError("--ctx must be positive")
|
||
if ram<0: raise ValueError("--ram cannot be negative")
|
||
if vram<0: raise ValueError("--vram cannot be negative")
|
||
except ValueError as error:
|
||
report={"schema_version":1,"status":"error","model":os.path.abspath(a.model),
|
||
"checks":[{"id":"config.arguments","status":"fail","summary":str(error)}],
|
||
"plan":None}
|
||
print(json.dumps(report,indent=2) if a.json else format_doctor(report))
|
||
return 2
|
||
report=run_doctor(a.model,ram,ctx,devices,vram,engine_path=GLM)
|
||
print(json.dumps(report,indent=2) if a.json else format_doctor(report))
|
||
return exit_code(report)
|
||
|
||
def cmd_run(a):
|
||
need_model(a.model)
|
||
prompt=" ".join(a.prompt) if a.prompt else sys.exit('usage: coli run "your prompt"')
|
||
banner("run")
|
||
# template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink)
|
||
e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|><think></think>"
|
||
sys.exit(subprocess.call([GLM, str(a.cap)], env=e))
|
||
|
||
def cmd_chat(a):
|
||
need_model(a.model)
|
||
banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
|
||
errlog=tempfile.NamedTemporaryFile(mode="w+", suffix=".log", delete=False)
|
||
e=env_for(a); e["SERVE"]="1"
|
||
p=subprocess.Popen([GLM,str(a.cap)], env=e, stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE, stderr=errlog, bufsize=0)
|
||
sp=Spinner("waking the giant (744B)…"); sp.start()
|
||
st=stream_turn(p, READY, lambda b: None)
|
||
sp.stop()
|
||
if st is None:
|
||
errlog.seek(0); print(errlog.read()[-1500:]); sys.exit("the engine exited while loading")
|
||
errlog.flush()
|
||
try:
|
||
elog=open(errlog.name).read()
|
||
mload=re.search(r"loaded in ([0-9.]+)s \| resident dense: ([0-9.]+) MB", elog)
|
||
if mload: print(f" {C.grn}✓{C.r} ready in {mload.group(1)}s {C.dim}· resident {float(mload.group(2))/1000:.1f} GB · RSS {st.get('rss','?')} GB{C.r}")
|
||
for l in elog.splitlines(): # una riga di stato per riga, senza path
|
||
if l.startswith(("[RAM_GB","[PIN]","[MTP]","[USAGE]","[DSA]","[KV]")):
|
||
l=re.sub(r" ?\(?/[^ )]+\)?","",l.strip()) # via i percorsi lunghi
|
||
l=re.sub(r" from$","",l)
|
||
for chunk in textwrap.wrap(l, term_w()-4) or [l]:
|
||
print(f" {C.dgray}{chunk}{C.r}")
|
||
except Exception: pass
|
||
print(f" {C.dim}type and press Enter · :more continues · :reset clears memory · :q exits{C.r}\n")
|
||
w=term_w()-4
|
||
def user_box(msg):
|
||
"""ri-disegna il messaggio dentro una box che si ADATTA su piu' righe:
|
||
l'input grezzo (che sborda) viene cancellato e sostituito dal testo avvolto."""
|
||
cols=shutil.get_terminal_size((80,20)).columns
|
||
used=max(1, (6+len(msg)+cols-1)//cols) # righe occupate dall'input (" │ › "+msg)
|
||
sys.stdout.write(f"\x1b[{used}A\x1b[0J") # su di N righe e pulisci fino in fondo
|
||
inner=w-3 # spazio utile: " │ › "+testo+"│" = w+4 colonne
|
||
lines=textwrap.wrap(msg, inner) or [""]
|
||
for i,ln in enumerate(lines):
|
||
pre = f"{C.teal}{C.b}›{C.r}" if i==0 else " "
|
||
print(f" {C.dgray}│{C.r} {pre} {ln}{' '*(inner-len(ln))}{C.dgray}│{C.r}")
|
||
print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||
try:
|
||
while True:
|
||
if TTY:
|
||
print(f" {C.dgray}╭{'─'*w}╮{C.r}")
|
||
try: msg=input(f" {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
|
||
except EOFError: print(); break
|
||
try: user_box(msg.strip())
|
||
except Exception: print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||
else:
|
||
try: msg=input()
|
||
except EOFError: break
|
||
msg=msg.strip()
|
||
if msg in (":q",":quit","exit"): break
|
||
if not msg: continue
|
||
if msg==":reset":
|
||
p.stdin.write(b"\x02RESET\n"); p.stdin.flush()
|
||
stream_turn(p, END, lambda b: None)
|
||
print(f" {C.dim}✦ memory cleared{C.r}\n"); continue
|
||
if msg in (":piu",":più",":more",":continua"):
|
||
p.stdin.write(b"\x02MORE\n"); p.stdin.flush()
|
||
else:
|
||
p.stdin.write((msg.replace("\n"," ")+"\n").encode()); p.stdin.flush()
|
||
print(f"\n {C.teal}◆ colibrì{C.r}")
|
||
dec=codecs.getincrementaldecoder("utf-8")("replace")
|
||
state={"first":True}
|
||
def prefill_tick(path=errlog.name):
|
||
try:
|
||
with open(path) as f:
|
||
f.seek(max(0, os.path.getsize(path)-1500)); tail=f.read()
|
||
pl=[l for l in tail.splitlines() if l.startswith("[prefill]")]
|
||
return pl[-1].replace("[prefill] ","prefill ") if pl else ""
|
||
except Exception: return ""
|
||
sp2=Spinner("thinking…", tick=prefill_tick); sp2.start()
|
||
md=MDStream(" ") # markdown -> terminale, in streaming
|
||
raw=os.environ.get("COLI_RAW")=="1"
|
||
def echo(bs, _dec=dec, _st=state):
|
||
if _st["first"]:
|
||
sp2.stop(); _st["first"]=False
|
||
if raw: sys.stdout.write(" ")
|
||
s=_dec.decode(bs)
|
||
if not s: return
|
||
if raw: sys.stdout.write(s.replace("\n","\n ")); sys.stdout.flush()
|
||
else: md.feed(s)
|
||
t0=time.time()
|
||
st=stream_turn(p, END, echo)
|
||
if not raw: md.close()
|
||
sp2.stop()
|
||
if st is None: print(f"\n {C.yel}[engine terminated]{C.r}"); break
|
||
el=time.time()-t0
|
||
if st.get("tok"):
|
||
print(f"\r {C.dgray}└─ {st['tok']} tok · {st['tps']:.2f} tok/s · hit {st['hit']:.0f}% · RSS {st['rss']:.1f} GB · {el:.0f}s{C.r}")
|
||
if st["tok"]>=a.ngen:
|
||
print(f" {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
|
||
print()
|
||
else:
|
||
print()
|
||
except KeyboardInterrupt:
|
||
print(f"\n {C.dim}interrupted{C.r}")
|
||
finally:
|
||
try: p.stdin.close(); p.terminate()
|
||
except Exception: pass
|
||
try: os.unlink(errlog.name)
|
||
except Exception: pass
|
||
print(f" {C.teal}goodbye{C.r} {C.dim}— the hummingbird returns to its nest{C.r} 🐦\n")
|
||
|
||
def cmd_serve(a):
|
||
need_model(a.model)
|
||
from openai_server import serve
|
||
serve(a.model, a.host, a.port, a.model_id, a.api_key,
|
||
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
|
||
a.max_queue,a.queue_timeout,a.kv_slots)
|
||
|
||
def cmd_bench(a):
|
||
need_model(a.model)
|
||
banner("bench")
|
||
# python con `tokenizers`: l'ambiente del progetto se c'e', altrimenti quello corrente
|
||
venv_py = os.path.join(HERE, "mio_env", "Scripts" if sys.platform == "win32" else "bin", "python3")
|
||
py = venv_py if os.path.exists(venv_py) else sys.executable
|
||
tasks = ",".join(a.tasks) if a.tasks else "hellaswag,arc_challenge,mmlu"
|
||
# dataset mancanti -> li scarica una volta (fetch_benchmarks.py li mette in --data come JSONL)
|
||
missing=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
|
||
if missing:
|
||
print(f" {C.dim}downloading missing datasets: {', '.join(missing)}{C.r}")
|
||
subprocess.call([py, os.path.join(TOOLS,"fetch_benchmarks.py"),
|
||
"--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
|
||
cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--snap",a.model,
|
||
"--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
|
||
if a.ram: cmd+=["--ram",str(a.ram)]
|
||
e=env_for(a)
|
||
print(f" {C.dim}decode is disk-bound: this takes HOURS on slow hardware. Raise --limit on faster machines.{C.r}\n")
|
||
sys.exit(subprocess.call(cmd, env=e))
|
||
|
||
def cmd_convert(a):
|
||
banner("convert")
|
||
# python con torch/safetensors: l'ambiente del progetto se c'e', altrimenti quello corrente
|
||
venv_py = os.path.join(HERE, "mio_env", "Scripts" if sys.platform == "win32" else "bin", "python3")
|
||
py = venv_py if os.path.exists(venv_py) else sys.executable
|
||
base=[py, os.path.join(TOOLS,"convert_fp8_to_int4.py"),
|
||
"--repo", a.repo, "--outdir", a.model, "--ebits", str(a.ebits), "--io-bits", str(a.io_bits)]
|
||
if a.xbits: base+=["--xbits",str(a.xbits)]
|
||
# passo 1: modello principale (78 layer). Resumabile: riparte dagli shard mancanti.
|
||
print(f" {C.dim}[1/2] model: {' '.join(base)}{C.r}")
|
||
rc=subprocess.call(base)
|
||
if rc!=0: sys.exit(rc)
|
||
if a.no_mtp: sys.exit(0)
|
||
# passo 2: testa MTP (layer 78). SEMPRE int8: a int4 i draft sbagliano quasi sempre
|
||
# (acceptance 0-4% vs 39-59%, misurato — issue #8) e la speculazione non parte mai.
|
||
mtp_cmd=list(base); i=mtp_cmd.index("--ebits"); mtp_cmd[i+1]=str(max(8,a.ebits))
|
||
print(f" {C.dim}[2/2] int8 MTP head (speculative drafts){C.r}")
|
||
sys.exit(subprocess.call(mtp_cmd+["--mtp"]))
|
||
|
||
def main():
|
||
common=argparse.ArgumentParser(add_help=False)
|
||
common.add_argument("--model", default=DEF_MODEL); common.add_argument("--ram", type=int, default=0) # 0 = auto (il motore usa l'88% della RAM disponibile)
|
||
common.add_argument("--auto-tier",action="store_true",help="automatically apply the RAM/VRAM plan")
|
||
common.add_argument("--ctx",type=int,default=0)
|
||
common.add_argument("--gpu",default=None,help="auto, none, or a device list such as 0,1")
|
||
common.add_argument("--vram",type=float,default=0,help="total VRAM budget in GB (0=auto)")
|
||
common.add_argument("--policy",choices=("quality","balanced","experimental-fast"),
|
||
default=os.environ.get("COLI_POLICY","quality"),
|
||
help="resource policy (explicit --topk/--topp overrides warn and proceed)")
|
||
common.add_argument("--repin", type=int, default=0, help="adapt RAM/VRAM experts every N tokens")
|
||
common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024) # rete di sicurezza: la fine vera la decidono gli stop token
|
||
common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0)
|
||
common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95)
|
||
ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — run GLM-5.2 locally")
|
||
sub=ap.add_subparsers(dest="cmd")
|
||
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
|
||
pp=sub.add_parser("plan",parents=[common])
|
||
pp.add_argument("--json",action="store_true")
|
||
pd=sub.add_parser("doctor",parents=[common])
|
||
pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
|
||
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
|
||
sub.add_parser("chat", parents=[common])
|
||
ps=sub.add_parser("serve", parents=[common])
|
||
ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
|
||
ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
|
||
ps.add_argument("--api-key",default=os.environ.get("COLI_API_KEY"))
|
||
ps.add_argument("--cors-origin",action="append",default=None)
|
||
ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
|
||
ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
|
||
ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1")))
|
||
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
|
||
pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
|
||
pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
|
||
pc.add_argument("--ebits",type=int,default=4); pc.add_argument("--io-bits",type=int,default=8); pc.add_argument("--xbits",type=int,default=0)
|
||
pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
|
||
a=ap.parse_args()
|
||
handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,
|
||
"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
|
||
"convert":cmd_convert}.get(a.cmd)
|
||
if handler: sys.exit(handler(a) or 0)
|
||
banner(); print(__doc__)
|
||
|
||
if __name__=="__main__":
|
||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||
main()
|