test-models: add --fp8 emission for the FP8->int4 converter test path

Both test-model generators (make_glm_oracle.py, make_glm_bench_model.py) can now emit weights as FP8 e4m3 + 128x128 block scale_inv, in the same layout as the real GLM-5.2-FP8 checkpoint. This lets convert_fp8_to_int4.py exercise its FP8->int4 dequant path on a local fixture without the 379 GB download. - New shared helper glm_fp8_emit.py: FP8 block quantize/dequantize (FBGEMM/TE scale=amax/448 convention) + state_dict emitter. Only exactly-2-D tensors are quantized; 1-D/3-D and norms/router/e_score_correction_bias are kept as f32, mirroring the converter's classify() + ndim!=2 guard. - make_glm_bench_model.py: opt-in --fp8 writes model.safetensors in FP8 layout (config.json written explicitly since the FP8 path bypasses save_pretrained); manifest gains a 'format' field. Default bf16 behavior unchanged. - make_glm_oracle.py: opt-in --fp8 round-trips quantizable weights through FP8 before computing ref_glm.json, so the reference reflects exactly the FP8 model the converter ingests. Default bf16 oracle contract unchanged. Verified end-to-end: FP8 model -> converter --indir -> int4 U8 + .qs F32 output, bit-identical dequant between helper and converter (maxdiff 0.0).
2026-07-15 01:02:21 -04:00
parent 21eb86a0dd
commit 2c6946c478
1 changed files with 95 additions and 0 deletions
@@ -0,0 +1,95 @@
+"""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 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)