File size: 3,918 Bytes
ac53c7a 1236014 ac53c7a 3a23a8f ac53c7a bc77947 ac53c7a 4ec61d8 3a23a8f ac53c7a 4ec61d8 ac53c7a 3a23a8f ac53c7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Tiny random Mistral4 with fp8 (attention + routed experts) and nvfp4 (shared experts) layers, saved pre-quantized
with a per-layer quantization plan in config.json. Usage: PYTHONPATH=<worktree>/src python make.py <out_dir>"""
import fnmatch
import json
import sys
import torch
from safetensors.torch import save_file
from transformers import FineGrainedFP8Config, Mistral4Config, Mistral4ForCausalLM, NVFP4Config
from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize
from transformers.integrations.mxfp4 import FP4_VALUES as FP4_E2M1_VALUES
from transformers.integrations.nvfp4 import SF_VEC_SIZE, NVFP4Dequantize
out = sys.argv[1]
torch.manual_seed(0)
config = Mistral4Config(
vocab_size=128, hidden_size=256, intermediate_size=512, moe_intermediate_size=256, num_hidden_layers=2,
num_attention_heads=2, num_key_value_heads=2, n_routed_experts=4, num_experts_per_tok=2, n_shared_experts=1,
q_lora_rank=64, kv_lora_rank=64, qk_rope_head_dim=32, qk_nope_head_dim=32, v_head_dim=32, first_k_dense_replace=0,
) # fmt: skip
model = Mistral4ForCausalLM(config).to(torch.bfloat16).eval()
fp8 = FineGrainedFP8Config(weight_block_size=(32, 32))
# The experts get the real-world 128x128 blocks: the grouped kernel needs them
plan = {"*.self_attn.*_proj*": fp8, "*.mlp.experts.*": FineGrainedFP8Config(weight_block_size=(128, 128)), "*.mlp.shared_experts.*": NVFP4Config()}
def quantize_nvfp4(weight):
"Packed e2m1 (low nibble first), row-major e4m3 block scales of 16, and the per-tensor global scale"
weight = weight.float()
global_scale = 448.0 * 6.0 / weight.abs().amax()
blocks = weight.view(weight.shape[0], -1, SF_VEC_SIZE)
scales = (blocks.abs().amax(-1) / 6.0 * global_scale).clamp(max=448.0).to(torch.float8_e4m3fn)
scaled = blocks * global_scale / scales.float().clamp(min=1e-12).unsqueeze(-1)
codes = (scaled.unsqueeze(-1) - torch.tensor(FP4_E2M1_VALUES)).abs().argmin(-1).to(torch.uint8).flatten(1)
packed = codes[:, 0::2] | (codes[:, 1::2] << 4)
return {"weight": packed, "weight_sf_rowmajor": scales.view(torch.uint8), "weight_global_scale": global_scale.reshape(1)}
state_dict, report = {}, []
for name, tensor in model.state_dict().items():
pattern = next((p for p in plan if fnmatch.fnmatch(name, p)), None)
if pattern is None or tensor.ndim < 2:
state_dict[name] = tensor.contiguous()
continue
prefix = name.removesuffix("weight")
if isinstance(plan[pattern], NVFP4Config):
quantized = quantize_nvfp4(tensor)
dense = NVFP4Dequantize(None).convert({"weight": quantized["weight"], "weight_sf_rowmajor": quantized["weight_sf_rowmajor"], "weight_global_scale": quantized["weight_global_scale"]}, model=model, full_layer_name=name)["weight"][0]
state_dict.update({prefix + key: value for key, value in quantized.items()})
else:
hf_quantizer = type("Q", (), {"quantization_config": plan[pattern]})()
quantized = Fp8Quantize(hf_quantizer).convert({name: tensor})
scale_key = next(key for key in quantized if key != name)
dense = Fp8Dequantize(hf_quantizer).convert({"weight": quantized[name], "weight_scale_inv": quantized[scale_key]}, full_layer_name=name)[name]
state_dict.update(quantized)
error = ((dense.float() - tensor.float()).norm() / tensor.float().norm()).item()
report.append(f"{type(plan[pattern]).__name__:22} {name:52} rel err {error:.3f}")
assert error < 0.2, (name, error)
print("\n".join(report))
model.config.quantization_config = {pattern: cfg.to_dict() for pattern, cfg in plan.items()}
model.config.save_pretrained(out)
save_file({k: v.contiguous() for k, v in state_dict.items()}, f"{out}/model.safetensors", metadata={"format": "pt"})
torch.save({k: v for k, v in model.state_dict().items()}, f"{out}/dense_reference.pt")
print(json.dumps(model.config.quantization_config, indent=1)[:800])
|