Files
esh-pfi-infrastructure/tools/mistral-small4-nvfp4/nvfp4_quant.py
T
vh dd3a5c93fd feat(tools): Mistral Small 4 NVFP4 build pipeline (quant + HF->native converter)
Quantize a HF-format Mistral Small 4 (Mistral3ForConditionalGeneration MoE) to
NVFP4 with the vision tower intact, then convert HF NVFP4 -> Mistral native so
vLLM can serve it (there is no HF Mistral4 serving path in any vLLM version).

Built + validated end-to-end on ana-ml2 for the abliterated character-model
successor (darkc0de/Mistral-Small-4-119B-2603-heretic): quant -> dry-run (clean
vs the official native NVFP4 reference) -> convert -> serve-test (loads on the
native loader, correct text, vision functional).

Converter scaffold came from worldtree-codex (bf16 bin maps + fused-expert
split); fixed here: NVFP4 layer regexes (keep the `model.` prefix) + non-mmap
shard reads (ZFS large-mmap ENOMEM). nvfp4_quant.py is local. README documents
the pipeline + every gotcha that cost a failed run. Homed here per operator
direction (not Worldtree).
2026-06-17 22:19:58 -07:00

27 lines
1.2 KiB
Python

import sys, torch
from transformers import AutoModelForImageTextToText, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.modeling.moe.linearize import load_quantizable_moe
MODEL, OUT = sys.argv[1], sys.argv[2]
NSAMPLES = int(sys.argv[3]) if len(sys.argv) > 3 else 64
# Mirror Mistral official NVFP4: quantize expert FFN (routed via MoE-linearize + shared),
# keep bf16: vision, all MLA attention, MoE router gate, embeddings, lm_head
IGNORE = [
"re:.*lm_head.*", "re:.*embed_tokens.*",
"re:.*vision_tower.*", "re:.*multi_modal_projector.*",
"re:.*self_attn.*", r"re:.*mlp\.gate$",
]
recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=IGNORE)
with load_quantizable_moe(AutoModelForImageTextToText):
model = AutoModelForImageTextToText.from_pretrained(MODEL, dtype="auto", device_map="cpu")
oneshot(model=model, recipe=recipe, processor=AutoTokenizer.from_pretrained(MODEL), dataset="ultrachat_200k",
splits={"calibration": f"train_gen[:{NSAMPLES}]"},
num_calibration_samples=NSAMPLES, max_seq_length=512)
model.save_pretrained(OUT, save_compressed=True)
print("SAVED", OUT)