diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 47d6a8b2c19..b54d701647d 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -525,6 +525,29 @@ def sparsity_main( mts.export(full_model) +def _restore_quantized_state_if_requested(args: argparse.Namespace, full_model: torch.nn.Module) -> bool: + """Restore a previously-calibrated ModelOpt state and skip calibration. + + Returns True if a restore happened (caller should skip mono_quantize). + """ + if args.restore_quantized_state is None: + return False + print( + f"Restoring quantized state from {args.restore_quantized_state}; skipping calibration." + ) + mto.restore(full_model, args.restore_quantized_state) + return True + + +def _save_quantized_state_if_requested(args: argparse.Namespace, full_model: torch.nn.Module) -> None: + """Save the just-calibrated ModelOpt state, if requested, so a later export-only retry can + restore it via --restore_quantized_state instead of repeating calibration.""" + if args.save_quantized_state is None: + return + print(f"Saving quantized state to {args.save_quantized_state}") + mto.save(full_model, args.save_quantized_state) + + def mono_quantize( args: argparse.Namespace, quant_cfg: dict[str, Any], @@ -943,6 +966,17 @@ def quantize_main( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, ): + if _restore_quantized_state_if_requested(args, full_model): + # Restore mode retries a failed/interrupted export from a previously calibrated state, + # so none of the calibration-only work below (batch-size probing, calibration + # dataloader construction, the pre-quantize generation preview) is needed -- go + # straight to export, with no before/after generation comparison. + export_quantized(args, full_model, language_model, model_type) + return + + # Detect if this is a Nemotron VL model using architecture-based detection. + is_nemotron_vl_model = is_nemotron_vl(full_model) + # Load the recipe up front so we can detect layerwise calibration before batch-size probing. recipe = None if args.recipe is not None: @@ -1072,9 +1106,6 @@ def quantize_main( ), ) - # Detect if this is a Nemotron VL model using architecture-based detection - is_nemotron_vl_model = is_nemotron_vl(full_model) - preview_input_ids, preview_attention_mask, generated_ids_before_ptq = pre_quantize( args, full_model, model_type, tokenizer, calib_dataloader, is_nemotron_vl_model ) @@ -1126,6 +1157,7 @@ def quantize_main( quant_cfg = _prepare_quant_cfg(args, quant_cfg, full_model) + quantized_this_run = bool(quant_cfg) if quant_cfg: mono_quantize( args, @@ -1150,6 +1182,10 @@ def quantize_main( # checkpoint directory main() downloaded. apply_cast_mxfp4_to_nvfp4(language_model, args.local_checkpoint_path) + # Save after the cast (if any) so a restore later reflects what actually gets exported. + if quantized_this_run: + _save_quantized_state_if_requested(args, full_model) + post_quantize( args, full_model, @@ -1374,6 +1410,27 @@ def parse_args() -> argparse.Namespace: "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe." ), ) + parser.add_argument( + "--save_quantized_state", + type=str, + default=None, + help=( + "Path to save the calibrated model (full weights plus ModelOpt state, so expect a " + "file as large as the model) after calibration completes, before export. Lets a " + "failed or interrupted export be retried via --restore_quantized_state without " + "repeating calibration. Plain (non-AutoQuantize) recipe/qformat path only." + ), + ) + parser.add_argument( + "--restore_quantized_state", + type=str, + default=None, + help=( + "Path to a ModelOpt state previously written by --save_quantized_state. When set, " + "calibration is skipped entirely and the saved state is restored onto the model " + "before proceeding straight to export." + ), + ) parser.add_argument( "--kv_auto_quantize_checkpoint", type=str, @@ -1476,6 +1533,19 @@ def parse_args() -> argparse.Namespace: "--low_memory_mode does not support --recipe; the low-memory loader initializes " "quantizers from --qformat/--kv_cache_qformat." ) + if args.save_quantized_state is not None and args.restore_quantized_state is not None: + parser.error( + "--save_quantized_state and --restore_quantized_state are mutually exclusive: " + "restoring a saved state skips calibration, so there is nothing new to save." + ) + if ( + args.save_quantized_state is not None or args.restore_quantized_state is not None + ) and _recipe_is_auto_quantize(args.recipe): + parser.error( + "--save_quantized_state/--restore_quantized_state are only supported for the plain " + "(non-AutoQuantize) recipe/qformat path; AutoQuantize's search state is not a single " + "quantized model state." + ) if args.use_fsdp2 and args.use_seq_device_map: warnings.warn("--use_seq_device_map is ignored when --use_fsdp2 is set.") args.use_seq_device_map = False diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 2f0b425de40..1a493c6d32b 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -68,6 +68,184 @@ def _parse_hf_ptq_args(monkeypatch, *args): return hf_ptq, parsed_args +def test_save_quantized_state_calls_mto_save_with_model_and_path(monkeypatch, tmp_path): + """--save_quantized_state PATH must hand mto.save exactly (model, PATH) -- this is the + round-trip a retried export depends on, so the call shape has to be exact.""" + state_path = str(tmp_path / "calib.pt") + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--save_quantized_state", state_path + ) + calls = [] + monkeypatch.setattr(hf_ptq.mto, "save", lambda model, path: calls.append((model, path))) + + dummy_model = object() + hf_ptq._save_quantized_state_if_requested(args, dummy_model) + + assert calls == [(dummy_model, state_path)] + + +def test_save_and_restore_quantized_state_together_is_rejected(monkeypatch): + """Both flags set is ambiguous -- restoring skips calibration, so there's nothing new to + save. Must fail loudly, at the CLI boundary, not silently prefer one and drop the other.""" + with pytest.raises(SystemExit): + _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--save_quantized_state", + "/tmp/out.pt", + "--restore_quantized_state", + "/tmp/in.pt", + ) + + +def test_save_quantized_state_with_auto_quantize_recipe_is_rejected(monkeypatch): + """AutoQuantize's search state is not a single quantized model state, so persistence flags + must be rejected before calibration starts, not discovered mid-run.""" + with pytest.raises(SystemExit): + _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--recipe", + "general/auto_quantize/nvfp4_fp8_at_5p4bits", + "--save_quantized_state", + "/tmp/out.pt", + ) + + +def test_save_quantized_state_round_trips_through_real_mto(monkeypatch, tmp_path): + """The mocked test above pins the call *shape*; this exercises the real ModelOpt + persistence path end to end: quantize a small local model, save its state through + hf_ptq's helper, restore it into a fresh copy, and confirm the restored model carries + the calibrated quantizer state and is usable (a forward pass runs cleanly).""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--save_quantized_state", str(tmp_path / "state.pt") + ) + + model = torch.nn.Linear(4, 4) + hf_ptq.mtq.quantize( + model, QUANT_CFG_CHOICES["fp8"], lambda m: m(torch.randn(2, 4)) + ) + expected_amax = model.weight_quantizer.amax.clone() + + hf_ptq._save_quantized_state_if_requested(args, model) + assert (tmp_path / "state.pt").exists() + + restored = torch.nn.Linear(4, 4) + restore_args = SimpleNamespace( + restore_quantized_state=str(tmp_path / "state.pt"), save_quantized_state=None + ) + assert hf_ptq._restore_quantized_state_if_requested(restore_args, restored) is True + + assert torch.equal(restored.weight_quantizer.amax, expected_amax) + restored(torch.randn(2, 4)) # usable: forward runs without error post-restore + + +def test_restore_quantized_state_skips_calibration_and_exports(monkeypatch): + """--restore_quantized_state must skip batch-size probing, calibration dataloader + construction, and the pre-quantize generation preview entirely, and go straight to + export -- retrying a failed export must not depend on the original calibration dataset + still being available.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--restore_quantized_state", "/tmp/in.pt" + ) + args.verbose = False + + monkeypatch.setattr(hf_ptq.mto, "restore", lambda model, path: None) + + def _fail(*_args, **_kwargs): + pytest.fail("calibration-only setup must be skipped in restore mode") + + monkeypatch.setattr(hf_ptq, "make_calib_dataloader", _fail) + monkeypatch.setattr(hf_ptq, "pre_quantize", _fail) + monkeypatch.setattr(hf_ptq, "mono_quantize", _fail) + monkeypatch.setattr(hf_ptq, "auto_quantize", _fail) + monkeypatch.setattr(hf_ptq, "get_max_batch_size", _fail) + + export_calls = [] + monkeypatch.setattr( + hf_ptq, + "export_quantized", + lambda args, full_model, language_model, model_type: export_calls.append(full_model), + ) + + full_model = object() + hf_ptq.quantize_main( + args, + full_model, + object(), # language_model + "llama", + False, # calibration_only + None, # processor + None, # tokenizer + "cuda", # device + ) + + assert export_calls == [full_model] + + +def test_save_quantized_state_happens_after_mxfp4_to_nvfp4_cast(monkeypatch): + """--cast_mxfp4_to_nvfp4 overwrites each weight quantizer's global_amax after calibration. + Saving before that cast would capture a pre-cast snapshot, so a later + --restore_quantized_state retry would export a model that never got the cast applied. + Save must run after the cast.""" + state_path = "/tmp/out.pt" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--kv_cache_qformat", + "none", + "--batch_size", + "1", + "--cast_mxfp4_to_nvfp4", + "--save_quantized_state", + state_path, + ) + args.verbose = False + + monkeypatch.setattr(hf_ptq, "is_nemotron_vl", lambda model: False) + monkeypatch.setattr( + hf_ptq, "make_calib_dataloader", lambda *a, **k: (object(), None) + ) + monkeypatch.setattr(hf_ptq, "pre_quantize", lambda *a, **k: (None, None, None)) + args.local_checkpoint_path = "dummy" + + call_order = [] + monkeypatch.setattr( + hf_ptq, "mono_quantize", lambda *a, **k: call_order.append("mono_quantize") + ) + monkeypatch.setattr( + hf_ptq, + "apply_cast_mxfp4_to_nvfp4", + lambda *a, **k: call_order.append("apply_cast_mxfp4_to_nvfp4"), + ) + monkeypatch.setattr( + hf_ptq, + "_save_quantized_state_if_requested", + lambda *a, **k: call_order.append("_save_quantized_state_if_requested"), + ) + monkeypatch.setattr(hf_ptq, "post_quantize", lambda *a, **k: None) + + hf_ptq.quantize_main( + args, + object(), # full_model + object(), # language_model + "llama", + False, # calibration_only + None, # processor + None, # tokenizer + "cuda", # device + ) + + assert call_order == [ + "mono_quantize", + "apply_cast_mxfp4_to_nvfp4", + "_save_quantized_state_if_requested", + ] + + def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys): hf_ptq = _import_hf_ptq(monkeypatch) monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"])