Fundamental rewrite: call hf_main() instead of rewriting the pipeline

The previous approach tried to reconstruct hf_ptq's pipeline by importing
individual functions and building a fake argparse.Namespace. This caused
repeated crashes from missing args (KV_QUANT_CFG_CHOICES, dataset,
calib_with_images, etc.).

New approach:
- Call hf_ptq.parse_args() with sys.argv replaced — gets ALL defaults
- Call hf_main(args) — the exact same entry point the shell script uses
- Hook export_quantized to add amax snapshot + state save before export
- No more missing args. No more diverging from the example script.

The only changes from the stock pipeline:
1. Runtime patches (load_calib_amax CPU, export_amax CPU, clamp)
2. Post-calibration hook (snapshot amax, save state, force CPU)

scripts/quantize_nvfp4.py

@@ -2,18 +2,16 @@
 """
 DeepSeek V4 Pro → NVFP4 quantization — defensive edition.
 
-Runs the full ModelOpt PTQ pipeline with maximum protection against GPU tensor
-corruption that crashes the export after 6 hours of calibration.
+This script:
+1. Applies runtime patches for GPU tensor safety (before modelopt runs)
+2. Calls the SAME hf_ptq.py pipeline that the shell script uses
+3. After calibration, snapshots amax to CPU and saves model state
 
-Key defense: immediately after calibration, every quantizer _amax tensor is
-snapshotted to CPU. Then the model state is saved to disk. If export crashes,
-the state can be reloaded and export retried without re-calibrating.
+The key insight: we don't rewrite the pipeline. We let hf_ptq do its thing
+with all its args, defaults, and edge cases handled correctly. We just add
+our defensive patches and post-calibration saves.
 
-The _amax tensors are tiny (scalars and small vectors). Snapshotting ~49K of them
-to CPU costs almost nothing in memory and guarantees we have valid calibration
-data regardless of what CUDA does to the GPU copies afterward.
-
-Must be run from the modelopt example directory for imports:
+Must be run from the modelopt example directory:
     cd /root/nvidia-meeting/modelopt-repo/examples/llm_ptq
     python3 /root/nvidia-meeting/deepseek-v4-quant/scripts/quantize_nvfp4.py
 
@@ -29,12 +27,10 @@ Usage:
 """
 
 import argparse
-import copy
 import gc
 import os
 import sys
 import time
-import warnings
 
 import torch
 
@@ -58,17 +54,16 @@ AMAX_SNAPSHOT_PATH = "/root/nvidia-meeting/v4_nvfp4_amax_snapshots.pt"
 
 
 def apply_patches():
-    """Apply runtime patches for V4 compatibility and GPU tensor safety."""
+    """Apply runtime patches for V4 compatibility and GPU tensor safety.
+
+    These patches are applied BEFORE hf_ptq runs, so they're active during
+    calibration and export. No modelopt source files are modified.
+    """
 
     from modelopt.torch.quantization.nn.modules import tensor_quantizer as tq_module
     from modelopt.torch.quantization.qtensor import nvfp4_tensor
 
     # ── Patch 1: load_calib_amax — force _amax to CPU after calibration ──
-    #
-    # load_calib_amax() is called by max_calibrate() after the forward loop
-    # finishes. It writes _amax to GPU by default. We patch it so _amax
-    # goes to CPU immediately, preventing GPU corruption during the long
-    # wait before export.
     orig_load_calib_amax = tq_module.TensorQuantizer.load_calib_amax
 
     def patched_load_calib_amax(self, *args, **kwargs):
@@ -80,7 +75,6 @@ def apply_patches():
     print("✓ Patched TensorQuantizer.load_calib_amax (force _amax to CPU)")
 
     # ── Patch 2: export_amax — CPU safety ──
-    # If any _amax is still on GPU at export time, move it before reading.
     orig_export_amax = tq_module.TensorQuantizer.export_amax
 
     def patched_export_amax(self):
@@ -122,13 +116,7 @@ def apply_patches():
 
 
 def snapshot_amax_to_cpu(model, snapshot_path):
-    """Walk all quantizers, copy their _amax to CPU, save to disk.
-
-    After calibration completes, the _amax tensors are fresh and valid on GPU.
-    We copy them to CPU immediately and save to disk. This costs almost nothing
-    (~50MB for ~49K quantizers) but guarantees we have valid calibration data
-    even if CUDA corrupts the GPU copies later.
-    """
+    """Walk all quantizers, copy _amax to CPU, save to disk."""
     from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
 
     print(f"\nSnapshotting quantizer _amax to CPU...")
