NVIDIA Model Optimizer branch: nvfp4_experts_only PTQ for DeepSeek V4 Pro

quantize_modelopt.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""NVIDIA Model Optimizer PTQ for DeepSeek V4 Pro → NVFP4.
+
+Uses nvidia-modelopt's official PTQ pipeline with NVFP4Experts-Only config,
+which quantizes only MoE expert layers while keeping attention QKV in higher
+precision — the recommended approach for DeepSeek MoE models.
+
+Output is a Unified HuggingFace checkpoint deployable on TRT-LLM / vLLM / SGLang.
+
+Usage:
+    python quantize_modelopt.py \
+        --model /root/nvidia-meeting/DeepSeek-V4-Pro \
+        --export_dir /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4-modelopt \
+        --qformat nvfp4_experts_only \
+        --tp 8 \
+        --calib_size 256
+
+For the FP8 source variant, just change --model path. modelopt handles
+dequantization internally.
+"""
+
+import argparse
+import os
+import random
+import time
+
+import numpy as np
+import torch
+
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export import export_hf_checkpoint
+from modelopt.torch.utils.dataset_utils import create_forward_loop
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+mto.enable_huggingface_checkpointing()
+
+
+QUANT_CONFIGS = {
+    "nvfp4": mtq.NVFP4_DEFAULT_CFG,
+    "nvfp4_experts_only": mtq.NVFP4_EXPERTS_ONLY_CFG,
+    "nvfp4_mlp_only": mtq.NVFP4_MLP_ONLY_CFG,
+    "nvfp4_omlp_only": mtq.NVFP4_OMLP_ONLY_CFG,
+    "fp8": mtq.FP8_DEFAULT_CFG,
+}
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Model Optimizer PTQ for DeepSeek V4 Pro")
+    ap.add_argument("--model", required=True, help="Path to HF model (BF16 or FP8)")
+    ap.add_argument("--export_dir", required=True, help="Output directory for quantized checkpoint")
+    ap.add_argument("--qformat", default="nvfp4_experts_only",
+                    choices=list(QUANT_CONFIGS.keys()),
+                    help="Quantization format (default: nvfp4_experts_only for MoE)")
+    ap.add_argument("--kv_cache_qformat", default="fp8_cast",
+                    help="KV cache quantization (default: fp8_cast, fast no-calib)")
+    ap.add_argument("--tp", type=int, default=8, help="Tensor parallelism for export")
+    ap.add_argument("--calib_size", type=int, nargs="+", default=[256],
+                    help="Calibration dataset size (per dataset)")
+    ap.add_argument("--batch_size", type=int, default=1, help="Calibration batch size")
+    ap.add_argument("--calib_seq", type=int, default=4096, help="Max calibration sequence length")
+    ap.add_argument("--trust_remote_code", action="store_true", default=True,
+                    help="Trust remote code (required for V4)")
+    ap.add_argument("--use_seq_device_map", action="store_true",
+                    help="Use sequential device map for low-memory calibration")
+    ap.add_argument("--low_memory_mode", action="store_true",
+                    help="Compress weights before calibration (FP8/NVFP4 only)")
+    args = ap.parse_args()
+
+    print(f"=== Model Optimizer PTQ ===")
+    print(f"  Model:    {args.model}")
+    print(f"  QFormat:  {args.qformat}")
+    print(f"  KV Cache: {args.kv_cache_qformat}")
+    print(f"  TP:       {args.tp}")
+    print(f"  Calib:    {args.calib_size} samples, seq_len={args.calib_seq}")
+    print()
+
+    # Seed everything
+    random.seed(1234)
+    np.random.seed(1234)
+    torch.manual_seed(1234)
+
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+        padding_side="left",
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load model
+    print("Loading model...")
+    model_kwargs = {
+        "trust_remote_code": args.trust_remote_code,
+        "torch_dtype": torch.bfloat16,
+    }
+    if args.use_seq_device_map:
+        model_kwargs["device_map"] = "auto"
+
+    model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
+
+    if not args.use_seq_device_map:
+        model = model.cuda()
+
+    # Build calibration dataloader
+    print("Building calibration dataset...")
+    calib_dataloader = create_forward_loop(
+        model,
+        dataloader=get_dataloader(
+            tokenizer=tokenizer,
+            calib_size=args.calib_size,
+            batch_size=args.batch_size,
+            calib_seq=args.calib_seq,
+        ),
+    )
+
+    # Quantize
+    quant_cfg = QUANT_CONFIGS[args.qformat]
+    print(f"Running PTQ with {args.qformat}...")
+    t0 = time.time()
+
+    model = mtq.quantize(model, quant_cfg, calib_dataloader)
+
+    elapsed = time.time() - t0
+    print(f"Quantization complete in {elapsed/60:.1f} min")
+
+    # Export
+    print(f"Exporting to {args.export_dir} ...")
+    with torch.inference_mode():
+        export_hf_checkpoint(
+            model,
+            args.export_dir,
+            tokenizer=tokenizer,
+            export_tensorrt_llm_plugins=True,
+        )
+
+    print(f"Done. Output at {args.export_dir}")
+
+
+def get_dataloader(tokenizer, calib_size, batch_size, calib_seq):
+    """Create calibration dataloader using modelopt's built-in dataset utils."""
+    from modelopt.torch.utils.dataset_utils import get_dataset_dataloader
+
+    return get_dataset_dataloader(
+        tokenizer=tokenizer,
+        num_samples=calib_size[0],
+        batch_size=batch_size,
+        seq_len=calib_seq,
+    )
+
+
+if __name__ == "__main__":
+    main()