JustVugg
1ae22a6135
Engine (c/glm.c): MLA attention with compressed KV, sigmoid noaux_tc router, int8/int4/int2 quant kernels (AVX2), per-layer LRU expert cache + pinned hot-store, batch-union MoE, native MTP speculative decoding (lossless), multi-stop + official chat template, RAM auto-budget from MemAvailable. Tokenizer: byte-level BPE in C. Tooling: coli CLI, disk-safe FP8→int4 converter, tiny-random oracle validation (TF 32/32, greedy 20/20). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
292 lines
13 KiB
Python
292 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
colibrì — piccolo motore, modello immenso.
|
||
CLI per far girare GLM-5.2 (744B) in locale, su CPU, in ~15-26 GB di RAM.
|
||
|
||
coli chat chat interattiva (carica il modello UNA volta)
|
||
coli run "prompt" generazione singola
|
||
coli info stato: modello, RAM, disco, config
|
||
coli bench [task...] benchmark di qualità (MMLU/HellaSwag/...)
|
||
coli convert converte GLM-5.2-FP8 -> int4 (streaming)
|
||
coli build compila il motore
|
||
|
||
Config via env o flag (validi anche dopo il sottocomando):
|
||
COLI_MODEL=<dir> modello (default /home/vincenzo/glm52_i4)
|
||
--ram N budget RAM in GB (auto-cap cache expert)
|
||
--topp P top-p adattivo sugli expert --topk N top-k fisso
|
||
--ngen N token massimi per risposta --cap N slot cache/layer
|
||
"""
|
||
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile
|
||
|
||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
GLM = os.path.join(HERE, "glm")
|
||
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}piccolo motore, modello immenso{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}modello non trovato:{C.r} {model}\n imposta COLI_MODEL o usa --model")
|
||
if not os.path.exists(os.path.join(model,"tokenizer.json")):
|
||
sys.exit(f"{C.yel}manca tokenizer.json in {model}{C.r}")
|
||
if not os.path.exists(GLM):
|
||
sys.exit(f"{C.yel}motore non compilato.{C.r} Esegui: coli build")
|
||
|
||
def env_for(a):
|
||
e = dict(os.environ, SNAP=a.model)
|
||
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)
|
||
return e
|
||
|
||
class Spinner:
|
||
FRAMES=["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
|
||
def __init__(self,label): self.label=label; 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
|
||
sys.stdout.write(f"\r {C.teal}{self.FRAMES[i%10]}{C.r} {C.dim}{self.label} {el:.0f}s{C.r}\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("modello", 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("shard", f"{len(sts)} file · {sz/1e9:.0f} GB su disco")
|
||
else:
|
||
print(f" {C.yel}config.json non presente (conversione incompleta?){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 totali · {av:.1f} GB disponibili")
|
||
except Exception: pass
|
||
fs=os.statvfs(a.model if os.path.isdir(a.model) else HERE)
|
||
row("disco", f"{fs.f_bavail*fs.f_frsize/1e9:.0f} GB liberi")
|
||
row("motore", "pronto ✓" if os.path.exists(GLM) else "da compilare (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_run(a):
|
||
need_model(a.model)
|
||
prompt=" ".join(a.prompt) if a.prompt else sys.exit('uso: coli run "il tuo 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("sveglio il gigante (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("il motore è uscito durante il load")
|
||
errlog.flush()
|
||
try:
|
||
elog=open(errlog.name).read()
|
||
mload=re.search(r"caricato in ([0-9.]+)s \| densa residente: ([0-9.]+) MB", elog)
|
||
extra=" · ".join(l.strip() for l in elog.splitlines() if l.startswith("[RAM_GB") or l.startswith("[PIN]"))
|
||
if mload: print(f" {C.grn}✓{C.r} pronto in {mload.group(1)}s {C.dim}· residente {float(mload.group(2))/1000:.1f} GB · RSS {st.get('rss','?')} GB{C.r}")
|
||
if extra: print(f" {C.dgray}{extra}{C.r}")
|
||
except Exception: pass
|
||
print(f" {C.dim}scrivi e premi invio · :reset memoria · :q esci{C.r}\n")
|
||
w=term_w()-4
|
||
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
|
||
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}✦ memoria azzerata{C.r}\n"); continue
|
||
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}
|
||
sp2=Spinner("pensa…"); sp2.start()
|
||
def echo(bs, _dec=dec, _st=state):
|
||
if _st["first"]:
|
||
sp2.stop(); _st["first"]=False
|
||
sys.stdout.write(" ")
|
||
s=_dec.decode(bs)
|
||
if s: sys.stdout.write(s.replace("\n","\n ")); sys.stdout.flush()
|
||
t0=time.time()
|
||
st=stream_turn(p, END, echo)
|
||
sp2.stop()
|
||
if st is None: print(f"\n {C.yel}[motore terminato]{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}\n")
|
||
else:
|
||
print()
|
||
except KeyboardInterrupt:
|
||
print(f"\n {C.dim}interrotto{C.r}")
|
||
finally:
|
||
try: p.stdin.close(); p.terminate()
|
||
except Exception: pass
|
||
try: os.unlink(errlog.name)
|
||
except Exception: pass
|
||
print(f" {C.teal}ciao{C.r} {C.dim}— il colibrì torna al nido{C.r} 🐦\n")
|
||
|
||
def cmd_bench(a):
|
||
need_model(a.model)
|
||
banner("bench")
|
||
cmd=[sys.executable, os.path.join(HERE,"eval_glm.py"), "--snap",a.model,
|
||
"--tasks", ",".join(a.tasks) if a.tasks else "hellaswag,arc_challenge,mmlu",
|
||
"--limit", str(a.limit), "--data", a.data]
|
||
if a.ram: cmd+=["--ram",str(a.ram)]
|
||
e=dict(os.environ)
|
||
if a.topp: e["TOPP"]=str(a.topp)
|
||
if a.topk: e["TOPK"]=str(a.topk)
|
||
sys.exit(subprocess.call(cmd, env=e))
|
||
|
||
def cmd_convert(a):
|
||
banner("convert")
|
||
cmd=[sys.executable, os.path.join(HERE,"convert_fp8_to_int4.py"),
|
||
"--repo", a.repo, "--outdir", a.model, "--ebits", str(a.ebits), "--io-bits", str(a.io_bits)]
|
||
if a.xbits: cmd+=["--xbits",str(a.xbits)]
|
||
print(f" {C.dim}{' '.join(cmd)}{C.r}")
|
||
sys.exit(subprocess.call(cmd))
|
||
|
||
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("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=256)
|
||
common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0)
|
||
ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — GLM-5.2 in locale")
|
||
sub=ap.add_subparsers(dest="cmd")
|
||
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
|
||
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
|
||
sub.add_parser("chat", parents=[common])
|
||
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)
|
||
a=ap.parse_args()
|
||
{"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"bench":cmd_bench,
|
||
"convert":cmd_convert}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)
|
||
|
||
if __name__=="__main__":
|
||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||
main()
|