| """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, |
| ) |
| model = Mistral4ForCausalLM(config).to(torch.bfloat16).eval() |
|
|
| fp8 = FineGrainedFP8Config(weight_block_size=(32, 32)) |
| |
| 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]) |
|
|
