""" 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 """ 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()