From dd3a5c93fdbf372bbfdf17959c3c70b8d9115729 Mon Sep 17 00:00:00 2001 From: Vuong Hoang Date: Wed, 17 Jun 2026 22:19:58 -0700 Subject: [PATCH] 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). --- tools/mistral-small4-nvfp4/README.md | 94 ++++ .../convert_hf_to_native.py | 525 ++++++++++++++++++ tools/mistral-small4-nvfp4/nvfp4_quant.py | 26 + 3 files changed, 645 insertions(+) create mode 100644 tools/mistral-small4-nvfp4/README.md create mode 100644 tools/mistral-small4-nvfp4/convert_hf_to_native.py create mode 100644 tools/mistral-small4-nvfp4/nvfp4_quant.py diff --git a/tools/mistral-small4-nvfp4/README.md b/tools/mistral-small4-nvfp4/README.md new file mode 100644 index 0000000..c9884c0 --- /dev/null +++ b/tools/mistral-small4-nvfp4/README.md @@ -0,0 +1,94 @@ +# Mistral Small 4 → NVFP4 (vision-intact) build tooling + +Quantize a **HF-format** `Mistral3ForConditionalGeneration` checkpoint +(Mistral Small 4, 119B-total / 6.5B-active MoE) to **NVFP4** with the vision +tower intact, then convert it to **Mistral native format** so vLLM can serve it. + +Built for the abliterated character-model successor +(`darkc0de/Mistral-Small-4-119B-2603-heretic`), validated end-to-end on ana-ml2 +(2026-06-17). The official `mistralai/Mistral-Small-4-119B-2603-NVFP4` is the +naming reference the converter diffs against. + +## Why both a quant *and* a convert step + +vLLM serves Mistral Small 4 **only** through its native loader +(`--config-format mistral --load-format mistral --tokenizer-mode mistral`) — +there is no HF `Mistral4` serving path in any vLLM version. But `llm-compressor` +quantizes the **HF** checkpoint. So the pipeline is: + +``` +HF bf16 ──quant──▶ HF NVFP4 ──convert──▶ native NVFP4 ──serve──▶ vLLM + nvfp4_quant.py convert_hf_to_native.py (native loader) +``` + +## Pipeline (on ana-ml2, `/tank/aimodels/quant-work`, in a uv venv) + +```bash +# 0. Pull the HF bf16 source (e.g. via huggingface-cli download). + +# 1. Quantize HF bf16 -> HF NVFP4 (~65 GB out). GPU0 for compute, CPU-resident model. +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True CUDA_VISIBLE_DEVICES=0 \ + python nvfp4_quant.py heretic-nvfp4 128 + +# 2. Dry-run the native convert (name-map check vs the official native reference). +python convert_hf_to_native.py --format nvfp4 \ + --hf-dir heretic-nvfp4 \ + --native-ref-dir \ + --out-dir heretic-native-nvfp4 --dry-run +# Expect: unmapped=0, missing_from_output=0, extra_in_output=0. + +# 3. Full native convert (~65 GB out, 5 shards). +python convert_hf_to_native.py --format nvfp4 \ + --hf-dir heretic-nvfp4 --native-ref-dir \ + --out-dir heretic-native-nvfp4 --max-shard-size-gb 15 + +# 4. Serve-test on vLLM (native loader, v0.22.0 = last vision-working pin). +docker run -d --name heretic-serve-test --ipc host --gpus '"device=0"' \ + -p 8099:8000 -v $PWD/heretic-native-nvfp4:/model:ro \ + vllm/vllm-openai:v0.22.0 /model --served-model-name heretic-test \ + --host 0.0.0.0 --port 8000 \ + --tokenizer-mode mistral --config-format mistral --load-format mistral \ + --tensor-parallel-size 1 --gpu-memory-utilization 0.93 \ + --max-model-len 16384 --attention-backend TRITON_MLA --max-num-seqs 8 --dtype auto +``` + +## Gotchas (each one cost a failed run) + +- **`device_map="cpu"`, not `"auto"`** in the quant. `auto` fills GPU0 with the + 205 GB model → OOM during MoE un-fusing; constraining with `max_memory` then + offloads experts to the *meta* device, which `copy_from_experts_module` can't + `.copy_()` (`Cannot copy out of meta tensor`). CPU-resident keeps every tensor + real; the sequential pipeline still onloads each layer to GPU0 for compute. +- **Non-mmap shard reads** in the converter. `safetensors.safe_open()` mmaps the + whole shard; on `/tank` (ZFS) a 50 GB shard mmap ENOMEMs regardless of free RAM + (MAP_SHARED never consults the commit limit). `read_tensor` reads with plain + `read()` + `safetensors.torch.load(bytes)`, caching one shard at a time — the + copy loop is sorted by shard so the cache doesn't thrash. +- **`vm.overcommit_memory=1`** on ana-ml2 (now durable — see + `playbooks/ana-ml2-overcommit-memory.yaml`). overcommit=0 + zero swap caps the + CommitLimit at ~RAM/2; the resident vLLM services eat the headroom and large + allocations fail despite free RAM. +- **NVFP4 output keeps the `model.` prefix.** llm-compressor's NVFP4 tensor names + are `model.language_model.model.layers.N...` (same prefix as bf16) — they are + *not* prefix-shifted. The only NVFP4 difference vs bf16 is per-expert-quantized + (`mlp.experts.E.{gate,up,down}_proj.{weight_packed,...}`) vs fused. +- **Vision tower stays bf16.** The IGNORE list excludes `vision_tower` + + `multi_modal_projector` (and all MLA attention, the MoE gate, embeddings, + lm_head) — only the expert FFN is NVFP4. So the vision encoder is byte-for-byte + full precision; any vision-quality nuance is the quantized LLM backbone, not the + tower. + +## Serve-test results (2026-06-17, heretic native NVFP4) + +- Loads on the vLLM native loader (v0.22.0), GPU0, ~91.9 GB at util 0.93. +- Text: correct (`2+2 = 4`, `capital of Japan = Tokyo`). +- Vision: tower functional — colors + spatial position accurate; exact shape + geometry fuzzy on small synthetic images (circle → pentagon). Evaluate real + vision quality during character tuning, against the official NVFP4 as baseline. + +## Provenance + +`convert_hf_to_native.py` originated with **worldtree-codex** (HF↔native bin +maps + the fused-expert split for the bf16 path). Fixed here: the NVFP4 layer +regexes (the `model.` prefix), and non-mmap shard reads. `nvfp4_quant.py` is +local. Homed in this infra repo per operator direction (not Worldtree). diff --git a/tools/mistral-small4-nvfp4/convert_hf_to_native.py b/tools/mistral-small4-nvfp4/convert_hf_to_native.py new file mode 100644 index 0000000..0c42bf2 --- /dev/null +++ b/tools/mistral-small4-nvfp4/convert_hf_to_native.py @@ -0,0 +1,525 @@ +#!/usr/bin/env python3 +"""Convert Mistral Small 4 HF weights to Mistral native consolidated format. + +The converter targets the BF16 path: + + HF Mistral3ForConditionalGeneration-style checkpoint + -> Mistral native `params.json` + `consolidated*.safetensors` + +It copies native runtime assets from an official native reference directory, +strips NVFP4 quantization config from `params.json`, rewrites tensor names, and +splits fused HF MoE expert tensors into native per-expert weights. +""" +from __future__ import annotations + +import argparse +import json +import math +import re +import shutil +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + + +TEXT_LAYER_BF16_RE = re.compile( + r"^model\.language_model\.model\.layers\.(?P\d+)\.(?P.+)$" +) +TEXT_LAYER_NVFP4_RE = re.compile( + r"^model\.language_model\.model\.layers\.(?P\d+)\.(?P.+)$" +) +VISION_LAYER_BF16_RE = re.compile( + r"^model\.vision_tower\.transformer\.layers\.(?P\d+)\.(?P.+)$" +) +VISION_LAYER_NVFP4_RE = re.compile( + r"^model\.vision_tower\.transformer\.layers\.(?P\d+)\.(?P.+)$" +) +NVFP4_EXPERT_RE = re.compile( + r"^mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight_packed|weight_scale|weight_global_scale|input_global_scale)$" +) + +TOP_LEVEL_MAP = { + "model.language_model.model.embed_tokens.weight": "tok_embeddings.weight", + "language_model.model.embed_tokens.weight": "tok_embeddings.weight", + "model.language_model.model.norm.weight": "norm.weight", + "language_model.model.norm.weight": "norm.weight", + "language_model.lm_head.weight": "output.weight", + "model.vision_tower.ln_pre.weight": "vision_encoder.ln_pre.weight", + "vision_tower.ln_pre.weight": "vision_encoder.ln_pre.weight", + "model.vision_tower.patch_conv.weight": "vision_encoder.patch_conv.weight", + "vision_tower.patch_conv.weight": "vision_encoder.patch_conv.weight", + "model.multi_modal_projector.linear_1.weight": "vision_language_adapter.w_in.weight", + "multi_modal_projector.linear_1.weight": "vision_language_adapter.w_in.weight", + "model.multi_modal_projector.linear_2.weight": "vision_language_adapter.w_out.weight", + "multi_modal_projector.linear_2.weight": "vision_language_adapter.w_out.weight", + "model.multi_modal_projector.norm.weight": "pre_mm_projector_norm.weight", + "multi_modal_projector.norm.weight": "pre_mm_projector_norm.weight", + "model.multi_modal_projector.patch_merger.merging_layer.weight": ( + "patch_merger.merging_layer.weight" + ), + "multi_modal_projector.patch_merger.merging_layer.weight": ( + "patch_merger.merging_layer.weight" + ), +} + +TEXT_LAYER_MAP = { + "input_layernorm.weight": "attention_norm.weight", + "post_attention_layernorm.weight": "ffn_norm.weight", + "self_attn.q_a_proj.weight": "attention.wq_a.weight", + "self_attn.q_b_proj.weight": "attention.wq_b.weight", + "self_attn.q_a_layernorm.weight": "attention.q_a_norm.weight", + "self_attn.kv_a_proj_with_mqa.weight": "attention.wkv_a_with_mqa.weight", + "self_attn.kv_a_layernorm.weight": "attention.kv_a_norm.weight", + "self_attn.kv_b_proj.weight": "attention.wkv_b.weight", + "self_attn.o_proj.weight": "attention.wo.weight", + "mlp.gate.weight": "gate.weight", + "mlp.shared_experts.gate_proj.weight": "shared_experts.w1.weight", + "mlp.shared_experts.up_proj.weight": "shared_experts.w3.weight", + "mlp.shared_experts.down_proj.weight": "shared_experts.w2.weight", +} + +VISION_LAYER_MAP = { + "attention.q_proj.weight": "attention.wq.weight", + "attention.k_proj.weight": "attention.wk.weight", + "attention.v_proj.weight": "attention.wv.weight", + "attention.o_proj.weight": "attention.wo.weight", + "attention_norm.weight": "attention_norm.weight", + "ffn_norm.weight": "ffn_norm.weight", + "feed_forward.gate_proj.weight": "feed_forward.w1.weight", + "feed_forward.up_proj.weight": "feed_forward.w3.weight", + "feed_forward.down_proj.weight": "feed_forward.w2.weight", +} + +NVFP4_EXPERT_PROJ_MAP = { + "gate_proj": "w1", + "down_proj": "w2", + "up_proj": "w3", +} + +ASSET_FILES = ( + "params.json", + "tekken.json", + "tokenizer_config.json", + "processor_config.json", + "chat_template.jinja", +) + + +@dataclass(frozen=True) +class TensorRef: + name: str + shard: str + + +@dataclass +class ConvertStats: + copied: int = 0 + split: int = 0 + bytes_written: int = 0 + shards_written: int = 0 + + +class ShardWriter: + def __init__(self, out_dir: Path, max_shard_bytes: int) -> None: + self.out_dir = out_dir + self.max_shard_bytes = max_shard_bytes + self.pending: dict[str, torch.Tensor] = {} + self.pending_bytes = 0 + self.weight_map: dict[str, str] = {} + self.shard_paths: list[Path] = [] + self.stats = ConvertStats() + + def add(self, name: str, tensor: torch.Tensor) -> None: + owned = tensor.detach().cpu().contiguous() + size = tensor_nbytes(owned) + if self.pending and self.pending_bytes + size > self.max_shard_bytes: + self.flush() + self.pending[name] = owned + self.pending_bytes += size + + def flush(self) -> None: + if not self.pending: + return + shard_idx = len(self.shard_paths) + 1 + shard_name = f"consolidated-{shard_idx:05d}.safetensors" + shard_path = self.out_dir / shard_name + save_file(self.pending, shard_path) + for name in self.pending: + self.weight_map[name] = shard_name + self.stats.bytes_written += self.pending_bytes + self.stats.shards_written += 1 + self.shard_paths.append(shard_path) + self.pending = {} + self.pending_bytes = 0 + + def write_index(self) -> None: + self.flush() + total_size = sum(path.stat().st_size for path in self.shard_paths) + index = { + "metadata": {"total_size": total_size}, + "weight_map": dict(sorted(self.weight_map.items())), + } + (self.out_dir / "consolidated.safetensors.index.json").write_text( + json.dumps(index, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def tensor_nbytes(tensor: torch.Tensor) -> int: + return tensor.numel() * tensor.element_size() + + +def parse_size_gib(value: str) -> int: + size = float(value) + if not math.isfinite(size) or size <= 0: + raise argparse.ArgumentTypeError("--max-shard-size-gb must be positive") + return int(size * 1024**3) + + +def load_weight_map(hf_dir: Path) -> list[TensorRef]: + index_path = hf_dir / "model.safetensors.index.json" + data = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = data.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"{index_path} does not contain a weight_map object") + return [TensorRef(name=name, shard=shard) for name, shard in sorted(weight_map.items())] + + +_SHARD_CACHE: dict = {} + + +def read_tensor(hf_dir: Path, ref: TensorRef) -> torch.Tensor: + # Non-mmap shard read: safetensors safe_open() mmaps the whole shard, which + # ENOMEMs on /tank (ZFS) for large shards (the 50 GB NVFP4 shard) regardless of + # free RAM or overcommit (a MAP_SHARED file mmap never consults the commit limit). + # Read the shard with a plain read() and deserialize from the in-memory buffer + # instead. Caches one shard at a time, so the copy loop must iterate refs grouped + # by shard (see the sorted(...) in convert()). + global _SHARD_CACHE + if _SHARD_CACHE.get("shard") != ref.shard: + from safetensors.torch import load as _st_load + with open(hf_dir / ref.shard, "rb") as _fh: + _SHARD_CACHE = {"shard": ref.shard, "tensors": _st_load(_fh.read())} + return _SHARD_CACHE["tensors"][ref.name] + + +def write_assets(native_ref_dir: Path, out_dir: Path, *, keep_quantization_config: bool) -> None: + for filename in ASSET_FILES: + src = native_ref_dir / filename + if not src.exists(): + continue + if filename == "params.json": + params = json.loads(src.read_text(encoding="utf-8")) + if not keep_quantization_config: + params.pop("quantization_config", None) + (out_dir / filename).write_text( + json.dumps(params, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + else: + shutil.copy2(src, out_dir / filename) + + +def mapped_name(name: str, *, output_format: str) -> str | None: + if name in TOP_LEVEL_MAP: + return TOP_LEVEL_MAP[name] + + text_match = ( + TEXT_LAYER_NVFP4_RE.match(name) + if output_format == "nvfp4" + else TEXT_LAYER_BF16_RE.match(name) + ) + if text_match: + layer = text_match.group("layer") + rest = text_match.group("rest") + suffix = TEXT_LAYER_MAP.get(rest) + if suffix is not None: + return f"layers.{layer}.{suffix}" + + if output_format == "nvfp4": + expert_match = NVFP4_EXPERT_RE.match(rest) + if expert_match: + expert = expert_match.group("expert") + native_proj = NVFP4_EXPERT_PROJ_MAP[expert_match.group("proj")] + artifact = expert_match.group("artifact") + return f"layers.{layer}.experts.{expert}.{native_proj}.{artifact}" + + for hf_proj, native_proj in NVFP4_EXPERT_PROJ_MAP.items(): + prefix = f"mlp.shared_experts.{hf_proj}." + if rest.startswith(prefix): + artifact = rest.removeprefix(prefix) + if artifact in { + "weight_packed", + "weight_scale", + "weight_global_scale", + "input_global_scale", + }: + return f"layers.{layer}.shared_experts.{native_proj}.{artifact}" + + return None + + vision_match = ( + VISION_LAYER_NVFP4_RE.match(name) + if output_format == "nvfp4" + else VISION_LAYER_BF16_RE.match(name) + ) + if vision_match: + layer = vision_match.group("layer") + rest = vision_match.group("rest") + suffix = VISION_LAYER_MAP.get(rest) + if suffix is not None: + return f"vision_encoder.transformer.layers.{layer}.{suffix}" + return None + + return None + + +def split_experts( + *, + writer: ShardWriter, + layer: int, + gate_up: torch.Tensor, + down: torch.Tensor, + expert_hidden_dim: int, +) -> None: + if gate_up.ndim != 3: + raise ValueError(f"layer {layer}: gate_up_proj must be rank 3, got {gate_up.shape}") + if down.ndim != 3: + raise ValueError(f"layer {layer}: down_proj must be rank 3, got {down.shape}") + if gate_up.shape[1] != expert_hidden_dim * 2: + raise ValueError( + f"layer {layer}: gate_up second dim {gate_up.shape[1]} != " + f"2 * expert_hidden_dim {expert_hidden_dim}" + ) + if gate_up.shape[0] != down.shape[0]: + raise ValueError( + f"layer {layer}: gate_up experts {gate_up.shape[0]} != down experts {down.shape[0]}" + ) + + for expert in range(gate_up.shape[0]): + writer.add( + f"layers.{layer}.experts.{expert}.w1.weight", + gate_up[expert, :expert_hidden_dim, :].clone(), + ) + writer.add( + f"layers.{layer}.experts.{expert}.w3.weight", + gate_up[expert, expert_hidden_dim:, :].clone(), + ) + writer.add(f"layers.{layer}.experts.{expert}.w2.weight", down[expert].clone()) + + +def text_layer_id(name: str, suffix: str) -> int | None: + match = TEXT_LAYER_BF16_RE.match(name) + if match and match.group("rest") == suffix: + return int(match.group("layer")) + return None + + +def convert( + *, + hf_dir: Path, + native_ref_dir: Path, + out_dir: Path, + max_shard_bytes: int, + expert_hidden_dim: int, + output_format: str, + dry_run: bool, +) -> dict[str, Any]: + refs = load_weight_map(hf_dir) + if output_format == "bf16": + gate_up_by_layer = { + layer: ref + for ref in refs + if (layer := text_layer_id(ref.name, "mlp.experts.gate_up_proj")) is not None + } + down_by_layer = { + layer: ref + for ref in refs + if (layer := text_layer_id(ref.name, "mlp.experts.down_proj")) is not None + } + else: + gate_up_by_layer = {} + down_by_layer = {} + + unmapped: list[str] = [] + mapped: dict[str, str] = {} + split_layers = sorted(set(gate_up_by_layer) | set(down_by_layer)) + + for ref in refs: + if output_format == "bf16" and ( + text_layer_id(ref.name, "mlp.experts.gate_up_proj") is not None + or text_layer_id(ref.name, "mlp.experts.down_proj") is not None + ): + continue + native_name = mapped_name(ref.name, output_format=output_format) + if native_name is None: + unmapped.append(ref.name) + else: + mapped[ref.name] = native_name + + if dry_run: + return { + "mapped_tensors": len(mapped), + "split_layers": split_layers, + "unmapped": unmapped, + "native_reference_diff": reference_name_diff( + native_ref_dir=native_ref_dir, + candidate_names=predicted_native_names(mapped.values(), split_layers), + output_format=output_format, + ), + } + + out_dir.mkdir(parents=True, exist_ok=True) + if any(out_dir.iterdir()): + raise FileExistsError(f"output directory is not empty: {out_dir}") + + write_assets( + native_ref_dir, + out_dir, + keep_quantization_config=output_format == "nvfp4", + ) + writer = ShardWriter(out_dir=out_dir, max_shard_bytes=max_shard_bytes) + + for ref in sorted(refs, key=lambda r: (r.shard, r.name)): + if ref.name not in mapped: + continue + writer.add(mapped[ref.name], read_tensor(hf_dir, ref)) + writer.stats.copied += 1 + + for layer in split_layers: + gate_ref = gate_up_by_layer.get(layer) + down_ref = down_by_layer.get(layer) + if gate_ref is None or down_ref is None: + raise ValueError(f"layer {layer}: missing gate_up or down fused expert tensor") + split_experts( + writer=writer, + layer=layer, + gate_up=read_tensor(hf_dir, gate_ref), + down=read_tensor(hf_dir, down_ref), + expert_hidden_dim=expert_hidden_dim, + ) + writer.stats.split += 1 + + writer.write_index() + + report = { + "copied_tensors": writer.stats.copied, + "split_layers": writer.stats.split, + "shards_written": writer.stats.shards_written, + "bytes_written": writer.stats.bytes_written, + "unmapped": unmapped, + "native_reference_diff": reference_name_diff( + native_ref_dir=native_ref_dir, + candidate_names=writer.weight_map.keys(), + output_format=output_format, + ), + } + (out_dir / "conversion_report.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return report + + +def predicted_native_names(mapped_names: Iterable[str], split_layers: Iterable[int]) -> set[str]: + names = set(mapped_names) + for layer in split_layers: + for expert in range(128): + for weight in ("w1", "w2", "w3"): + names.add(f"layers.{layer}.experts.{expert}.{weight}.weight") + return names + + +def iter_native_reference_names(native_ref_dir: Path, *, output_format: str) -> Iterable[str]: + index_path = native_ref_dir / "consolidated.safetensors.index.json" + if not index_path.exists(): + return () + data = json.loads(index_path.read_text(encoding="utf-8")) + names = data.get("weight_map", {}).keys() + if output_format == "nvfp4": + return set(names) + + normalized = set() + for name in names: + if name.endswith(".weight_packed"): + normalized.add(name.removesuffix("weight_packed") + "weight") + elif not ( + name.endswith(".weight_scale") + or name.endswith(".weight_global_scale") + or name.endswith(".input_global_scale") + ): + normalized.add(name) + return normalized + + +def reference_name_diff( + *, native_ref_dir: Path, candidate_names: Iterable[str], output_format: str +) -> dict[str, list[str]]: + reference_names = set( + iter_native_reference_names(native_ref_dir, output_format=output_format) + ) + if not reference_names: + return {"missing_from_output": [], "extra_in_output": []} + + candidates = set(candidate_names) + return { + "missing_from_output": sorted(reference_names - candidates), + "extra_in_output": sorted(candidates - reference_names), + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--native-ref-dir", type=Path, required=True) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument( + "--max-shard-size-gb", + type=parse_size_gib, + default=parse_size_gib("20"), + help="Approximate max safetensors shard size in GiB before flushing.", + ) + parser.add_argument( + "--expert-hidden-dim", + type=int, + default=2048, + help="Mistral Small 4 routed expert hidden dim used to split gate_up_proj.", + ) + parser.add_argument( + "--format", + choices=("bf16", "nvfp4"), + default="bf16", + help="Output checkpoint format. BF16 splits fused HF experts; NVFP4 renames per-expert quant artifacts.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Only report mapped/unmapped tensors; do not read or write weight shards.", + ) + return parser + + +def main() -> int: + args = build_parser().parse_args() + report = convert( + hf_dir=args.hf_dir, + native_ref_dir=args.native_ref_dir, + out_dir=args.out_dir, + max_shard_bytes=args.max_shard_size_gb, + expert_hidden_dim=args.expert_hidden_dim, + output_format=args.format, + dry_run=args.dry_run, + ) + print(json.dumps(report, indent=2, sort_keys=True)) + return 1 if report.get("unmapped") else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/mistral-small4-nvfp4/nvfp4_quant.py b/tools/mistral-small4-nvfp4/nvfp4_quant.py new file mode 100644 index 0000000..2f93b2f --- /dev/null +++ b/tools/mistral-small4-nvfp4/nvfp4_quant.py @@ -0,0 +1,26 @@ +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)