woolcoxm
2c6946c478
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).
96 lines
5.0 KiB
Python
96 lines
5.0 KiB
Python
"""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)
|