Generators: FP8 emission + expert unfusing + download tool

c/download_fp8.py

@@ -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")

c/tools/glm_fp8_emit.py

@@ -0,0 +1,137 @@
+"""Helper: salva pesi in FP8 e4m3 + scale a blocchi 128x128, nello STESSO layout del
+checkpoint reale GLM-5.2-FP8 che `convert_fp8_to_int4.py` legge.
+
+Layout (deve combaciare col `dequant()` del converter, convert_fp8_to_int4.py:164-169):
+  - `name`            F8_E4M3  [O, I]
+  - `name_scale_inv`  F32      [ceil(O/128), ceil(I/128)]   (NOTA: '_scale_inv', underscore)
+  dequant: W = q.float() * scale.repeat_interleave(128,0).repeat_interleave(128,1)[:O,:I]
+
+Convenzione FBGEMM/TransformerEngine: scale = amax(blocco)/448 (448 = max e4m3),
+si MEMORIZZA il valore e si MOLTIPLICA in dequant. Malgrado il nome "_scale_inv" il
+checkpoint memorizza la scala (non il reciproco): e' un MOLTIPLIER.
+
+EN: Helper that writes weights as FP8 e4m3 with 128x128 block scales, in the SAME layout
+EN: as the real GLM-5.2-FP8 checkpoint that `convert_fp8_to_int4.py` reads.
+EN: FBGEMM/TransformerEngine convention: scale = amax(block)/448, stored (not its
+EN: reciprocal) and MULTIPLIED on dequant. Despite the name "_scale_inv" it is a multiplier.
+"""
+import torch
+
+E4M3_MAX = 448.0   # max valore rappresentabile in float8_e4m3fn / max representable value
+BLOCK = 128         # granularita' delle scale a blocchi del checkpoint FP8 / FP8 block scale granularity
+
+
+def keep_f32(name, t):
+    """Stesso set F32 di `classify()` in convert_fp8_to_int4.py (norme, router, bias 1-D).
+    Tutti gli altri tensori 2-D vengono quantizzati FP8 (attn/mlp/shared/expert/embed/lm_head).
+    EN: Same F32 set as the converter's classify(): norms, router, 1-D biases. All other 2-D
+    EN: tensors are FP8-quantized (attn/mlp/shared/expert/embed/lm_head)."""
+    if t.dim() < 2:
+        return True                                  # bias 1-D, e_score_correction_bias
+    if name.endswith("e_score_correction_bias"):
+        return True
+    if name.endswith("mlp.gate.weight"):
+        return True                                  # router (NON gate_proj): tenuto F32 / kept F32
+    if name.endswith("norm.weight") or name == "model.norm.weight":
+        return True                                  # RMSNorm
+    return False
+
+
+def fp8_block_quantize(w):
+    """w: [O,I] f32 -> (w_fp8 float8_e4m3fn [O,I], scale_inv f32 [ceil(O/128),ceil(I/128)]).
+    Identica matematica al `--selftest` del converter (scale = amax(blocco)/448). Padda a
+    multipli di 128 internamente (gli zeri non alzano l'amax) e fa slice al risultato.
+    EN: same math as the converter's --selftest. Pads to 128 multiples internally (zeros do
+    EN: not raise amax), slices the result back to [O,I]."""
+    O, I = w.shape
+    nbO, nbI = (O + BLOCK - 1) // BLOCK, (I + BLOCK - 1) // BLOCK
+    Op, Ip = nbO * BLOCK, nbI * BLOCK
+    wpad = torch.zeros(Op, Ip, dtype=torch.float32, device=w.device)
+    wpad[:O, :I] = w
+    wb = wpad.view(nbO, BLOCK, nbI, BLOCK)               # [nbO, BLOCK, nbI, BLOCK]
+    amax = wb.abs().amax(dim=(1, 3))                      # [nbO, nbI]
+    scale = amax / E4M3_MAX                               # FBGEMM/TE: memorizza la scala / store the scale
+    scale = torch.where(scale == 0, torch.ones_like(scale), scale)  # blocco tutto-zero -> no div0
+    scale = scale.to(torch.float32)
+    q = (wpad / scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)).clamp(-E4M3_MAX, E4M3_MAX)
+    w_fp8 = q.to(torch.float8_e4m3fn)
+    return w_fp8[:O, :I].contiguous(), scale.contiguous()
+
+
+def fp8_block_dequantize(w_fp8, scale):
+    """Esatto inverso di fp8_block_quantize, e identico al `dequant()` del converter.
+    EN: exact inverse of fp8_block_quantize, identical to the converter's dequant()."""
+    O, I = w_fp8.shape
+    qf = w_fp8.to(torch.float32)
+    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);
+    norme/router/bias e qualsiasi tensore NON 2-D (es. pesi MLA impaccati 3-D) restano nel
+    dtype originale. Questo rispecchia il guard `w.ndim != 2 -> f32` del converter
+    (convert_fp8_to_int4.py:184). EN: builds the real-checkpoint FP8 layout. Only exactly-2-D
+    tensors are FP8-quantized; anything else (1-D, 3-D packed MLA weights, ...) is kept, exactly
+    like the converter's `ndim != 2 -> f32` guard."""
+    out = {}
+    for name, t in sd.items():
+        if keep_f32(name, t) or t.dim() != 2:
+            out[name] = t                               # f32 / 1-D / 3-D+: tieni / keep
+        else:
+            w_fp8, scale = fp8_block_quantize(t.float())
+            out[name] = w_fp8
+            out[name + "_scale_inv"] = scale
+    return out
+
+
+def save_fp8_safetensors(sd, path):
+    """Quantizza a blocchi FP8 e salva in un singolo safetensors leggibile dal converter
+    via `--indir`. EN: block-quantize to FP8 and save a single safetensors for the converter."""
+    from safetensors.torch import save_file
+    out = state_dict_to_fp8(sd)
+    save_file({k: v.contiguous() for k, v in out.items()}, str(path))
+    n_fp8 = sum(1 for v in out.values() if v.dtype == torch.float8_e4m3fn)
+    return n_fp8, len(out)