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).
This commit is contained in:
@@ -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 <hf-bf16-dir> 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 <official native NVFP4 snapshot 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 <ref> \
|
||||
--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).
|
||||
@@ -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<layer>\d+)\.(?P<rest>.+)$"
|
||||
)
|
||||
TEXT_LAYER_NVFP4_RE = re.compile(
|
||||
r"^model\.language_model\.model\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
|
||||
)
|
||||
VISION_LAYER_BF16_RE = re.compile(
|
||||
r"^model\.vision_tower\.transformer\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
|
||||
)
|
||||
VISION_LAYER_NVFP4_RE = re.compile(
|
||||
r"^model\.vision_tower\.transformer\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
|
||||
)
|
||||
NVFP4_EXPERT_RE = re.compile(
|
||||
r"^mlp\.experts\.(?P<expert>\d+)\."
|
||||
r"(?P<proj>gate_proj|up_proj|down_proj)\."
|
||||
r"(?P<artifact>weight_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())
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user