#!/usr/bin/env python3 """ NVFP4 quantization recipe for MiniMax-M2 family with the GB10 ignore list, using a 6-dataset agentic calibration mix. This is the script that produced `saricles/MiniMax-M2.7-REAP-172B-A10B-NVFP4-GB10` (REAP-pruned 172B variant). The same script can be run against the full 230B BF16 source by changing INPUT_DIR — identical recipe, different input. WHAT THIS DOES -------------- 1. Loads a BF16 model (e.g., saricles/MiniMax-M2.7-REAP-172B-A10B-BF16 or the full dequantized MiniMax-M2.7). 2. Runs `mtq.quantize` with NVFP4_DEFAULT_CFG on every Linear in the model, calibrating against a 6-dataset agentic mix (evol-codealpaca, xlam-function- calling, Mixture-of-Thoughts code/math/science, SWE-smith-trajectories `tool` split). 3. Disables quantizers on the GB10 ignore list (lm_head + MoE router gate, plus optionally embed_tokens and first/last layer) POST-quantize via `mtq.disable_quantizer`. Self-attention STAYS quantized — that's the GB10 specialization vs. the standard NVFP4 reference variant that keeps attention in BF16. 4. Force-populates `weight_quantizer.amax` on every enabled quantizer that calibration missed (see "MoE expert calibration gotcha" below). 5. Exports via `export_hf_checkpoint` to a HuggingFace-loadable NVFP4 checkpoint in compressed-tensors format. THE MoE EXPERT CALIBRATION GOTCHA --------------------------------- The REAP'd input has 192 experts and uses top-K=8 routing (the full M2.7 has 256 experts with the same top-K=8). With 64 samples × 2048 tokens × 8 experts- per-token, that's roughly 1M expert activations per layer; even across 192-256 experts, most see plenty of calibration. BUT MoE routing is heavily skewed — popular experts dominate while tail experts may be undersampled or never fire during short calibration runs. Experts that never activate have their `weight_quantizer.amax` unset → `export_hf_checkpoint` asserts on them. Newer modelopt (>= a recent version) auto-handles this via `_calibrate_weight_quantizer_if_needed` (see modelopt source `modelopt/torch/quantization/quant_utils.py:275-317`). Older versions don't. This script's Phase 2.5 manually replicates that fix: - Find every enabled `weight_quantizer` with `amax is None` - Wrap a forward pass in `enable_stats_collection` / `finish_stats_collection` (without the wrap, after `mtq.quantize()` flips `_if_calib=False` on every quantizer, a bare `wq(weight)` call takes the QUANT branch — reads amax — instead of the CALIB branch — writes amax. Silently no-ops.) - This populates amax from weight statistics so export sees a valid quantizer. Math is the same as activation-derived amax (just slightly less precise for those experts since they were never seen during real routing — but they weren't seen, so any reasonable scale is fine). ENVIRONMENT ----------- INPUT_DIR — path to BF16 source model (default: /model) OUTPUT_DIR — path to write NVFP4 output (default: /tmp/nvfp4_model) OFFLOAD_DIR — accelerate offload folder for big models (default: /tmp/offload) NUM_CALIB — number of calibration samples (default: 64) MAX_SEQ — calibration max sequence length (default: 2048) CALIB_DATASET — HF dataset for calibration (default: HuggingFaceH4/ultrachat_200k) IGNORE_LM_HEAD — keep lm_head in BF16 (default: 1) IGNORE_ROUTER_GATE — keep MoE router gate in BF16 (default: 1) USAGE ----- Single-host (model fits in one GPU + offload): INPUT_DIR=/path/to/MiniMax-M2.7-BF16 \\ OUTPUT_DIR=./minimax-m2.7-nvfp4-gb10 \\ python quantize-nvfp4-gb10.py Multi-GPU (HF Jobs a100x8 was our reference target): Same env vars; transformers' device_map="auto" handles sharding. ADAPTING FOR ANOTHER ARCHITECTURE --------------------------------- Most NVFP4 quantization works the same. Two things to change: 1. The ignore-list patterns (line ~140) for your model's router/lm_head names. M2 uses `*block_sparse_moe.gate`. Mixtral uses `*block_sparse_moe.gate`. Qwen3.5 uses `*mlp.gate`. Adjust accordingly. 2. The Conv1D shim (line ~125) is only needed for transformers 4.57+ models that have removed Conv1D from modeling_utils. Harmless to leave for others. The amax-populate phase (Phase 2.5) is architecture-agnostic and worth keeping for any MoE model. REQUIREMENTS ------------ pip install nvidia-modelopt 'transformers>=4.57' \\ 'huggingface_hub>=0.30' 'tokenizers>=0.22' \\ accelerate datasets License: Other (inherited from MiniMax-M2.7 — see base model card). """ import os import shutil import time from pathlib import Path import torch start = time.time() # --------------------------------------------------------------------------- # Configuration (env-driven) # --------------------------------------------------------------------------- INPUT_DIR = Path(os.environ.get("INPUT_DIR", "/model")) OUTPUT_DIR = Path(os.environ.get("OUTPUT_DIR", "/tmp/nvfp4_model")) OFFLOAD_DIR = Path(os.environ.get("OFFLOAD_DIR", "/tmp/offload")) NUM_CALIB_PER_DS = int(os.environ.get("NUM_CALIB_PER_DS", "64")) MAX_SEQ = int(os.environ.get("MAX_SEQ", "2048")) # Agentic calibration mix — matches the 6 datasets used in REAP preservation. # Each tuple: (dataset_name, subset_or_None, split). # Schemas vary; extractor below handles messages / trajectory / instruction / query. CALIB_DATASETS = [ ("theblackcat102/evol-codealpaca-v1", None, "train"), ("Salesforce/xlam-function-calling-60k", None, "train"), ("open-r1/Mixture-of-Thoughts", "code", "train"), ("open-r1/Mixture-of-Thoughts", "math", "train"), ("open-r1/Mixture-of-Thoughts", "science", "train"), ("SWE-bench/SWE-smith-trajectories", None, "tool"), ] CRITICAL_DATASETS = {"Salesforce/xlam-function-calling-60k", "SWE-bench/SWE-smith-trajectories"} IGNORE_LM_HEAD = bool(int(os.environ.get("IGNORE_LM_HEAD", "1"))) IGNORE_ROUTER_GATE = bool(int(os.environ.get("IGNORE_ROUTER_GATE", "1"))) def log(msg): print(f"[{time.time()-start:7.1f}s] {msg}", flush=True) # --------------------------------------------------------------------------- # Phase 0 — Preflight # --------------------------------------------------------------------------- log("=== PHASE 0: Preflight ===") assert INPUT_DIR.exists(), f"INPUT_DIR not found: {INPUT_DIR}" assert (INPUT_DIR / "config.json").exists(), f"config.json missing in {INPUT_DIR}" # CUDA init retry — H200 / cloud pools sometimes have brief Error 802 races cuda_ok = False for attempt in range(10): if torch.cuda.is_available(): cuda_ok = True break log(f"CUDA not ready (attempt {attempt+1}/10) — sleeping 3s") time.sleep(3) assert cuda_ok, "CUDA still not available after retries" log(f"Preflight OK — {torch.cuda.device_count()} GPUs, CUDA {torch.version.cuda}") log(f"INPUT: {INPUT_DIR}") log(f"OUTPUT: {OUTPUT_DIR}") OFFLOAD_DIR.mkdir(parents=True, exist_ok=True) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) # --------------------------------------------------------------------------- # Phase 1 — Load BF16 model # --------------------------------------------------------------------------- log("=== PHASE 1: Load BF16 ===") # transformers 4.57+ moved Conv1D out of modeling_utils; some custom modeling # code still imports from the old path. Harmless shim. import transformers.modeling_utils if not hasattr(transformers.modeling_utils, "Conv1D"): from transformers.pytorch_utils import Conv1D as _Conv1D transformers.modeling_utils.Conv1D = _Conv1D from transformers import AutoModelForCausalLM, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(str(INPUT_DIR), trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( str(INPUT_DIR), dtype=torch.bfloat16, device_map="auto", trust_remote_code=True, offload_folder=str(OFFLOAD_DIR), ) devices = sorted({str(p.device) for p in model.parameters()}) log(f"Model loaded across: {devices}") # --------------------------------------------------------------------------- # Phase 2 — modelopt NVFP4 with GB10 ignore list # --------------------------------------------------------------------------- log("=== PHASE 2: modelopt NVFP4 quantization ===") import modelopt.torch.quantization as mtq from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, ) from modelopt.torch.export import export_hf_checkpoint from datasets import load_dataset nvfp4_cfg = mtq.NVFP4_DEFAULT_CFG log("Config: NVFP4_DEFAULT_CFG straight (post-quantize disable for ignore list)") log(f"Loading agentic calibration mix: {NUM_CALIB_PER_DS} samples × {len(CALIB_DATASETS)} datasets") all_texts = [] datasets_loaded = 0 failed_datasets = set() for ds_name, subset, split in CALIB_DATASETS: label = f"{ds_name}" + (f"[{subset}]" if subset else "") try: load_kwargs = {"split": f"{split}[:{NUM_CALIB_PER_DS}]"} if subset is not None: load_kwargs["name"] = subset ds = load_dataset(ds_name, **load_kwargs) import json as _json loaded_from_ds = 0 for sample in ds: text = None # messages — may be a list-of-dicts OR a JSON-encoded string # (SWE-smith-trajectories stores messages as JSON string!) if "messages" in sample: msgs = sample["messages"] if isinstance(msgs, str): try: msgs = _json.loads(msgs) except Exception: msgs = None if isinstance(msgs, list): text = "\n".join( m.get("content", "") for m in msgs if isinstance(m, dict) and m.get("content") ) elif "trajectory" in sample: traj = sample["trajectory"] if isinstance(traj, str): try: traj = _json.loads(traj) except Exception: traj = None if isinstance(traj, list): text = "\n".join( m.get("content", "") for m in traj if isinstance(m, dict) and m.get("content") ) elif traj is not None: text = str(traj) elif "instruction" in sample: text = sample.get("instruction", "") + "\n" + sample.get("output", "") elif "query" in sample: text = sample.get("query", "") + "\n" + str(sample.get("answers", "")) if text and text.strip(): all_texts.append(text) loaded_from_ds += 1 # BUG-CATCH: silent 0-text extraction is how SWE-smith was missed in the # original run (see Calibration notes in model card). Hard-fail now. assert loaded_from_ds > 0, ( f"EXTRACTOR BUG: {label} returned 0 texts despite loading. " f"Schema drift — investigate sample fields." ) datasets_loaded += 1 log(f" {label}: {loaded_from_ds} texts collected") except AssertionError: raise # fail-fast, don't hide except Exception as e: log(f" {label}: FAILED — {e}") failed_datasets.add(ds_name) critical_failed = failed_datasets & CRITICAL_DATASETS assert not critical_failed, f"CRITICAL dataset(s) failed: {critical_failed} — aborting" assert datasets_loaded >= 5, f"Only {datasets_loaded}/6 datasets loaded — need >= 5" assert len(all_texts) >= 300, f"Only {len(all_texts)} texts — need >= 300" log(f"Total: {len(all_texts)} texts from {datasets_loaded}/6 datasets") def forward_loop(m): m.eval() with torch.no_grad(): for i, text in enumerate(all_texts): toks = tokenizer( text, return_tensors="pt", max_length=MAX_SEQ, truncation=True ) if toks.input_ids.shape[1] < 32: continue ids = toks.input_ids.to(next(m.parameters()).device) m(input_ids=ids) if (i + 1) % 25 == 0: log(f"Calibrated {i+1}/{len(all_texts)}") log("Starting mtq.quantize") model = mtq.quantize(model, nvfp4_cfg, forward_loop) log("Quantization complete — applying ignore list post-quantize") # Ignore list — adapt the patterns for your model's router/lm_head module names. # Defaults match MiniMax-M2 family. The GB10 specialization vs. standard NVFP4 # is that self_attn STAYS QUANTIZED here (we don't add `*self_attn*` to ignore). if IGNORE_LM_HEAD: mtq.disable_quantizer(model, "lm_head*") log("ignore: lm_head (kept BF16)") if IGNORE_ROUTER_GATE: mtq.disable_quantizer(model, "*block_sparse_moe.gate") log("ignore: *block_sparse_moe.gate (MoE router gate kept BF16)") # --------------------------------------------------------------------------- # Phase 2.5 — Force-populate amax on never-calibrated MoE experts # --------------------------------------------------------------------------- # See module docstring for the full explanation. tl;dr: with N samples × top-K # routing across E experts, when N×K << E, many experts never activate during # calibration and their weight_quantizer.amax stays None → export asserts. # We populate amax from weight statistics, mirroring newer modelopt's # `_calibrate_weight_quantizer_if_needed` helper. log("=== PHASE 2.5: Force-populate amax on never-calibrated quantizers ===") populated = 0 skipped_disabled = 0 already_ok = 0 for name, module in model.named_modules(): wq = getattr(module, "weight_quantizer", None) if wq is None: continue if not getattr(wq, "is_enabled", True): skipped_disabled += 1 continue if getattr(wq, "amax", None) is not None: already_ok += 1 continue with torch.no_grad(): if hasattr(wq, "reset_amax"): wq.reset_amax() # CRITICAL: enable_stats_collection routes the forward through the # calibration branch (writes amax). Without this wrap, after # mtq.quantize completes, _if_calib=False and the forward takes the # quant branch — reads amax — and silently no-ops. enable_stats_collection(wq) wq(module.weight) finish_stats_collection(wq) populated += 1 log(f"amax populate: {populated} filled, {already_ok} already set, {skipped_disabled} disabled-skipped") # --------------------------------------------------------------------------- # Phase 3 — Export to HuggingFace-compatible NVFP4 checkpoint # --------------------------------------------------------------------------- log("=== PHASE 3: Export ===") with torch.inference_mode(): export_hf_checkpoint(model, export_dir=str(OUTPUT_DIR)) tokenizer.save_pretrained(str(OUTPUT_DIR)) # Carry custom modeling code + chat template forward so trust_remote_code works # downstream (vLLM, transformers loaders). for f in os.listdir(INPUT_DIR): if f.startswith(("modeling_", "configuration_", "tokenization_")) and f.endswith(".py"): shutil.copy2(INPUT_DIR / f, OUTPUT_DIR / f) if f == "chat_template.jinja": shutil.copy2(INPUT_DIR / f, OUTPUT_DIR / f) out_gb = sum(p.stat().st_size for p in OUTPUT_DIR.rglob("*") if p.is_file()) / 1e9 log(f"Export size: {out_gb:.1f} GB") total = time.time() - start log(f"ALL DONE — {total:.0f}s ({total/60:.1f} min)") log(f"NVFP4 written to: {OUTPUT_DIR}")