@@ -149,7 +137,6 @@ def snapshot_amax_to_cpu(model, snapshot_path):
     size_mb = os.path.getsize(snapshot_path) / (1024**2)
     print(f"✓ Snapshotted {n_moved} quantizer _amax tensors to CPU ({time.time()-t0:.1f}s)")
     print(f"  Saved to: {snapshot_path} ({size_mb:.1f} MB)")
-
     return snapshots
 
 
@@ -172,7 +159,7 @@ def restore_amax_from_snapshot(model, snapshot_path):
 
 
 def force_all_amax_to_cpu(model):
-    """Force ALL quantizer tensors to CPU. Nuclear option after calibration."""
+    """Force ALL quantizer tensors to CPU."""
     from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
 
     count = 0
@@ -195,12 +182,10 @@ def save_calibrated_state(model, path):
     print(f"{'='*60}")
 
     start = time.time()
-
     state = {
         'model_state_dict': model.state_dict(),
         'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
     }
-
     torch.save(state, path)
     size_gb = os.path.getsize(path) / (1024**3)
     print(f"✓ Saved calibrated state: {size_gb:.1f} GB ({time.time()-start:.0f}s)")
@@ -209,7 +194,7 @@ def save_calibrated_state(model, path):
 
 
 def run_calibration(model_path, export_dir, calib_save_path, amax_snapshot_path, calib_size, calib_seq):
-    """Full pipeline: load → quantize → calibrate → snapshot → save → export."""
+    """Full pipeline: parse args via hf_ptq → load → quantize → snapshot → save → export."""
 
     os.chdir(EXAMPLE_DIR)
     sys.path.insert(0, EXAMPLE_DIR)
@@ -217,150 +202,68 @@ def run_calibration(model_path, export_dir, calib_save_path, amax_snapshot_path,
     os.environ["HF_TOKEN"] = HF_TOKEN
     os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
 
-    # Import from hf_ptq and modelopt — all verified against the example script
-    from example_utils import get_model, get_tokenizer
-    from hf_ptq import (
-        make_calib_dataloader,
-        build_quant_cfg,
-        load_mtp_weights,
-        copy_custom_model_files,
-        QUANT_CFG_CHOICES,
-        KV_QUANT_CFG_CHOICES,
-    )
-    from modelopt.torch import quantization as mtq
-    from modelopt.torch.quantization.config import need_calibration
-    from modelopt.torch.utils.dataset_utils import get_max_batch_size
-    from modelopt.torch.export import export_hf_checkpoint
+    from hf_ptq import parse_args, main as hf_main
 
     apply_patches()
 
-    # ── Load model ──
-    # Use modelopt's get_model() — handles max_memory properly for 3TB model.
-    # Raw AutoModelForCausalLM.from_pretrained OOMs during expert weight conversion.
-    print(f"\nLoading model from {model_path}...")
-    t0 = time.time()
+    # ── Build args using hf_ptq's own parser ──
+    # This guarantees ALL attributes exist with correct defaults.
+    # We temporarily replace sys.argv so parse_args() sees our config.
+    saved_argv = sys.argv
+    sys.argv = [
+        "hf_ptq.py",
+        "--model", model_path,
+        "--quant", QUANT,
+        "--calib", str(calib_size),
+        "--calib_seq", str(calib_seq),
+        "--kv_cache_quant", KV_CACHE_QUANT,
+        "--tp", str(TP),
+        "--export_path", export_dir,
+        "--trust_remote_code",
+        "--use_seq_device_map",
+        "--gpu_max_mem_percentage", str(GPU_MEM_PCT),
+        "--batch_size", "0",
+    ]
+    args = parse_args()
+    sys.argv = saved_argv
 
-    model = get_model(
-        model_path,
-        gpu_mem_percentage=GPU_MEM_PCT,
-        trust_remote_code=True,
-        use_seq_device_map=True,
-    )
-    tokenizer = get_tokenizer(model_path, trust_remote_code=True)
-    print(f"✓ Model loaded in {time.time()-t0:.0f}s")
+    # ── Post-calibration hook ──
+    # We monkey-patch export_quantized to add our defensive saves before export.
+    import hf_ptq
 
-    # ── Setup quantization config ──
-    # Same flow as hf_ptq's quantize_main()
-    quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[QUANT])
-    quant_cfg = build_quant_cfg(QUANT, quant_cfg, None, None, None)
+    orig_export_quantized = hf_ptq.export_quantized
 
