diff --git a/scripts/quantize_nvfp4.py b/scripts/quantize_nvfp4.py index 1687536..c1462ac 100644 --- a/scripts/quantize_nvfp4.py +++ b/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: