vh
dd3a5c93fd
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).
526 lines
18 KiB
Python
526 lines
18 KiB
Python
#!/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())
|