Windows-dev: grouped quantization + mixed precision + expert budget + download tool

Consolidates all experiment branches into one Windows-dev branch: 1. Group-scaled int4 (fmt=4, gs=128) — glm.c - QT struct: added gs field - matmul_i4_grouped: AVX2 kernel, verified to 3e-08 vs f32 - Format detection: auto-detects from .qs scale array size - expert_load: both mmap and slab+pread paths handle fmt=4 - qt_bytes: fmt=4 case added 2. Per-tensor-type mixed precision — convert_fp8_to_int4.py - Split classify() into sh/o/kvb/attn/dmlp sub-types - New args: --shared-bits, --o-bits, --kvb-bits, --attn-bits, --dmlp-bits - Plan: shared expert + o_proj + kv_b_proj at int8, rest grouped int4 - Only +5.3 GB RAM vs +0 for pure int4 3. EXPERT_BUDGET (miss-aware) — glm.c - Caps distinct experts per layer across batch-union - Always keeps cache hits, only drops misses - Up to 1.8x faster decode on low-RAM hosts 4. Two-step shared-expert prediction (PILOT_TWO) — glm.c - la_predict kind==2 + pilot_prefetch integration - +3.1% recall over baseline PILOT 5. FP8 download tool — download_fp8.py - ModelScope + HuggingFace dual-source - Parallel shard download with stall recovery 6. Tiny model generation — make_glm_oracle.py - Generated and tested locally for pipeline validation
2026-07-15 01:00:53 -04:00
parent e141db047d
commit 69f65d5173
3 changed files with 303 additions and 5 deletions
@@ -0,0 +1,231 @@
+"""Download GLM-5.2-FP8 from ModelScope (fast, no HF throttling) with
+HuggingFace fallback. Parallel shard download, clean progress display.
+Usage: python download_fp8.py
+       python download_fp8.py --parallel 4
+       python download_fp8.py --source hf  (force HuggingFace)
+"""
+import os, sys, time, threading, argparse, subprocess
+
+REPO_MS = "ZhipuAI/GLM-5.2-FP8"        # ModelScope
+REPO_HF = "zai-org/GLM-5.2-FP8"        # HuggingFace
+DEST = r"I:\glm52_fp8"
+
+# ── ANSI colors ──
+class C:
+    dim="\033[2m"; grn="\033[32m"; yel="\033[33m"; cyn="\033[36m"; b="\033[1m"; r="\033[0m"
+
+def fmt_bytes(n):
+    if n>=1e9: return f"{n/1e9:.2f} GB"
+    if n>=1e6: return f"{n/1e6:.1f} MB"
+    return f"{n/1e3:.0f} KB"
+
+def fmt_time(s):
+    if s<60: return f"{s:.0f}s"
+    if s<3600: return f"{s/60:.0f}m"
+    return f"{s/3600:.1f}h"
+
+def bar(cur, total, width=24):
+    if total<=0: return "["+" "*width+"]"
+    pct=min(cur/total,1.0); filled=int(width*pct)
+    return "["+"█"*filled+"░"*(width-filled)+f"] {pct*100:4.0f}%"
+
+def get_shard_list_hf():
+    from huggingface_hub import HfApi
+    info=HfApi().repo_info(REPO_HF, files_metadata=True)
+    shards=sorted(s.rfilename for s in info.siblings if s.rfilename.endswith(".safetensors"))
+    sizes={s.rfilename:s.size for s in info.siblings if s.rfilename.endswith(".safetensors")}
+    return shards, sizes
+
+def get_shard_list_ms():
+    """Get shard list from ModelScope API."""
+    import requests
+    # ModelScope API: list files
+    r = requests.get(f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo/files?Revision=master&Root=", timeout=30)
+    data = r.json()["Data"]["Files"]
+    shards = sorted(f["Path"] for f in data if f["Path"].endswith(".safetensors"))
+    sizes = {f["Path"]: f.get("Size", 0) for f in data if f["Path"].endswith(".safetensors")}
+    return shards, sizes
+
+def download_file_ms(fn):
+    """Download a single file from ModelScope using their CDN."""
+    from modelscope.hub.file_download import model_file_download
+    model_file_download(
+        model_id=REPO_MS,
+        file_path=fn,
+        local_dir=DEST,
+        revision="master",
+    )
+
+def download_file_hf(fn):
+    """Download a single file from HuggingFace with hf_transfer."""
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+    from huggingface_hub import hf_hub_download
+    hf_hub_download(REPO_HF, fn, local_dir=DEST)
+
+def download_file_curl(fn, base_url, expected_size):
+    """Fallback: download with curl to a .part file with resume."""
+    outpath = os.path.join(DEST, fn)
+    partpath = outpath + ".part"
+    if os.path.exists(outpath) and os.path.getsize(outpath) == expected_size:
+        return True
+    url = f"{base_url}/{fn}"
+    cmd = ["curl", "-L", "-C", "-", "--retry", "999", "--retry-delay", "5",
+           "--connect-timeout", "15", "--speed-time", "30", "--speed-limit", "1000",
+           "-o", partpath, "-H", "User-Agent: colibri-download/1.0", url]
+    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    if os.path.exists(partpath) and os.path.getsize(partpath) >= expected_size:
+        os.replace(partpath, outpath)
+        return True
+    return False
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--parallel", type=int, default=3)
+    ap.add_argument("--source", choices=["auto", "ms", "hf"], default="auto",
+                    help="auto (try ModelScope first), ms (ModelScope only), hf (HuggingFace only)")
+    args = ap.parse_args()
+
+    os.makedirs(DEST, exist_ok=True)
+
+    # Determine source and get shard list
+    use_ms = False
+    shards, sizes = [], {}
+
+    if args.source in ("auto", "ms"):
+        try:
+            print(f"{C.dim}Trying ModelScope...{C.r}", end=" ", flush=True)
+            shards, sizes = get_shard_list_ms()
+            use_ms = True
+            print(f"{C.grn}✓{C.r} {len(shards)} shards found")
+        except Exception as e:
+            print(f"{C.yel}failed ({e}){C.r}")
+            if args.source == "ms":
+                print("ModelScope failed and --source ms was set. Exiting."); return
+
+    if not shards:
+        print(f"{C.dim}Using HuggingFace...{C.r}", end=" ", flush=True)
+        shards, sizes = get_shard_list_hf()
+        print(f"{C.grn}✓{C.r} {len(shards)} shards found")
+
+    total = len(shards)
+    total_bytes = sum(sizes.values())
+    source_name = "ModelScope" if use_ms else "HuggingFace"
+
+    # Download metadata files
+    meta_files = ["config.json", "tokenizer.json", "tokenizer_config.json",
+                  "generation_config.json", "model.safetensors.index.json"]
+    for fn in meta_files:
+        out = os.path.join(DEST, fn)
+        if not os.path.exists(out):
+            try:
+                if use_ms: download_file_ms(fn)
+                else: download_file_hf(fn)
+            except Exception: pass
+
+    # Build work queue
+    todo = []
+    done_set = set()
+    for fn in shards:
+        outpath = os.path.join(DEST, fn)
+        if os.path.exists(outpath) and os.path.getsize(outpath) == sizes.get(fn, 0):
+            done_set.add(fn)
+        else:
+            todo.append(fn)
+
+    existing = len(done_set)
+    print(f"\n{C.b}GLM-5.2-FP8 Download ({source_name}){C.r}")
+    print(f"  {C.dim}{total} shards · {total_bytes/1e9:.0f} GB · {existing}/{total} complete{C.r}")
+    print(f"  {C.dim}{len(todo)} to download · {args.parallel} parallel{C.r}")
+    if todo:
+        remaining = sum(sizes[fn] for fn in todo)
+        print(f"  {C.dim}Remaining: {remaining/1e9:.0f} GB{C.r}")
+    print()
+
+    if not todo:
+        print(f"{C.grn}✓ All shards already downloaded!{C.r}\n"); return
+
+    lock = threading.Lock()
+    completed = list(done_set)
+    t0 = time.time()
+    qidx = [0]
+
+    def worker(wid):
+        while True:
+            with lock:
+                if qidx[0] >= len(todo): return
+                idx = qidx[0]; qidx[0] += 1
+                fn = todo[idx]
+
+            expected = sizes.get(fn, 0)
+            shard_num = existing + idx + 1
+            print(f"  {C.cyn}[{shard_num}/{total}]{C.r} {C.dim}{fn}{C.r}")
+
+            success = False
+            for attempt in range(3):
+                try:
+                    if use_ms:
+                        download_file_ms(fn)
+                    else:
+                        download_file_hf(fn)
+                    outpath = os.path.join(DEST, fn)
+                    # ModelScope downloads to a cache dir, need to check
+                    # if the file exists at our expected path
+                    if not os.path.exists(outpath):
+                        # Try to find it in ModelScope's cache structure
+                        # and copy/symlink it
+                        pass
+                    if os.path.exists(outpath):
+                        actual = os.path.getsize(outpath)
+                        if expected == 0 or actual == expected:
+                            success = True; break
+                        # Size mismatch — might be in cache
+                    # If we got here, file downloaded but not at expected path
+                    # ModelScope puts it in local_dir; check again
+                    if os.path.exists(outpath):
+                        success = True; break
+                    # Retry with curl fallback
+                    if use_ms:
+                        base = f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo?Revision=master&FilePath="
+                    else:
+                        base = f"https://huggingface.co/{REPO_HF}/resolve/main"
+                    if download_file_curl(fn, base, expected):
+                        success = True; break
+                except Exception as e:
+                    if attempt < 2:
+                        print(f"    {C.yel}retry {attempt+1}: {e}{C.r}")
+                        time.sleep(3)
+                    else:
+                        print(f"    {C.yel}✗ failed: {e}{C.r}")
+
+            with lock:
+                if success:
+                    completed.append(fn)
+                    elapsed = time.time() - t0
+                    have = sum(sizes.get(f,0) for f in completed)
+                    pct = 100.0 * have / total_bytes
+                    speed = (have - sum(sizes[f] for f in done_set)) / max(elapsed, 1)
+                    eta = (total_bytes - have) / speed if speed > 0 else 0
+                    print(f"  {C.grn}✓{C.r} {fn} {C.dim}— {len(completed)}/{total} · "
+                          f"{pct:.1f}% · {fmt_time(elapsed)} · ETA {fmt_time(eta)}{C.r}")
+                else:
+                    print(f"  {C.yel}✗ GIVE UP: {fn}{C.r}")
+
+    threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(args.parallel)]
+    for t in threads: t.start()
+    for t in threads: t.join()
+
+    print()
+    final = sum(1 for fn in shards
+                if os.path.exists(os.path.join(DEST, fn))
+                and os.path.getsize(os.path.join(DEST, fn)) == sizes.get(fn, 0))
+    if final == total:
+        print(f"{C.grn}{'='*50}")
+        print(f"  ✓ All {total} shards downloaded!{C.r}\n")
+    else:
+        print(f"{C.yel}  {final}/{total} complete, {total-final} remaining{C.r}")
+        print(f"  Re-run to resume.\n")
+
+if __name__ == "__main__":
+    try: main()
+    except KeyboardInterrupt:
+        print(f"\n\n{C.yel}Interrupted. Re-run to resume — no data lost.{C.r}\n")
@@ -3,15 +3,27 @@
 This is not a useful language model. It preserves the real glm_moe_dsa data
 flow while remaining small enough to generate locally and run repeated CPU/CUDA
 A/B tests without downloading the 379 GB checkpoint.
