diskio: KV write batching + persistent handle; generators: unfuse experts

Two independent fixes validated end-to-end on fresh fixtures:

1. KV cache disk I/O (issue_diskio.md opportunities #1 + #4):
   - kv_disk_append: fopen/fclose every turn -> persistent FILE* kept open
     for the engine lifetime, lazy open on first append, closed in
     serve_ctx_free. Eliminates per-turn handle creation overhead.
   - kv_disk_append: ~157 small fwrites per position -> one contiguous
     record memcpy'd into a staging buffer then a single fwrite per
     position. The staging buffer grows on demand via realloc.
   - kv_disk_truncate: closes the persistent handle before truncating
     so the file actually shrinks on disc, then reopens lazily.
   - KVState gains disk_fp, disk_buf, disk_buf_cap fields.
   - Verified: serve-mode round-trip, write 11 tokens then reload and
     resume with no re-prefill, then append 8 more and reload to 19.

2. Expert weight unfusing in test-model generators:
   - The real GLM-5.2-FP8 checkpoint stores routed experts UNFUSED as
     per-expert 2-D tensors, each with its own _scale_inv. HF fuses
     gate+up into a single 3-D gate_up_proj for compute efficiency.
   - The converter and C engine both expect the unfused layout. The
     fused 3-D tensors were silently skipped by the converter, and the
     engine crashed with missing-tensor errors.
   - New unfuse_experts in glm_fp8_emit.py splits gate_up_proj and
     down_proj into per-expert 2-D tensors. Called after reference
     generation but before saving, in both generators, both FP8 and bf16.
   - Also fixed: make_glm_oracle.py FP8 round-trip guard used p.dim()<2
     which let 3-D fused experts through and crashed fp8_block_quantize.
     Changed to p.dim()!=2 to match the converter ndim!=2 guard.

Validated full chain on fresh fixtures:
  generator --fp8 -> 570 e4m3 tensors + 629 scale_inv, was 90 when fused
  converter --group-size 0  -> per-row int4 fmt=2, engine loads clean
  converter --group-size 128 -> grouped int4 fmt=4, 8-16x more scales,
    engine loads clean, fmt=4 auto-detected in both mmap and slab paths
  dequant error: grouped 1.14-1.22x lower than per-row vs FP8 source
This commit is contained in:
woolcoxm
2026-07-15 01:40:02 -04:00
parent e71d4fbe29
commit 71c262ce1a
4 changed files with 127 additions and 27 deletions
+11 -4
View File
@@ -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"