#!/usr/bin/env python3 """NVFP4-quantize the MTP-grafted Heretic2 model via NVIDIA nvidia-modelopt (MODELOPT format). Sibling of quant_nvfp4.py (llm-compressor / compressed-tensors) but produces the **modelopt** NVFP4 format instead. Why it exists: the compressed-tensors path (quant_nvfp4.py) serves COHERENT but its MTP is 0% acceptance — vLLM's `Qwen3_5MTP` speculative-decode drafter only loads the bf16 MTP head off a **modelopt** main-model checkpoint (the mtp tensors are byte-identical between the two formats; the difference is purely how the main model's quantized weights/scales are stored). So working native MTP (the ~2-4x goal) requires this format. (2026-07-14 modelopt pivot.) Reference = AEON `/tank/aimodels/qwen36-27b-aeon-nvfp4` (served by vllm-aeon-rp). Its hf_quant_config.json says: quant_algo NVFP4 (W4A4, group_size 16), targets Linear, exclude `lm_head` + `model.visual*` + every `linear_attn*` (the GDN). This script matches that EXACTLY. `embed_tokens` is an Embedding (not a Linear target) so it is never quantized — no need to exclude. Pipeline (unchanged except THIS quant step swaps llm-compressor -> modelopt): graft_mtp.py -> quant_modelopt.py (this) -> splice_mtp.py (bf16 mtp) -> serve serve: vllm --quantization modelopt --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":3}' (AEON vllm-aeon-rp is the ref) The MTP head is NOT in the module tree at load (transformers builds no mtp module for either Qwen3_5 class), so the 15 bf16 mtp.* tensors are spliced back AFTER export — same as pantheon/AEON. Load class is AutoModelForImageTextToText (= Qwen3_5ForConditionalGeneration) so weight keys are born `model.language_model.*` + `model.visual.*` — the namespace vLLM's ConditionalGeneration loader expects. (The compressed-tensors-path gibberish was a text-only-namespace bug; same fix.) Run in a vLLM container on the freed GPU0 (nvidia-modelopt[hf] per the transformers-compat warning): docker run --gpus '"device=0"' --ipc host -v /tank/aimodels:/tank/aimodels -v /home/lkraven:/lk \ --entrypoint bash vllm/vllm-openai:v0.24.0 -c \ "pip install -q 'nvidia-modelopt[hf]' tiktoken sentencepiece && python3 /lk/quant_modelopt.py \ --model /tank/aimodels/heretic2-nvfp4-work/heretic2-mtp-bf16 \ --calib-mode chat --calib /tank/aimodels/heretic2-nvfp4-work/production_calib_512.jsonl \ --num-samples 512 --seqlen 8192 \ --out /tank/aimodels/heretic2-nvfp4-work/heretic2-modelopt-nvfp4" """ import argparse import copy import json import sys def load_calib_text(name, tokenizer, n, seqlen): from datasets import load_dataset ds = load_dataset(name, split="train").shuffle(seed=42).select(range(min(n, 100000))) col = "text" if "text" in ds.column_names else ds.column_names[0] return [tokenizer(x[col], truncation=True, max_length=seqlen) for x in ds.select(range(n))] def load_calib_chat(path, tokenizer, seqlen, n): # Identical to quant_nvfp4.py: render each row via the Qwen3.6 chat template with thinking on # so the forward-pass sees the seat's native tool-call XML activations. tool_call arguments may # arrive as OpenAI wire-form JSON strings; the template does .items() on them -> parse to dict. rows = [json.loads(l) for l in open(path) if l.strip()][:n] out = [] for r in rows: for m in r["messages"]: for tc in (m.get("tool_calls") or []): a = tc.get("function", {}).get("arguments") if isinstance(a, str): tc["function"]["arguments"] = json.loads(a) text = tokenizer.apply_chat_template( r["messages"], tools=r.get("tools"), tokenize=False, add_generation_prompt=False, enable_thinking=True, ) out.append(tokenizer(text, truncation=True, max_length=seqlen)) return out def build_nvfp4_cfg(mtq): """NVFP4 W4A4 (group_size 16) matching AEON's exclusions. NVFP4_DEFAULT_CFG already disables lm_head + linear_attn.conv1d/in_proj_a/in_proj_b + mlp.gate; append the FULL linear_attn (GDN) and the vision tower so only the standard attn/MLP Linears get NVFP4 (later entries override).""" cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG) cfg["quant_cfg"].append({"quantizer_name": "*linear_attn*", "enable": False}) cfg["quant_cfg"].append({"quantizer_name": "*visual*", "enable": False}) cfg["quant_cfg"].append({"quantizer_name": "*mtp*", "enable": False}) # moot (not in tree); belt+suspenders return cfg def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--model", required=True) ap.add_argument("--calib-mode", choices=["text", "chat"], required=True) ap.add_argument("--calib", required=True) ap.add_argument("--out", required=True) ap.add_argument("--num-samples", type=int, default=512) ap.add_argument("--seqlen", type=int, default=8192) args = ap.parse_args() import torch from transformers import AutoModelForImageTextToText, AutoTokenizer import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint # modelopt 0.45 + transformers 5.12.1 compat guard. transformers 5.x exposes `FusedMoE` as a # FUNCTION, but modelopt registers it in QuantModuleRegistry expecting an nn.Module class, so the # registry scan (register_fused_experts_on_the_fly -> _get_registered_nn_class) does # `issubclass(nn_cls, )` and dies with "arg 2 must be a class". Our model is # DENSE (no FusedMoE) so skipping non-class registry entries is safe. Guard the scan: from modelopt.torch.opt import dynamic as _mo_dyn def _grnc_safe(self, nn_cls): for nn_cls_ in self._registry: if (isinstance(nn_cls_, type) and issubclass(nn_cls, nn_cls_) and nn_cls.forward is nn_cls_.forward): return nn_cls_ return None _mo_dyn._DMRegistryCls._get_registered_nn_class = _grnc_safe print(f"loading grafted model (multimodal ConditionalGeneration): {args.model}", flush=True) model = AutoModelForImageTextToText.from_pretrained( args.model, torch_dtype="auto", device_map="auto", trust_remote_code=True, ) model.eval() tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) print(f"building calibration ({args.calib_mode}, up to {args.num_samples} @ seq {args.seqlen})", flush=True) if args.calib_mode == "text": calib = load_calib_text(args.calib, tok, args.num_samples, args.seqlen) else: calib = load_calib_chat(args.calib, tok, args.seqlen, args.num_samples) print(f" {len(calib)} calibration rows", flush=True) def forward_loop(m): with torch.no_grad(): for i, row in enumerate(calib): ids = torch.tensor([row["input_ids"]], device=m.device) m(input_ids=ids) if (i + 1) % 64 == 0: print(f" calib {i + 1}/{len(calib)}", flush=True) cfg = build_nvfp4_cfg(mtq) print("running modelopt NVFP4 PTQ (W4A4 g16; lm_head/linear_attn/visual kept BF16)", flush=True) mtq.quantize(model, cfg, forward_loop=forward_loop) print(f"exporting modelopt HF checkpoint -> {args.out}", flush=True) # Force a SINGLE shard (default max_shard_size 10GB would split the ~14GB output into 3 shards, # but splice_mtp.py expects a single /model.safetensors to add the bf16 mtp.* into). export_hf_checkpoint(model, export_dir=args.out, max_shard_size="1TB") tok.save_pretrained(args.out) print("DONE. Next: splice_mtp.py then serve " "--quantization modelopt --speculative-config " "'{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":3}' " "--reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice", flush=True) return 0 if __name__ == "__main__": sys.exit(main())