-    if KV_CACHE_QUANT != "none":
-        enable_quant_kv_cache = True
-        print(f"✓ KV cache quantization: {KV_CACHE_QUANT}")
-        quant_cfg = mtq.update_quant_cfg_with_kv_cache_quant(
-            quant_cfg,
-            getattr(mtq, KV_QUANT_CFG_CHOICES[KV_CACHE_QUANT])["quant_cfg"],
-        )
-    else:
-        enable_quant_kv_cache = False
+    def patched_export_quantized(exp_args, full_model, language_model, model_type,
+                                  tokenizer, default_padding_side, default_pad_token):
+        """Wrapper that snapshots amax and saves state before calling the real export."""
+        print("\n" + "="*60)
+        print("POST-CALIBRATION: Snapshotting amax and saving state")
+        print("="*60)
 
-    # ── Detect batch size ──
-    # Same as hf_ptq's quantize_main()
-    print("\nDetecting max calibration batch size...")
-    batch_size = get_max_batch_size(
-        model,
-        max_sample_length=calib_seq,
-        sample_memory_usage_ratio=1.1,
-    )
-    batch_size = min(batch_size, calib_size)
-    print(f"✓ Using calibration batch_size={batch_size}")
+        # Snapshot amax to CPU
+        snapshot_amax_to_cpu(language_model, amax_snapshot_path)
 
-    # ── Prepare dataloader ──
-    # Same args structure as hf_ptq
-    args = argparse.Namespace(
-        calib_size=[calib_size],
-        calib_seq=calib_seq,
-        calib_dataset="",
-        dataset=None,  # None triggers default: ["cnn_dailymail", "nemotron-post-training-dataset-v2"]
-        batch_size=batch_size,
-        calib_batch_size=0,
-        calib_with_images=False,
-        auto_quantize_bits=None,
-        auto_quantize_method=None,
-        specdec_offline_dataset=None,
-        inference_pipeline_parallel=1,
-    )
-    calib_dataloader, _ = make_calib_dataloader(
-        args, model, None, tokenizer, torch.device("cuda"), None,
-    )
+        # Force all quantizer state to CPU
+        force_all_amax_to_cpu(language_model)
 
-    # ── Quantize + Calibrate ──
-    print(f"\n{'='*60}")
-    print(f"QUANTIZING: {QUANT} with {calib_size} calibration samples")
-    print(f"{'='*60}")
-    t0 = time.time()
+        # Free GPU memory
+        torch.cuda.empty_cache()
+        gc.collect()
 
-    model = mtq.quantize(model, quant_cfg, forward_loop=calib_dataloader)
+        # Save calibrated state
+        save_calibrated_state(language_model, calib_save_path)
 
-    print(f"✓ Quantization + calibration complete in {time.time()-t0:.0f}s")
+        # Now run the real export
+        orig_export_quantized(exp_args, full_model, language_model, model_type,
+                             tokenizer, default_padding_side, default_pad_token)
 
-    # ── IMMEDIATELY snapshot all _amax to CPU ──
-    snapshots = snapshot_amax_to_cpu(model, amax_snapshot_path)
+    hf_ptq.export_quantized = patched_export_quantized
+    print("✓ Hooked export_quantized with amax snapshot + state save")
 