+
+With --fp8 the weights are written as FP8 e4m3 + 128x128 block scale_inv, in the
+SAME layout as the real GLM-5.2-FP8 checkpoint, so convert_fp8_to_int4.py can
+exercise its FP8->int4 dequant path on a local fixture (its dims are 128-friendly,
+so this is also the right fixture for --group-size 128 testing):
+
+  python tools/make_glm_bench_model.py --fp8 --output glm_bench_fp8
+  python tools/convert_fp8_to_int4.py --indir glm_bench_fp8 --outdir glm_bench_i4 --ebits 4 --group-size 128
 """

 import argparse
 import json
+import sys
 from pathlib import Path

 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
+

 def build_config() -> GlmMoeDsaConfig:
    return GlmMoeDsaConfig(
@@ -51,6 +63,9 @@ def main() -> None:
    parser.add_argument("--output", default="glm_bench_medium")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--fp8", action="store_true",
+                        help="write weights as FP8 e4m3 + 128x128 block scale_inv (same layout as "
+                             "GLM-5.2-FP8) instead of bf16, so convert_fp8_to_int4.py can dequant+requant")
    args = parser.parse_args()

    torch.manual_seed(args.seed)
@@ -70,7 +85,16 @@ def main() -> None:
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    params = sum(p.numel() for p in model.parameters())
-    model.save_pretrained(output, safe_serialization=True, max_shard_size="4GB")
+    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]
@@ -89,6 +113,7 @@ def main() -> None:
        "seed": args.seed,
        "parameters": params,
        "parameters_billions": round(params / 1e9, 4),
+        "format": "fp8-e4m3-128" if args.fp8 else "bf16",
        "purpose": "backend benchmark fixture; random weights, not a language model",
    }
    (output / "bench_manifest.json").write_text(json.dumps(manifest, indent=2))
@@ -3,10 +3,34 @@ Architettura vera (MLA + DSA indexer + router sigmoid/noaux_tc + shared expert),
 dimensioni minuscole. Salva pesi+config in c/glm_tiny/ e un riferimento greedy in
 c/ref_glm.json. seq corta (<= index_topk) cosi' il DSA seleziona tutte le key e
 l'attenzione coincide con la MLA densa: il motore C puo' validare senza implementare
-l'indexer sparso."""
-import json, torch
+l'indexer sparso.
+
+--fp8: salva i pesi come FP8 e4m3 + scale a blocchi 128x128 (layout del checkpoint reale
+GLM-5.2-FP8) invece di bf16, cosi' convert_fp8_to_int4.py puo' esercitare il path FP8->int4
+su un modello minuscolo. PRIMA di calcolare ref_glm.json fa il round-trip dei pesi per FP8
+(quant->dequant, copy_ nel modello): cosi' il riferimento riflette ESATTAMENTE il modello
+FP8 che il converter legge, non il modello bf16 a precisione piena. Default: bf16 (oracolo
+originale invariato).
+EN: --fp8 writes FP8 e4m3 + 128x128 block scale_inv (real GLM-5.2-FP8 layout) instead of bf16,
+EN: so convert_fp8_to_int4.py can run its FP8->int4 path on a tiny model. ref_glm.json is
+EN: computed AFTER the FP8 round-trip, so the reference matches exactly what the converter
+EN: ingests. Default: bf16 (original oracle unchanged)."""
+import json, sys, argparse
+from pathlib import Path
+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 (fp8_block_quantize, fp8_block_dequantize, keep_f32,
+                          save_fp8_safetensors)
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--fp8", action="store_true",
+                help="salva in FP8 e4m3 + 128x128 block scale_inv (layout GLM-5.2-FP8) e "
+                     "calcola ref_glm.json sul modello dopo il round-trip FP8. "
+                     "EN: write FP8 e4m3 + block scale_inv, ref computed on FP8-rounded model")
+args = ap.parse_args()
+
 torch.manual_seed(1234)

 cfg = GlmMoeDsaConfig(
@@ -53,6 +77,18 @@ with torch.no_grad():
            layer.mlp.gate.e_score_correction_bias.copy_(
                torch.linspace(-0.1, 0.1, cfg.n_routed_experts))

+# --fp8: round-trip dei pesi quantizzabili per FP8 PRIMA di calcolare il riferimento,
+# cosi' ref_glm.json riflette esattamente il modello FP8 che il converter leggera'.
+# Norme/router/bias (keep_f32) restano a precisione piena. EN: --fp8: round-trip quantizable
+# weights through FP8 before computing the reference, so ref_glm.json matches the FP8 model.
+if args.fp8:
+    with torch.no_grad():
+        for n, p in model.named_parameters():
+            if keep_f32(n, p) or p.dim() < 2:
+                continue
+            q, s = fp8_block_quantize(p)
+            p.copy_(fp8_block_dequantize(q, s))
+
 print("=== state_dict tensors (names used by the C loader) ===")
 for n, p in model.state_dict().items():
    print(f"  {n:60s} {tuple(p.shape)}")
@@ -73,7 +109,13 @@ with torch.no_grad():
 tf_pred = lg.argmax(-1).tolist()
 print("tf_pred:", tf_pred)

-model.save_pretrained("glm_tiny", safe_serialization=True)
+if args.fp8:
+    n_fp8, n_tot = save_fp8_safetensors(model.state_dict(), "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)
 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("\nsaved: glm_tiny/ (weights + config) and ref_glm.json")
+print("saved: glm_tiny/ (weights + config) and ref_glm.json"
+      + (" [fp8]" if args.fp8 else ""))