diff --git a/c/glm.c b/c/glm.c index fb93e59..3e24c31 100644 --- a/c/glm.c +++ b/c/glm.c @@ -142,6 +142,9 @@ typedef struct { int *kv_start, max_t; int disk_nrec; char disk_path[2048]; + FILE *disk_fp; /* kept-open handle: fopen once, fwrite per turn, fclose at exit (#4) */ + uint8_t *disk_buf; /* staging buffer: one contiguous record per position (#1) */ + int64_t disk_buf_cap; } KVState; typedef struct { @@ -3971,36 +3974,73 @@ static void kv_hdr(Model *m, int32_t *h, int nrec){ h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope; h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0; } +/* Bytes of one on-disk record: [tok i32][Lc+Rc per layer][Ic per DSA layer]. + * Layout matches what kv_disk_append writes and kv_disk_load reads. */ +static int64_t kv_rec_bytes(Model *m){ + Cfg *c=&m->c; + int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4; + if(m->has_dsa) for(int i=0;in_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4; + return rec; +} +/* Open the persistent handle lazily; write the header if the file is new. After + * this returns successfully, k->disk_fp is valid for the engine's lifetime and + * positioned at end-of-header (nrec==0 case) or wherever the caller seeks. */ +static int kv_disk_open(Model *m){ + KVState *k=m->kv; + if(k->disk_fp) return 1; + k->disk_fp=fopen(k->disk_path,"r+b"); + if(!k->disk_fp){ /* not there yet -> create + header */ + k->disk_fp=fopen(k->disk_path,"wb"); + if(!k->disk_fp) return 0; + int32_t h[8]; kv_hdr(m,h,0); + fwrite(KV_MAGIC,1,8,k->disk_fp); fwrite(h,4,8,k->disk_fp); + fflush(k->disk_fp); + fclose(k->disk_fp); + k->disk_fp=fopen(k->disk_path,"r+b"); /* reopen r+b for append */ + if(!k->disk_fp) return 0; + } + return 1; +} static void kv_disk_truncate(Model *m, int nrec){ if(!g_kvsave) return; KVState *k=m->kv; + if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; } /* drop to shrink on disc */ FILE *f=fopen(k->disk_path,"r+b"); if(!f){ k->disk_nrec=0; return; } k->disk_nrec=nrec; - int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f); + int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); + fflush(f); fclose(f); } static void kv_disk_reset(Model *m){ kv_disk_truncate(m,0); } static void kv_disk_append(Model *m, const int *hist, int len){ KVState *k=m->kv; if(!g_kvsave || len<=k->disk_nrec) return; Cfg *c=&m->c; - FILE *f=fopen(k->disk_path,"r+b"); - if(!f){ f=fopen(k->disk_path,"wb"); if(!f) return; - int32_t h[8]; kv_hdr(m,h,0); fwrite(KV_MAGIC,1,8,f); fwrite(h,4,8,f); } - int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4; - if(m->has_dsa) for(int i=0;in_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4; + if(!kv_disk_open(m)) return; + FILE *f=k->disk_fp; + int64_t rec = kv_rec_bytes(m); + /* grow the contiguous staging buffer if the record is larger (#1 batching) */ + if(rec > k->disk_buf_cap){ + uint8_t *nb=realloc(k->disk_buf, rec); + if(!nb) return; /* OOM: skip this turn, retry next */ + k->disk_buf=nb; k->disk_buf_cap=rec; + } fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET); for(int p=k->disk_nrec;pdisk_buf; /* pack token + every layer into one record */ + *(int32_t*)b = hist[p]; b+=4; for(int i=0;in_layers;i++){ - fwrite(m->Lc[i]+(int64_t)p*c->kv_lora, 4, c->kv_lora, f); - fwrite(m->Rc[i]+(int64_t)p*c->qk_rope, 4, c->qk_rope, f); + memcpy(b, m->Lc[i]+(int64_t)p*c->kv_lora, (size_t)c->kv_lora*4); b+=c->kv_lora*4; + memcpy(b, m->Rc[i]+(int64_t)p*c->qk_rope,(size_t)c->qk_rope*4); b+=c->qk_rope*4; } - if(m->has_dsa) for(int i=0;in_layers;i++) if(m->Ic[i]) - fwrite(m->Ic[i]+(int64_t)p*c->index_hd, 4, c->index_hd, f); + if(m->has_dsa) for(int i=0;in_layers;i++) if(m->Ic[i]){ + memcpy(b, m->Ic[i]+(int64_t)p*c->index_hd, (size_t)c->index_hd*4); b+=c->index_hd*4; + } + fwrite(k->disk_buf, 1, (size_t)rec, f); /* one fwrite per position (was ~157) */ } fflush(f); /* dati prima, contatore poi */ - int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f); + int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); + fflush(f); /* persist the counter too */ k->disk_nrec=len; } static int kv_disk_load(Model *m, int *hist, int maxctx){ @@ -4055,6 +4095,8 @@ static void serve_ctx_init(Model *m, ServeCtx *s, const char *snap, int slot, in static void serve_ctx_free(Model *m, ServeCtx *s){ KVState *k=&s->kv; int NR=m->c.n_layers+1; + if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; } + free(k->disk_buf); k->disk_buf=NULL; if(k->Lc) for(int i=0;iLc[i]); free(k->Rc[i]); } if(k->Ic) for(int i=0;ic.n_layers;i++) free(k->Ic[i]); free(k->Lc); free(k->Rc); free(k->Ic); free(k->kv_start); free(s->hist); diff --git a/c/tools/glm_fp8_emit.py b/c/tools/glm_fp8_emit.py index 4ae05d6..45bb304 100644 --- a/c/tools/glm_fp8_emit.py +++ b/c/tools/glm_fp8_emit.py @@ -66,6 +66,48 @@ def fp8_block_dequantize(w_fp8, scale): return qf * scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)[:O, :I] +def unfuse_experts(sd): + """Split HF's fused 3-D `experts.gate_up_proj` [E, 2*M, I] into per-expert 2-D + `experts.{e}.gate_proj` [M, I] + `experts.{e}.up_proj` [M, I], and + `experts.down_proj` [E, I, M] -> `experts.{e}.down_proj` [M_out, I]. + + The real GLM-5.2-FP8 checkpoint stores experts UNFUSED as per-expert 2-D tensors + (gate_proj, up_proj, down_proj), each with its own _scale_inv. HF's + GlmMoeDsaForCausalLM fuses gate+up into a single 3-D gate_up_proj for efficiency. + The converter (classify + ndim!=2 guard) and the C engine both expect the unfused + layout, so we split before saving. + + Idempotent: if experts are already unfused (no 3-D gate_up_proj), returns sd as-is. + EN: split HF's fused 3-D expert weights into the per-expert 2-D layout that the real + EN: checkpoint uses and the converter/engine expect. No-op if already unfused.""" + keys_to_remove = [] + new_entries = {} + for name, t in sd.items(): + if not name.endswith(".mlp.experts.gate_up_proj"): + continue + # prefix = everything before ".mlp.experts.gate_up_proj" + prefix = name[:-len(".mlp.experts.gate_up_proj")] + E, twoM, I = t.shape # [E, 2*intermediate, input] + M = twoM // 2 + for e in range(E): + new_entries[f"{prefix}.mlp.experts.{e}.gate_proj.weight"] = t[e, :M, :].contiguous() + new_entries[f"{prefix}.mlp.experts.{e}.up_proj.weight"] = t[e, M:, :].contiguous() + keys_to_remove.append(name) + # down_proj may be 3-D [E, I, M] in the fused form, or already per-expert + for name, t in sd.items(): + if not name.endswith(".mlp.experts.down_proj") or t.dim() != 3: + continue + prefix = name[:-len(".mlp.experts.down_proj")] + E = t.shape[0] + for e in range(E): + new_entries[f"{prefix}.mlp.experts.{e}.down_proj.weight"] = t[e].contiguous() + keys_to_remove.append(name) + for k in keys_to_remove: + sd.pop(k, None) + sd.update(new_entries) + return sd + + def state_dict_to_fp8(sd): """Converte uno state_dict HuggingFace nel layout FP8 del checkpoint reale: per ogni tensore quantizzabile 2-D scrive `{name}` (F8_E4M3) + `{name}_scale_inv` (F32); diff --git a/c/tools/make_glm_bench_model.py b/c/tools/make_glm_bench_model.py index 54a1a3f..5ec3457 100644 --- a/c/tools/make_glm_bench_model.py +++ b/c/tools/make_glm_bench_model.py @@ -22,7 +22,7 @@ import torch from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/ -from glm_fp8_emit import save_fp8_safetensors +from glm_fp8_emit import save_fp8_safetensors, unfuse_experts def build_config() -> GlmMoeDsaConfig: @@ -85,16 +85,6 @@ def main() -> None: output = Path(args.output) output.mkdir(parents=True, exist_ok=True) params = sum(p.numel() for p in model.parameters()) - if args.fp8: - n_fp8, n_tot = save_fp8_safetensors(model.state_dict(), output / "model.safetensors") - # save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo - # a mano (serve al converter e al motore C). EN: save_pretrained writes config.json; - # the FP8 path bypasses it, so write it manually (converter + C engine need it). - (output / "config.json").write_text(json.dumps(cfg.to_dict())) - print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) " - f"-> {output / 'model.safetensors'}") - else: - model.save_pretrained(output, safe_serialization=True, max_shard_size="4GB") model.to(args.device) prompt = [3, 14, 159, 26, 53, 58, 200, 11, 77, 240, 5, 99] @@ -103,6 +93,25 @@ def main() -> None: full = model.generate(ids, max_new_tokens=8, do_sample=False, use_cache=True)[0] logits = model(full.unsqueeze(0), use_cache=False).logits[0] + # Unfuse experts AFTER reference generation (model needs fused weights for + # forward/generate) but BEFORE saving — the real checkpoint and the converter + # + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors. + sd = model.state_dict() + unfuse_experts(sd) + + if args.fp8: + n_fp8, n_tot = save_fp8_safetensors(sd, output / "model.safetensors") + # save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo + # a mano (serve al converter e al motore C). EN: save_pretrained writes config.json; + # the FP8 path bypasses it, so write it manually (converter + C engine need it). + (output / "config.json").write_text(json.dumps(cfg.to_dict())) + print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) " + f"-> {output / 'model.safetensors'}") + else: + from safetensors.torch import save_file + save_file({k: v.contiguous() for k, v in sd.items()}, str(output / "model.safetensors")) + (output / "config.json").write_text(json.dumps(cfg.to_dict())) + ref = { "prompt_ids": prompt, "full_ids": full.cpu().tolist(), diff --git a/c/tools/make_glm_oracle.py b/c/tools/make_glm_oracle.py index 67ad5a5..b623faf 100644 --- a/c/tools/make_glm_oracle.py +++ b/c/tools/make_glm_oracle.py @@ -22,7 +22,7 @@ from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/ from glm_fp8_emit import (fp8_block_quantize, fp8_block_dequantize, keep_f32, - save_fp8_safetensors) + save_fp8_safetensors, unfuse_experts) ap = argparse.ArgumentParser() ap.add_argument("--fp8", action="store_true", @@ -84,7 +84,7 @@ with torch.no_grad(): if args.fp8: with torch.no_grad(): for n, p in model.named_parameters(): - if keep_f32(n, p) or p.dim() < 2: + if keep_f32(n, p) or p.dim() != 2: continue q, s = fp8_block_quantize(p) p.copy_(fp8_block_dequantize(q, s)) @@ -109,12 +109,19 @@ with torch.no_grad(): tf_pred = lg.argmax(-1).tolist() print("tf_pred:", tf_pred) +# Unfuse experts AFTER reference generation (model needs fused weights for +# forward/generate) but BEFORE saving — the real checkpoint and the converter +# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors. +sd = model.state_dict() +unfuse_experts(sd) + if args.fp8: - n_fp8, n_tot = save_fp8_safetensors(model.state_dict(), "glm_tiny/model.safetensors") + n_fp8, n_tot = save_fp8_safetensors(sd, "glm_tiny/model.safetensors") print(f"\nsaved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) " f"-> glm_tiny/model.safetensors") else: - model.save_pretrained("glm_tiny", safe_serialization=True) + from safetensors.torch import save_file + save_file({k: v.contiguous() for k, v in sd.items()}, "glm_tiny/model.safetensors") json.dump(cfg.to_dict(), open("glm_tiny/config.json", "w")) json.dump({"prompt_ids": prompt, "full_ids": full, "tf_pred": tf_pred}, open("ref_glm.json", "w")) print("saved: glm_tiny/ (weights + config) and ref_glm.json"