-    # ── Force ALL quantizer state to CPU ──
-    force_all_amax_to_cpu(model)
-
-    # ── Free GPU memory ──
-    torch.cuda.empty_cache()
-    gc.collect()
-
-    # ── SAVE STATE ──
-    save_calibrated_state(model, calib_save_path)
-
-    # ── Export ──
-    run_export(model, tokenizer, model_path, export_dir, amax_snapshot_path)
-
-
-def run_export(model, tokenizer, model_path, export_dir, amax_snapshot_path=None):
-    """Export the quantized model to HF safetensors format."""
-    from modelopt.torch.export import export_hf_checkpoint
-    from hf_ptq import load_mtp_weights, copy_custom_model_files
-
-    print(f"\n{'='*60}")
-    print(f"EXPORTING → {export_dir}")
-    print(f"{'='*60}")
-
-    force_all_amax_to_cpu(model)
-    if amax_snapshot_path and os.path.exists(amax_snapshot_path):
-        restore_amax_from_snapshot(model, amax_snapshot_path)
-
-    torch.cuda.empty_cache()
-    gc.collect()
-
-    t0 = time.time()
-
-    try:
-        mtp_layer_prefixes, mtp_state_dict = load_mtp_weights(model, model_path)
-        if mtp_layer_prefixes:
-            model._mtp_layer_prefixes = mtp_layer_prefixes
-
-        export_hf_checkpoint(
-            model,
-            export_dir=export_dir,
-            extra_state_dict=mtp_state_dict,
-        )
-
-        tokenizer.save_pretrained(export_dir)
-        copy_custom_model_files(model_path, export_dir, True)
-
-        print(f"\n✓ Export complete in {time.time()-t0:.0f}s → {export_dir}")
-
-    except Exception as e:
-        print(f"\n✗ EXPORT FAILED: {e}")
-        print(f"  Calibrated state: {CALIB_SAVE_PATH}")
-        print(f"  Amax snapshots: {AMAX_SNAPSHOT_PATH}")
-        print(f"  Re-run with --export-only to retry")
-        raise
+    # ── Run hf_ptq's full pipeline ──
+    # This handles model loading, quantization, calibration, and export
+    # using the exact same code path as the shell script.
+    hf_main(args)
 
 
 def run_export_only(calib_save_path, amax_snapshot_path, model_path, export_dir):
@@ -389,7 +292,35 @@ def run_export_only(calib_save_path, amax_snapshot_path, model_path, export_dir)
     model.load_state_dict(state['model_state_dict'])
     print(f"✓ Loaded calibrated state (saved at {state['timestamp']})")
 
-    run_export(model, tokenizer, model_path, export_dir, amax_snapshot_path)
+    force_all_amax_to_cpu(model)
+    if amax_snapshot_path and os.path.exists(amax_snapshot_path):
+        restore_amax_from_snapshot(model, amax_snapshot_path)
+
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    from modelopt.torch.export import export_hf_checkpoint
+    from hf_ptq import load_mtp_weights, copy_custom_model_files
+
+    print(f"\n{'='*60}")
+    print(f"EXPORTING → {export_dir}")
+    print(f"{'='*60}")
+
+    t0 = time.time()
+    try:
+        mtp_layer_prefixes, mtp_state_dict = load_mtp_weights(model, model_path)
+        if mtp_layer_prefixes:
+            model._mtp_layer_prefixes = mtp_layer_prefixes
+
+        export_hf_checkpoint(model, export_dir=export_dir, extra_state_dict=mtp_state_dict)
+        tokenizer.save_pretrained(export_dir)
+        copy_custom_model_files(model_path, export_dir, True)
+        print(f"\n✓ Export complete in {time.time()-t0:.0f}s → {export_dir}")
+    except Exception as e:
+        print(f"\n✗ EXPORT FAILED: {e}")
+        print(f"  Calibrated state: {CALIB_SAVE_PATH}")
+        print(f"  Amax snapshots: {AMAX_SNAPSHOT_PATH}")
+        raise
 
 
 def run_validate(calib_save_path, amax_snapshot_path):
@@ -399,10 +330,7 @@ def run_validate(calib_save_path, amax_snapshot_path):
     if os.path.exists(amax_snapshot_path):
         snapshots = torch.load(amax_snapshot_path, map_location='cpu')
         n_total = len(snapshots)
-        n_valid = 0
-        n_zero = 0
-        n_nan = 0
-        n_neg = 0
+        n_valid = n_zero = n_nan = n_neg = 0
 
         for name, amax in snapshots.items():
             if torch.any(torch.isnan(amax)):
@@ -415,12 +343,7 @@ def run_validate(calib_save_path, amax_snapshot_path):
                 n_valid += 1
 
         print(f"\nAmax snapshot validation:")
-        print(f"  Total quantizers: {n_total}")
-        print(f"  Valid: {n_valid}")
-        print(f"  All zeros: {n_zero}")
-        print(f"  Negative: {n_neg}")
-        print(f"  NaN: {n_nan}")
-
+        print(f"  Total: {n_total}  Valid: {n_valid}  Zero: {n_zero}  Neg: {n_neg}  NaN: {n_nan}")
         if n_valid == n_total:
             print(f"\n✓ All {n_total} amax snapshots are valid!")
         else: