quantization_recipe/03_ptq_modelopt.py

"""
NVFP4 PTQ of llm-jp-4-8b-instruct with NVIDIA Model Optimizer — the recipe
used to produce this checkpoint.

Mixed precision:
  - MLP linears (gate/up/down): NVFP4 (FP4 block-16, FP8 block scales)
  - Attention projections (q/k/v/o): FP8 (E4M3 per-tensor)
  - First 2 and last 2 decoder layers, embeddings, lm_head, norms: BF16

Calibration: calib_full.jsonl (bundled) — chat-templated conversations plus
raw-text strict-format samples. Conversation samples go through
tokenizer.apply_chat_template so activation ranges see deployment-realistic
Harmony token streams; {"text": ...} samples are encoded as-is (raw
zero-shot completion form).

Requires nvidia-modelopt==0.45.0. Run with the GPU otherwise idle (needs
~20GB+ headroom for the BF16 model plus calibration activations):
  python3 03_ptq_modelopt.py --calib calib_full.jsonl --export-dir <export-dir>
"""

import argparse
import copy
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import modelopt.torch.quantization as mtq
from modelopt.torch.export import export_hf_checkpoint

DEFAULT_MODEL = "llm-jp/llm-jp-4-8b-instruct"
DEFAULT_REVISION = "098f2b2cf33021eba19a6d3582aa3d071ccc0aff"  # the revision this checkpoint was built from

BF16_LAYERS = (0, 1, 30, 31)  # Nemotron-style: first/last 2 layers stay BF16


def build_calib_batches(calib_path: str, tokenizer, n_samples: int, max_seq_len: int):
    """Tokenize calibration samples.

    NOTE: apply_chat_template(tokenize=True)'s return type varies across
    transformers versions (list vs BatchEncoding — the latter silently broke
    len() checks on 5.5.x). Always template to a string, then encode.
    """
    batches = []
    with open(calib_path) as f:
        for i, line in enumerate(f):
            if i >= n_samples:
                break
            rec = json.loads(line)
            if "messages" in rec:
                text = tokenizer.apply_chat_template(
                    rec["messages"], add_generation_prompt=False, tokenize=False
                )
            else:
                # Raw-completion sample (strict-format slice): encode as-is —
                # these prompts are served with no chat template.
                text = rec["text"]
            ids = tokenizer.encode(text, add_special_tokens=False)[:max_seq_len]
            batches.append(torch.tensor([ids], dtype=torch.long))
    return batches


def make_forward_loop(batches):
    def forward_loop(model):
        with torch.no_grad():
            for i, input_ids in enumerate(batches):
                model(input_ids=input_ids.to(model.device))
                if (i + 1) % 64 == 0:
                    print(f"  calib {i + 1}/{len(batches)}", flush=True)

    return forward_loop


def build_quant_config():
    # modelopt 0.45's quant_cfg is an ORDERED LIST of {quantizer_name, ...}
    # entries where later entries override earlier pattern matches.
    # NVFP4_DEFAULT_CFG already disables lm_head, routers, MoE gates, and
    # norm/BN layers.
    cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG)

    # Attention projections to per-tensor FP8 (E4M3); MLP linears stay NVFP4
    # from the base config. Appended entries win over the base patterns.
    fp8_w = {"num_bits": (4, 3), "axis": None}
    for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
        cfg["quant_cfg"].append(
            {"quantizer_name": f"*self_attn.{proj}*weight_quantizer", "cfg": dict(fp8_w)}
        )
        cfg["quant_cfg"].append(
            {"quantizer_name": f"*self_attn.{proj}*input_quantizer", "cfg": dict(fp8_w)}
        )

    # Most quantization-sensitive layers stay BF16.
    for layer in BF16_LAYERS:
        cfg["quant_cfg"].append({"quantizer_name": f"*layers.{layer}.*", "enable": False})

    # Embeddings are not in the default ignore list — exclude explicitly.
    cfg["quant_cfg"].append({"quantizer_name": "*embed_tokens*", "enable": False})
    return cfg


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=DEFAULT_MODEL,
                        help="HF id or local path of the BF16 base model")
    parser.add_argument("--revision", default=DEFAULT_REVISION,
                        help="HF revision of the base model (pinned to the "
                             "revision this checkpoint was built from; "
                             "ignored for local paths)")
    parser.add_argument("--calib", default="calib_full.jsonl")
    parser.add_argument("--n-samples", type=int, default=588)
    parser.add_argument("--max-seq-len", type=int, default=4096)
    parser.add_argument("--export-dir", required=True)
    args = parser.parse_args()

    print(f"Loading tokenizer + model from {args.model}...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(
        args.model, revision=args.revision, trust_remote_code=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        args.model, revision=args.revision, torch_dtype=torch.bfloat16,
        device_map="cuda", trust_remote_code=True,
    )
    model.eval()

    calib_path = Path(args.calib)
    if not calib_path.exists():
        # default calib_full.jsonl ships next to this script — resolve
        # relative to the script dir so the command works from anywhere
        calib_path = Path(__file__).parent / args.calib
    print(f"Building calibration batches from {calib_path}...", flush=True)
    batches = build_calib_batches(str(calib_path), tokenizer, args.n_samples, args.max_seq_len)
    print(f"  {len(batches)} samples, total tokens "
          f"{sum(b.numel() for b in batches)}", flush=True)

    print("Quantizing...", flush=True)
    model = mtq.quantize(model, build_quant_config(), make_forward_loop(batches))

    # Mandatory gate: confirm what actually got quantized before any
    # export/eval investment. Saved alongside the checkpoint for review.
    export_dir = Path(args.export_dir)
    export_dir.mkdir(parents=True, exist_ok=True)
    import contextlib, io

    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        mtq.print_quant_summary(model)
    summary = buf.getvalue()
    (export_dir / "quant_summary.txt").write_text(summary)
    print(summary[:4000], flush=True)
    for banned in ("lm_head", "embed_tokens"):
        for line in summary.splitlines():
            if banned in line and "TensorQuantizer" in line and "disabled" not in line:
                raise SystemExit(f"GATE FAILED: {banned} appears quantized: {line}")
    print("Quant summary gate passed (lm_head/embed_tokens not quantized).", flush=True)

    print(f"Exporting HF checkpoint to {export_dir}...", flush=True)
    export_hf_checkpoint(model, export_dir=str(export_dir))
    print("Done. Now run 04_export_and_patch.py on the export dir.", flush=True)


if __name__ == "__main__":
    main()