hf_ptq: save/restore quantized state to retry export without recalibrating

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

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