esh-pfi-infrastructure/services/heretic2-nvfp4-quant/quant_modelopt.py

#!/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, <function FusedMoE>)` 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 <out>/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 <out> <graft>  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())