JustVugg
193d2ce92d
- Faithful to the official modeling (transformers glm_moe_dsa): q from the q_a latent via wq_b (32 heads x 128), k = LayerNorm(wk(h)) shared across heads, interleaved RoPE on the first 64 dims, ReLU(q.k/sqrt(128)) weighted by weights_proj(h)/sqrt(32), causal top-2048 per query. - 'full' layers compute the selection (+ maintain the indexer k-cache from token 0); 'shared' layers reuse it (IndexShare, index_topk_freq=4). - Selection restricts both attention paths (absorbed decode + prefill reconstruction). MTP row stays dense. - Auto-detected like MTP: if out-idx-* weights are present for all full layers, DSA arms itself; DSA=0 disables; DSA_FORCE/DSA_TOPK for testing. - Validated on the tiny oracle (which ships indexer weights): selection machinery forced on with keep=all keys reproduces dense attention exactly (TF 32/32, gen 20/20); sparse smoke runs clean; kill switch verified. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
340 lines
16 KiB
Python
340 lines
16 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, textwrap
|
||
|
||
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)
|
||
if a.temp is not None: e["TEMP"]=str(a.temp) # 0 = greedy; default motore: 1.0 + nucleus 0.95
|
||
return e
|
||
|
||
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("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)
|
||
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}")
|
||
for l in elog.splitlines(): # una riga di stato per riga, senza path
|
||
if l.startswith(("[RAM_GB","[PIN]","[MTP]","[USAGE]","[DSA]")):
|
||
l=re.sub(r" ?\(?/[^ )]+\)?","",l.strip()) # via i percorsi lunghi
|
||
l=re.sub(r" da$| in$","",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}scrivi e premi invio · :piu continua risposta · :reset memoria · :q esci{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}✦ memoria azzerata{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("pensa…", tick=prefill_tick); 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}")
|
||
if st["tok"]>=a.ngen:
|
||
print(f" {C.yel}…troncato al limite --ngen ({a.ngen}): scrivi :piu per far continuare la risposta{C.r}")
|
||
print()
|
||
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")
|
||
# python con torch/safetensors: l'ambiente del progetto se c'e', altrimenti quello corrente
|
||
venv_py=os.path.join(HERE,"mio_env","bin","python3")
|
||
py = venv_py if os.path.exists(venv_py) else sys.executable
|
||
base=[py, 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: base+=["--xbits",str(a.xbits)]
|
||
# passo 1: modello principale (78 layer). Resumabile: riparte dagli shard mancanti.
|
||
print(f" {C.dim}[1/2] modello: {' '.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) -> abilita la decodifica speculativa nativa (piu' veloce)
|
||
print(f" {C.dim}[2/2] testa MTP (draft speculativi){C.r}")
|
||
sys.exit(subprocess.call(base+["--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("--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ì — 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)
|
||
pc.add_argument("--no-mtp",action="store_true",help="salta la testa MTP (niente draft speculativi)")
|
||
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()
|