From 00f8d40ebc222a968a6dabcfa07def7c6c40a5e5 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Mon, 10 Aug 2026 08:49:45 -0400
Subject: [PATCH 1/9] PTQ reliability fixes: checkpoint resume +
 exclude_modules sentinel corruption

Three fixes, found while quantizing nvidia/Nemotron-Cascade-2-30B-A3B to NVFP4:

1. examples/hf_ptq/hf_ptq.py: --save_quantized_state/--restore_quantized_state
   flags. Lets a failed *export* be retried without redoing calibration.

2. modelopt/torch/utils/dataset_utils.py: checkpoint_every/checkpoint_fn hook in
   _forward_loop/create_forward_loop. Periodic-checkpoint mechanism (not yet
   wired into hf_ptq.py's calibration call -- follow-up, not done here).

3. modelopt/torch/export/quant_aware_conversion.py: build_reverse_name_mapper's
   sentinel-strip logic (_map) silently accepted a mangled sentinel instead of
   raising, when a reverse-rename rule's `.` (matched as "any path separator
   char") ate into the appended sentinel before it could be stripped. In our
   export this corrupted 77/77 exclude_modules entries with an unsubstituted
   placeholder (`\x00backbone.pt_name_sentinel`, not matching the real sentinel
   `\x00modelopt_name_sentinel` -- direct evidence a rename rule rewrote it).
   vLLM couldn't match any excluded layer, defaulted every quant-aware fused
   linear (e.g. Mamba's in_proj, a MergedColumnParallelLinear) to quantized-width
   allocation, and crashed loading the correctly-unquantized-but-unlabeled
   weight. Extracted the strip logic into _strip_sentinel_or_raise(), which now
   raises QuantConversionUnsupportedError on a failed strip -- exactly the
   exception build_reverse_name_mapper's own docstring already promised for
   this case. No caller changes needed: both call sites in unified_export_hf.py
   already wrap this in a broad try/except that falls back to safe in-memory
   names with a warning: this fix just makes that safety net actually trigger,
   converting a silent, deployment-time failure (bare AssertionError in vLLM,
   no useful diagnostic) into a loud, immediate one at export time.

All three built test-first (Red -> Green). 26 total tests passing across the
two touched test files (2 new for the checkpoint flags, 3 new for the sentinel
fix), no regressions on pre-existing tests.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BZby2EsEkVQKwwNKFDpwBm
---
 examples/hf_ptq/hf_ptq.py                     | 54 ++++++++++++++++++-
 .../torch/export/quant_aware_conversion.py    | 30 ++++++++++-
 modelopt/torch/utils/dataset_utils.py         | 20 ++++++-
 tests/examples/hf_ptq/test_hf_ptq_args.py     | 33 ++++++++++++
 .../export/test_quant_aware_conversion.py     | 29 ++++++++++
 tests/unit/torch/utils/test_dataset_utils.py  | 51 ++++++++++++++++++
 6 files changed, 213 insertions(+), 4 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 7a2328d10f7..9ad94564ac2 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -755,6 +755,34 @@ 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
+    if args.save_quantized_state is not None:
+        raise ValueError(
+            "--save_quantized_state and --restore_quantized_state are mutually exclusive: "
+            "restoring a saved state skips calibration, so there is nothing new to save."
+        )
+    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],
@@ -1324,7 +1352,9 @@ def _is_layerwise(obj):
             quant_cfg = copy.deepcopy(quant_cfg)
             force_weight_quantizers_static(quant_cfg["quant_cfg"])
 
-        if quant_cfg:
+        if _restore_quantized_state_if_requested(args, full_model):
+            pass
+        elif quant_cfg:
             mono_quantize(
                 args,
                 quant_cfg,
@@ -1335,6 +1365,7 @@ def _is_layerwise(obj):
                 calib_dataloader,
                 is_nemotron_vl_model,
             )
+            _save_quantized_state_if_requested(args, full_model)
         else:
             assert model_type != "dbrx", f"Does not support export {model_type} without quantizaton"
             print(f"qformat: {args.qformat}. No quantization applied, export {device} model")
@@ -1563,6 +1594,27 @@ def parse_args() -> argparse.Namespace:
             "deprecated --auto_quantize_bits CLI path."
         ),
     )
+    parser.add_argument(
+        "--save_quantized_state",
+        type=str,
+        default=None,
+        help=(
+            "Path to save the calibrated/quantized model's ModelOpt state 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."
+        ),
+    )
     # Deprecated AutoQuantize CLI flags: kept as a backward-compat shim that converts them into an
     # AutoQuantizeConfig on the fly (see _auto_quantize_config_from_cli). Prefer --recipe. The old
     # CLI also lives on the 0.45 branch.
diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 2e32f123869..050b3cef603 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -233,6 +233,34 @@ def revert_weight_conversion_quant_aware(model, state_dict: dict[str, torch.Tens
     return apply_reverse_rules(state_dict, split_rules, rename_rules)
 
 
+def _strip_sentinel_or_raise(mapped: str, sentinel: str, original: str) -> str:
+    """Strip ``sentinel`` from the end of ``mapped``, or raise if a rename rule mangled it.
+
+    ``mapped`` is the result of running the reverse rename rules against
+    ``base + sentinel`` (see :func:`build_reverse_name_mapper`). Rename patterns use
+    ``.`` as "any path-separator char", so a rule whose match extends further than
+    intended -- most commonly a greedy ``.`` -- can consume or rewrite part of the
+    sentinel instead of leaving it as an untouched trailing segment. When that happens
+    ``str.removesuffix`` is a silent no-op (it only strips an *exact* suffix match), and
+    the mangled remnant would otherwise leak into the exported ``exclude_modules``
+    name, which no real checkpoint tensor can ever match.
+
+    Raises:
+        QuantConversionUnsupportedError: the sentinel did not survive as a clean
+            suffix, so this name mapping cannot be trusted. The caller
+            (:func:`build_reverse_name_mapper`'s callers) already treats this
+            exception as "reverse conversion failed, fall back to in-memory names for
+            both weights and config" -- so raising here converts a silent, downstream
+            (deployment-time) failure into a loud, immediate one at export time.
+    """
+    if not mapped.endswith(sentinel):
+        raise QuantConversionUnsupportedError(
+            f"reverse rename mangled the internal sentinel while mapping {original!r} "
+            f"(got {mapped!r}); a rename rule likely matched past its intended boundary"
+        )
+    return mapped.removesuffix(sentinel)
+
+
 def build_reverse_name_mapper(model):
     """Build a ``str -> str`` mapper that applies the quant-aware reverse *rename* rules.
 
@@ -271,7 +299,7 @@ def _map(name: str) -> str:
         elif name.endswith("*"):
             base, suffix = name[:-1], "*"
         mapped = _apply(base + _sentinel)
-        mapped = mapped.removesuffix(_sentinel)
+        mapped = _strip_sentinel_or_raise(mapped, _sentinel, name)
         return mapped + suffix
 
     return _map
diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py
index 789308a0b5b..cf762d40f8f 100644
--- a/modelopt/torch/utils/dataset_utils.py
+++ b/modelopt/torch/utils/dataset_utils.py
@@ -1158,6 +1158,8 @@ def _forward_loop(
     model: torch.nn.Module,
     dataloader: DataLoader,
     allowed_non_tensor_keys: set | None = None,
+    checkpoint_every: int = 0,
+    checkpoint_fn: Callable[[], None] | None = None,
 ) -> None:
     """Runs forward passes through the model using data from the dataloader.
 
@@ -1165,6 +1167,11 @@ def _forward_loop(
         model: The PyTorch model to run inference on
         dataloader: DataLoader containing the batched input data
         allowed_non_tensor_keys: Set of key names whose values may be non-tensor types
+        checkpoint_every: If > 0, call `checkpoint_fn` after every this-many batches.
+            0 (the default) disables checkpointing entirely -- `checkpoint_fn` is never
+            called, matching prior behavior for existing callers.
+        checkpoint_fn: No-arg callback invoked periodically per `checkpoint_every`. Ignored
+            if `checkpoint_every` is 0.
     """
     with _disable_use_cache(model), torch.no_grad():
         is_enc_dec = model_type_is_enc_dec(model)
@@ -1173,11 +1180,13 @@ def _forward_loop(
         infer_method = model.generate if is_enc_dec else model
         max_working_batch_size = None  # Initialize max working batch size as None
 
-        for _, data in enumerate(tqdm(dataloader)):
+        for step, data in enumerate(tqdm(dataloader)):
             # Process batch and update max working batch size
             max_working_batch_size = _process_batch(
                 data, infer_method, max_working_batch_size, allowed_non_tensor_keys
             )
+            if checkpoint_every and (step + 1) % checkpoint_every == 0:
+                checkpoint_fn()
 
 
 def create_forward_loop(
@@ -1191,6 +1200,8 @@ def create_forward_loop(
     include_labels: bool = False,
     dataloader: DataLoader | None = None,
     allowed_non_tensor_keys: set | None = None,
+    checkpoint_every: int = 0,
+    checkpoint_fn: Callable[[], None] | None = None,
 ) -> Callable:
     """Creates and returns a forward loop function configured for a specific model, dataset, and tokenizer.
 
@@ -1212,6 +1223,9 @@ def create_forward_loop(
         allowed_non_tensor_keys: Set of key names whose batch values may be non-tensor types.
             Useful when the dataloader yields batches with non-standard fields (e.g., nested
             model outputs).
+        checkpoint_every: If > 0, checkpoint_fn is called after every this-many batches. 0
+            (the default) disables checkpointing.
+        checkpoint_fn: No-arg callback invoked periodically per checkpoint_every.
 
     Example usage for quantization:
 
@@ -1251,7 +1265,9 @@ def create_forward_loop(
             include_labels=include_labels,
         )
 
-    return lambda model: _forward_loop(model, dataloader, allowed_non_tensor_keys)
+    return lambda model: _forward_loop(
+        model, dataloader, allowed_non_tensor_keys, checkpoint_every, checkpoint_fn
+    )
 
 
 def model_type_is_enc_dec(model):
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 419ec4c3025..d9156a31eb4 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -55,6 +55,39 @@ 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):
+    """--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."""
+    hf_ptq, args = _parse_hf_ptq_args(
+        monkeypatch, "--pyt_ckpt_path", "dummy", "--save_quantized_state", "/tmp/calib.pt"
+    )
+    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, "/tmp/calib.pt")]
+
+
+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, not silently prefer one and drop the other."""
+    hf_ptq, args = _parse_hf_ptq_args(
+        monkeypatch,
+        "--pyt_ckpt_path",
+        "dummy",
+        "--save_quantized_state",
+        "/tmp/out.pt",
+        "--restore_quantized_state",
+        "/tmp/in.pt",
+    )
+    monkeypatch.setattr(hf_ptq.mto, "restore", lambda model, path: pytest.fail("must not restore"))
+
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        hf_ptq._restore_quantized_state_if_requested(args, object())
+
+
 def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs."""
     hf_ptq, args = _parse_hf_ptq_args(
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index e3aa22bf3d8..cab1c03c69c 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -32,6 +32,7 @@
     RenameRule,
     SplitRule,
     _assert_experts_pre_expanded,
+    _strip_sentinel_or_raise,
     apply_reverse_rules,
     build_reverse_name_mapper,
     revert_quant_config_names,
@@ -492,6 +493,34 @@ def test_stacked_experts_guard():
     _assert_experts_pre_expanded(bad, [])
 
 
+def test_strip_sentinel_clean_strip():
+    """A rename rule that leaves the sentinel untouched strips cleanly, no raise."""
+    sentinel = ".\x00modelopt_name_sentinel"
+    mapped = "backbone.layers.4.mixer.in_proj" + sentinel
+    assert _strip_sentinel_or_raise(mapped, sentinel, "original.name") == "backbone.layers.4.mixer.in_proj"
+
+
+def test_strip_sentinel_mangled_raises():
+    """A rename rule that eats into the appended sentinel (most commonly a greedy `.`,
+    since rename patterns treat `.` as "any path-separator char") must raise instead of
+    silently emitting a corrupted name.
+
+    Regression for the real failure: ModelOpt main (0.47.0.dev41) exported an NVFP4
+    Nemotron-Cascade-2-30B checkpoint whose exclude_modules were 100% corrupted with
+    this exact kind of mangled-sentinel remnant (`\\x00backbone.pt_name_sentinel`,
+    not matching the real sentinel `\\x00modelopt_name_sentinel` -- proof a rename rule
+    partially rewrote it). vLLM couldn't match any of the 77 entries, defaulted every
+    quant-aware fused linear to quantized-width allocation, and crashed loading the
+    correctly-unquantized-but-unlabeled weights.
+    """
+    sentinel = ".\x00modelopt_name_sentinel"
+    # Simulates a greedy rule (e.g. `re.sub(r"\.experts\..*", "", key)`) that consumed
+    # past its intended boundary and ate the sentinel instead of leaving it as a suffix.
+    mangled = "backbone.layers.4.mixer.in_proj.\x00backbone.pt_name_sentinel"
+    with pytest.raises(QuantConversionUnsupportedError, match="mangled the internal sentinel"):
+        _strip_sentinel_or_raise(mangled, sentinel, "backbone.layers.4.mixer.in_proj.*")
+
+
 def test_revert_quant_config_names_mapper():
     """exclude_modules / quantized_layers keys revert to hub names, preserving wildcards.
 
diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py
index ebdbfa04f02..6a998729168 100644
--- a/tests/unit/torch/utils/test_dataset_utils.py
+++ b/tests/unit/torch/utils/test_dataset_utils.py
@@ -28,6 +28,7 @@
     _iter_use_cache_configs,
     _pack_documents_into_rows,
     _process_batch,
+    create_forward_loop,
     get_dataset_dataloader,
     get_dataset_samples,
     get_max_batch_size,
@@ -289,6 +290,56 @@ def _collate(samples):
     assert model.config.use_cache is True
 
 
+def _tiny_loader(num_batches: int):
+    class _Model(torch.nn.Module):
+        def forward(self, **kwargs):
+            pass
+
+    def _collate(samples):
+        return {"input_ids": torch.stack([s["input_ids"] for s in samples])}
+
+    data = [{"input_ids": torch.zeros(4, dtype=torch.long)} for _ in range(num_batches)]
+    loader = DataLoader(data, batch_size=1, collate_fn=_collate)
+    return _Model(), loader
+
+
+def test_forward_loop_checkpoints_every_n_steps():
+    """checkpoint_fn fires every `checkpoint_every` batches -- 5 batches at every=2 means
+    two fires (after batch 2 and batch 4), not one per batch and not a trailing fire for
+    the final partial group."""
+    calls: list[int] = []
+    model, loader = _tiny_loader(5)
+
+    _forward_loop(model, loader, checkpoint_every=2, checkpoint_fn=lambda: calls.append(1))
+
+    assert len(calls) == 2
+
+
+def test_forward_loop_checkpoint_every_zero_never_fires():
+    """checkpoint_every=0 (the default) must never call checkpoint_fn -- and must not raise,
+    guarding against a naive `step % checkpoint_every` hitting ZeroDivisionError."""
+    calls: list[int] = []
+    model, loader = _tiny_loader(3)
+
+    _forward_loop(model, loader, checkpoint_fn=lambda: calls.append(1))
+
+    assert calls == []
+
+
+def test_create_forward_loop_passes_checkpoint_args_through():
+    """create_forward_loop's returned closure must forward checkpoint_every/checkpoint_fn
+    to _forward_loop -- otherwise the capability is inert for actual callers like hf_ptq.py."""
+    calls: list[int] = []
+    model, loader = _tiny_loader(4)
+
+    forward_loop = create_forward_loop(
+        dataloader=loader, checkpoint_every=2, checkpoint_fn=lambda: calls.append(1)
+    )
+    forward_loop(model)
+
+    assert len(calls) == 2
+
+
 def test_disable_use_cache_restores_on_exception():
     """Restore must run even if the with-block raises."""
     model = torch.nn.Linear(4, 4)

From 79c93985890ea5a42d8d647299f2690d277861b5 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Mon, 10 Aug 2026 14:06:01 -0400
Subject: [PATCH 2/9] Restore NemotronH's legacy
 hybrid_override_pattern/num_hidden_layers on export

Root cause: transformers now ships native NemotronH support
(transformers/models/nemotron_h/configuration_nemotron_h.py) with a newer
config schema -- layers_block_type as a real stored field, with
hybrid_override_pattern/num_hidden_layers as read-only properties derived
from it. Many NemotronH checkpoints on the Hub (e.g.
nvidia/Nemotron-Cascade-2-30B-A3B) still ship their own older, bundled
trust_remote_code configuration_nemotron_h.py written against the *opposite*
schema: hybrid_override_pattern/num_hidden_layers as real fields, with
layers_block_type as a computed property that has no setter.

When ModelOpt loads such a checkpoint for PTQ, transformers' AutoConfig
resolves the model via its own newer native class (no trust_remote_code
needed once a model_type is natively registered), silently converting the
config to the new schema in memory. save_pretrained's plain to_dict() only
serializes __dict__, so neither hybrid_override_pattern nor num_hidden_layers
(both properties on the native class) makes it into config.json -- but
layers_block_type does, since it's a real field there. The result: an
exported config.json in the new schema, sitting next to a copied
configuration_nemotron_h.py file that only understands the old one. Loading
the export back through its own bundled class then fails outright
(AttributeError: property 'layers_block_type' has no setter) or silently
loses the pattern/layer-count metadata.

Reproduced on three separate nvidia/Nemotron-Cascade-2-30B-A3B NVFP4 exports
in a row; each one required hand-patching config.json afterward.

Fix: sanitize_hf_config_for_deployment now reconstructs
hybrid_override_pattern from layers_block_type using the exact inverse of
transformers' own NemotronHConfig._list_to_pattern mapping (linear_attention
-> M, moe -> E, full_attention -> *, mlp -> -), restores num_hidden_layers as
len(layers_block_type), and drops the now-redundant layers_block_type /
mtp_layers_block_type fields the old schema's class doesn't expect. No-op for
non-NemotronH exports, no-op if hybrid_override_pattern is already present
(nothing was dropped), and warns without guessing if layers_block_type
contains an unrecognized value.

Built test-first (Red -> Green): 4 new tests (restores correctly, ignores
other model types, no-ops when already legacy schema, warns on unrecognized
layer type), 19/19 passing in the touched test file, 165/165 passing across
the full export test suite -- no regressions.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BZby2EsEkVQKwwNKFDpwBm
---
 .../export/plugins/hf_checkpoint_utils.py     | 63 +++++++++++++++-
 .../torch/export/test_hf_checkpoint_utils.py  | 74 +++++++++++++++++++
 2 files changed, 136 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
index ddae8e2b409..b7b392386dd 100644
--- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py
+++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
@@ -104,6 +104,63 @@ def _sanitize_llama3_rope_config(config_data: dict[str, Any], model: Any) -> Non
             rope_config["rope_theta"] = rope_theta
 
 
+# transformers' native NemotronHConfig (registered as model_type "nemotron_h") stores
+# layer types as a list and derives ``hybrid_override_pattern``/``num_hidden_layers`` as
+# read-only properties (transformers/models/nemotron_h/configuration_nemotron_h.py).
+# Plain ``to_dict()`` serialization -- used by ``save_pretrained`` -- only captures
+# ``__dict__``, so neither property makes it into the exported config.json. Many
+# NemotronH checkpoints on the Hub (e.g. nvidia/Nemotron-Cascade-2-30B-A3B) still ship
+# their own older, bundled trust_remote_code configuration_nemotron_h.py written against
+# the *opposite* schema: ``hybrid_override_pattern``/``num_hidden_layers`` as real
+# fields, with ``layers_block_type`` as a computed property that has no setter. Loading
+# such a checkpoint's exported config.json back through its own bundled class then fails
+# (or, if it doesn't fail outright, silently loses the pattern/layer-count metadata).
+# Reproduced three times in a row on real nvidia/Nemotron-Cascade-2-30B-A3B NVFP4
+# exports. This mapping is the exact inverse of transformers' own
+# ``NemotronHConfig._list_to_pattern``, so it reconstructs precisely what a native
+# transformers load would have derived.
+_NEMOTRON_H_PATTERN_CHAR_BY_LAYER_TYPE = {
+    "linear_attention": "M",
+    "moe": "E",
+    "full_attention": "*",
+    "mlp": "-",
+}
+
+
+def _restore_nemotron_h_legacy_schema_if_dropped(config_data: dict[str, Any]) -> None:
+    """Reconstruct NemotronH's legacy ``hybrid_override_pattern``/``num_hidden_layers``.
+
+    No-op for anything but a NemotronH export (``model_type == "nemotron_h"``), and a
+    no-op if ``hybrid_override_pattern`` is already present (nothing was dropped). If
+    ``layers_block_type`` contains a value outside the four known block types, this
+    warns and leaves ``config.json`` in the new schema rather than guessing.
+    """
+    if config_data.get("model_type") != "nemotron_h":
+        return
+    layer_types = config_data.get("layers_block_type")
+    if not isinstance(layer_types, list) or "hybrid_override_pattern" in config_data:
+        return
+
+    try:
+        pattern = "".join(
+            _NEMOTRON_H_PATTERN_CHAR_BY_LAYER_TYPE[layer_type] for layer_type in layer_types
+        )
+    except KeyError as e:
+        warnings.warn(
+            f"Cannot reconstruct NemotronH's legacy hybrid_override_pattern: unrecognized "
+            f"layers_block_type entry {e}. Leaving config.json in the new schema; a "
+            "bundled trust_remote_code configuration_nemotron_h.py written against the "
+            "legacy schema may fail to load it.",
+            stacklevel=2,
+        )
+        return
+
+    config_data["hybrid_override_pattern"] = pattern
+    config_data["num_hidden_layers"] = len(layer_types)
+    del config_data["layers_block_type"]
+    config_data.pop("mtp_layers_block_type", None)
+
+
 def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) -> None:
     """Sanitize exported Hugging Face config metadata for deployment runtimes.
 
@@ -111,9 +168,13 @@ def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) -
 
     * add missing llama3 ``rope_theta`` metadata when available;
     * trim trailing MTP/next-token-prediction ``layer_types`` entries only when
-      the mismatch is exactly explained by next-token-prediction metadata.
+      the mismatch is exactly explained by next-token-prediction metadata;
+    * reconstruct NemotronH's legacy ``hybrid_override_pattern``/``num_hidden_layers``
+      when a native-transformers export dropped them (see
+      :func:`_restore_nemotron_h_legacy_schema_if_dropped`).
     """
     _sanitize_llama3_rope_config(config_data, model)
+    _restore_nemotron_h_legacy_schema_if_dropped(config_data)
 
     num_hidden_layers = _as_nonnegative_int(config_data.get("num_hidden_layers"))
     layer_types = config_data.get("layer_types")
diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py
index f3be2564312..0dcfb5836b9 100644
--- a/tests/unit/torch/export/test_hf_checkpoint_utils.py
+++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py
@@ -261,6 +261,80 @@ def test_sanitize_hf_config_for_deployment_ignores_broad_mtp_prefix_only():
     assert config_data["layer_types"] == ["full_attention", "linear_attention", "nextn_predict"]
 
 
+def test_sanitize_hf_config_for_deployment_restores_nemotron_h_legacy_schema():
+    """Reconstruct hybrid_override_pattern/num_hidden_layers from layers_block_type.
+
+    transformers' native NemotronHConfig (registered as model_type "nemotron_h")
+    stores layer types as a list and derives hybrid_override_pattern/num_hidden_layers
+    as read-only properties, so plain to_dict() serialization drops both. Many
+    NemotronH checkpoints on the Hub still ship their own older, bundled
+    trust_remote_code configuration_nemotron_h.py that expects the opposite schema
+    (hybrid_override_pattern/num_hidden_layers as real fields, layers_block_type as a
+    computed property with no setter) -- loading such a checkpoint's config.json back
+    through its own bundled class then fails. Reproduced on a real
+    nvidia/Nemotron-Cascade-2-30B-A3B NVFP4 export three times in a row.
+    """
+    config_data = {
+        "model_type": "nemotron_h",
+        "layers_block_type": [
+            "linear_attention",
+            "moe",
+            "full_attention",
+            "mlp",
+        ],
+        "mtp_layers_block_type": ["full_attention", "moe"],
+    }
+
+    sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace())
+
+    assert config_data["hybrid_override_pattern"] == "ME*-"
+    assert config_data["num_hidden_layers"] == 4
+    assert "layers_block_type" not in config_data
+    assert "mtp_layers_block_type" not in config_data
+
+
+def test_sanitize_hf_config_for_deployment_ignores_non_nemotron_h_layers_block_type():
+    """Only NemotronH's specific legacy-schema gap is patched, not other models."""
+    config_data = {
+        "model_type": "some_other_model",
+        "layers_block_type": ["full_attention", "mlp"],
+    }
+
+    sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace())
+
+    assert "hybrid_override_pattern" not in config_data
+    assert config_data["layers_block_type"] == ["full_attention", "mlp"]
+
+
+def test_sanitize_hf_config_for_deployment_leaves_nemotron_h_alone_if_already_legacy():
+    """No-op when config.json already carries hybrid_override_pattern (nothing dropped)."""
+    config_data = {
+        "model_type": "nemotron_h",
+        "hybrid_override_pattern": "M-*",
+        "num_hidden_layers": 3,
+    }
+
+    sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace())
+
+    assert config_data["hybrid_override_pattern"] == "M-*"
+    assert config_data["num_hidden_layers"] == 3
+    assert "layers_block_type" not in config_data
+
+
+def test_sanitize_hf_config_for_deployment_warns_on_unrecognized_nemotron_h_layer_type():
+    """Unknown layer-type strings should warn and leave config.json in the new schema."""
+    config_data = {
+        "model_type": "nemotron_h",
+        "layers_block_type": ["full_attention", "some_future_block_type"],
+    }
+
+    with pytest.warns(UserWarning, match="Cannot reconstruct"):
+        sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace())
+
+    assert "hybrid_override_pattern" not in config_data
+    assert config_data["layers_block_type"] == ["full_attention", "some_future_block_type"]
+
+
 def test_sanitize_hf_config_for_deployment_keeps_unexplained_layer_type_mismatch():
     """Do not rewrite config when extra layer types are not explained by nextn metadata."""
     config_data = {

From dedbf0ca920f1f845b97b1b834c85f9fa6f00dac Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:27:32 -0400
Subject: [PATCH 3/9] hf_ptq: validate save/restore quantized-state flags at
 parse_args()

Reject --save_quantized_state and --restore_quantized_state together, and
reject either with an AutoQuantize recipe (or the deprecated
--auto_quantize_bits CLI path), at the CLI boundary instead of only at
runtime inside _restore_quantized_state_if_requested. AutoQuantize's search
state is not a single quantized model state, so persistence flags for it
must fail before calibration work begins, not mid-run.

Also exercises the save/restore rejection through parse_args() instead of
only the private helper, per review feedback on PR #2129.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JjBHecV84Y6tcMJWKiBRzn
Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 examples/hf_ptq/hf_ptq.py                 | 13 ++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py | 53 +++++++++++++++++------
 2 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 58d0383e845..e169b7841dc 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -1734,6 +1734,19 @@ def parse_args() -> argparse.Namespace:
             "--low_memory_mode does not support --recipe or AutoQuantize (--auto_quantize_bits); "
             "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 (
+        args.auto_quantize_bits is not None or _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 d9156a31eb4..ad36c588bda 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -72,20 +72,47 @@ def test_save_quantized_state_calls_mto_save_with_model_and_path(monkeypatch):
 
 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, not silently prefer one and drop the other."""
-    hf_ptq, args = _parse_hf_ptq_args(
-        monkeypatch,
-        "--pyt_ckpt_path",
-        "dummy",
-        "--save_quantized_state",
-        "/tmp/out.pt",
-        "--restore_quantized_state",
-        "/tmp/in.pt",
-    )
-    monkeypatch.setattr(hf_ptq.mto, "restore", lambda model, path: pytest.fail("must not restore"))
+    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",
+        )
+
 
-    with pytest.raises(ValueError, match="mutually exclusive"):
-        hf_ptq._restore_quantized_state_if_requested(args, object())
+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_restore_quantized_state_with_deprecated_auto_quantize_cli_is_rejected(monkeypatch):
+    """The deprecated --auto_quantize_bits CLI path is also AutoQuantize; it must be rejected
+    the same way as an AutoQuantize --recipe."""
+    with pytest.raises(SystemExit):
+        _parse_hf_ptq_args(
+            monkeypatch,
+            "--pyt_ckpt_path",
+            "dummy",
+            "--auto_quantize_bits",
+            "5.4",
+            "--restore_quantized_state",
+            "/tmp/in.pt",
+        )
 
 
 def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):

From 4efe68e393f202aef4a02b7cdd85d897f90bdf78 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:28:21 -0400
Subject: [PATCH 4/9] export: handle non-hashable layers_block_type entries in
 NemotronH restore

A malformed config.json (e.g. layers_block_type containing a nested list or
dict) raised TypeError from the dict lookup, which the existing KeyError
handler did not catch. That aborted export instead of warning and leaving
the schema untouched, same as the already-handled unrecognized-string case.

Catch TypeError alongside KeyError and add a regression test.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 .../torch/export/plugins/hf_checkpoint_utils.py  |  4 +++-
 .../torch/export/test_hf_checkpoint_utils.py     | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
index b7b392386dd..7d1e9fe9a00 100644
--- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py
+++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
@@ -145,7 +145,9 @@ def _restore_nemotron_h_legacy_schema_if_dropped(config_data: dict[str, Any]) ->
         pattern = "".join(
             _NEMOTRON_H_PATTERN_CHAR_BY_LAYER_TYPE[layer_type] for layer_type in layer_types
         )
-    except KeyError as e:
+    except (KeyError, TypeError) as e:
+        # KeyError: a recognized-but-unlisted string block type. TypeError: a malformed
+        # entry (list, dict, ...) that isn't even hashable for the dict lookup.
         warnings.warn(
             f"Cannot reconstruct NemotronH's legacy hybrid_override_pattern: unrecognized "
             f"layers_block_type entry {e}. Leaving config.json in the new schema; a "
diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py
index 0dcfb5836b9..6925d846a47 100644
--- a/tests/unit/torch/export/test_hf_checkpoint_utils.py
+++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py
@@ -335,6 +335,22 @@ def test_sanitize_hf_config_for_deployment_warns_on_unrecognized_nemotron_h_laye
     assert config_data["layers_block_type"] == ["full_attention", "some_future_block_type"]
 
 
+def test_sanitize_hf_config_for_deployment_warns_on_non_hashable_layer_type():
+    """A malformed entry (e.g. a nested list or dict from a corrupted config.json) isn't
+    hashable, so the dict lookup raises TypeError instead of KeyError. Must warn and leave
+    config.json in the new schema, not raise and abort export."""
+    config_data = {
+        "model_type": "nemotron_h",
+        "layers_block_type": ["full_attention", ["mlp"]],
+    }
+
+    with pytest.warns(UserWarning, match="Cannot reconstruct"):
+        sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace())
+
+    assert "hybrid_override_pattern" not in config_data
+    assert config_data["layers_block_type"] == ["full_attention", ["mlp"]]
+
+
 def test_sanitize_hf_config_for_deployment_keeps_unexplained_layer_type_mismatch():
     """Do not rewrite config when extra layer types are not explained by nextn metadata."""
     config_data = {

From 7f4c04b73454c06fe6242574bf6fedd7a646be0a Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:29:29 -0400
Subject: [PATCH 5/9] dataset_utils: validate checkpoint_every/checkpoint_fn at
 create_forward_loop boundary

A negative checkpoint_every stays truthy under the modulo check inside
_forward_loop, and a positive checkpoint_every with checkpoint_fn=None
would call None once the first interval completes -- both failures only
surface deep inside the loop, mid-run. Validate both at the public
create_forward_loop boundary instead, before the dataloader/closure are
built, and document that checkpoint_fn runs on every process and must be
rank-safe or collective if it persists shared state.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 modelopt/torch/utils/dataset_utils.py        | 14 ++++++++++++--
 tests/unit/torch/utils/test_dataset_utils.py | 15 +++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py
index cf762d40f8f..bdea6805115 100644
--- a/modelopt/torch/utils/dataset_utils.py
+++ b/modelopt/torch/utils/dataset_utils.py
@@ -1171,7 +1171,8 @@ def _forward_loop(
             0 (the default) disables checkpointing entirely -- `checkpoint_fn` is never
             called, matching prior behavior for existing callers.
         checkpoint_fn: No-arg callback invoked periodically per `checkpoint_every`. Ignored
-            if `checkpoint_every` is 0.
+            if `checkpoint_every` is 0. Runs on every process that executes the loop; must be
+            rank-safe or collective if it persists shared state.
     """
     with _disable_use_cache(model), torch.no_grad():
         is_enc_dec = model_type_is_enc_dec(model)
@@ -1225,7 +1226,11 @@ def create_forward_loop(
             model outputs).
         checkpoint_every: If > 0, checkpoint_fn is called after every this-many batches. 0
             (the default) disables checkpointing.
-        checkpoint_fn: No-arg callback invoked periodically per checkpoint_every.
+        checkpoint_fn: No-arg callback invoked periodically per checkpoint_every. Runs on every
+            process that executes the forward loop; under a multi-process/distributed run,
+            the callback itself must be rank-safe (e.g. guard writes with a rank check) or
+            collective (e.g. an all-reduce/barrier) if it persists shared state such as a
+            checkpoint file.
 
     Example usage for quantization:
 
@@ -1249,6 +1254,11 @@ def create_forward_loop(
         A forward loop function that can be called with no arguments. When called, this function iterates over
             the dataset specified by `dataset_name`.
     """
+    if checkpoint_every < 0:
+        raise ValueError(f"checkpoint_every must be non-negative, got {checkpoint_every}")
+    if checkpoint_every > 0 and not callable(checkpoint_fn):
+        raise ValueError("checkpoint_fn must be callable when checkpoint_every > 0")
+
     if dataloader is None:
         if batch_size == 0:
             # We let the system to determine the max data batch for each forward.
diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py
index 6a998729168..02ea5dcc791 100644
--- a/tests/unit/torch/utils/test_dataset_utils.py
+++ b/tests/unit/torch/utils/test_dataset_utils.py
@@ -340,6 +340,21 @@ def test_create_forward_loop_passes_checkpoint_args_through():
     assert len(calls) == 2
 
 
+def test_create_forward_loop_rejects_negative_checkpoint_every():
+    """A negative checkpoint_every stays truthy under `% checkpoint_every`, so it must be
+    rejected explicitly rather than silently mis-triggering (or ZeroDivisionError-adjacent
+    surprises) deep inside the loop."""
+    with pytest.raises(ValueError, match="non-negative"):
+        create_forward_loop(dataloader=Mock(), checkpoint_every=-1, checkpoint_fn=lambda: None)
+
+
+def test_create_forward_loop_rejects_missing_checkpoint_fn():
+    """checkpoint_every > 0 with no callable checkpoint_fn would call None after the first
+    completed interval; reject it up front instead of failing mid-run."""
+    with pytest.raises(ValueError, match="checkpoint_fn"):
+        create_forward_loop(dataloader=Mock(), checkpoint_every=2, checkpoint_fn=None)
+
+
 def test_disable_use_cache_restores_on_exception():
     """Restore must run even if the with-block raises."""
     model = torch.nn.Linear(4, 4)

From 645261eba8e8cf05c09b1d40720d3bea0ba6d3fc Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:30:05 -0400
Subject: [PATCH 6/9] test(dataset_utils): give _tiny_loader's _Model an
 explicit __init__

Matches the other nn.Module test double in this file, which does call
super().__init__() explicitly rather than relying on the implicit default.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 tests/unit/torch/utils/test_dataset_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py
index 02ea5dcc791..d5d5b8a6d18 100644
--- a/tests/unit/torch/utils/test_dataset_utils.py
+++ b/tests/unit/torch/utils/test_dataset_utils.py
@@ -292,6 +292,9 @@ def _collate(samples):
 
 def _tiny_loader(num_batches: int):
     class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
         def forward(self, **kwargs):
             pass
 

From 8a08919526cf4bee84e18a8d7890ceab8a99c6e3 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:30:38 -0400
Subject: [PATCH 7/9] test(dataset_utils): assert checkpoint callback
 positions, not just count

len(calls) == 2 also passes if checkpoint_fn fires at the wrong batches
(e.g. after 1 and 3 instead of 2 and 4). Record the completed forward count
in the callback and assert calls == [2, 4] so the test actually pins down
where the callback fires.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 tests/unit/torch/utils/test_dataset_utils.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py
index d5d5b8a6d18..9ddc6d2d41c 100644
--- a/tests/unit/torch/utils/test_dataset_utils.py
+++ b/tests/unit/torch/utils/test_dataset_utils.py
@@ -294,9 +294,10 @@ def _tiny_loader(num_batches: int):
     class _Model(torch.nn.Module):
         def __init__(self):
             super().__init__()
+            self.forward_calls = 0
 
         def forward(self, **kwargs):
-            pass
+            self.forward_calls += 1
 
     def _collate(samples):
         return {"input_ids": torch.stack([s["input_ids"] for s in samples])}
@@ -308,14 +309,20 @@ def _collate(samples):
 
 def test_forward_loop_checkpoints_every_n_steps():
     """checkpoint_fn fires every `checkpoint_every` batches -- 5 batches at every=2 means
-    two fires (after batch 2 and batch 4), not one per batch and not a trailing fire for
-    the final partial group."""
+    two fires, after batch 2 and batch 4, not one per batch and not a trailing fire for
+    the final partial group. Recording the completed forward count (not just a fixed
+    marker) catches fires at the wrong positions, e.g. after batches 1 and 3."""
     calls: list[int] = []
     model, loader = _tiny_loader(5)
 
-    _forward_loop(model, loader, checkpoint_every=2, checkpoint_fn=lambda: calls.append(1))
+    _forward_loop(
+        model,
+        loader,
+        checkpoint_every=2,
+        checkpoint_fn=lambda: calls.append(model.forward_calls),
+    )
 
-    assert len(calls) == 2
+    assert calls == [2, 4]
 
 
 def test_forward_loop_checkpoint_every_zero_never_fires():

From fb472506755bc33d1bd538b022e72548bd1361e7 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:33:46 -0400
Subject: [PATCH 8/9] hf_ptq: move restore-quantized-state handling before
 calibration setup

--restore_quantized_state retries a failed or interrupted export from a
previously calibrated state, so it should not depend on calibration inputs
at all. Previously the restore check ran deep inside the mono-quantization
branch, after batch-size probing, calibration dataloader construction, and
the pre-quantize generation preview had already run -- defeating the
export-retry path when the original calibration dataset is unavailable.

Check for a requested restore at the top of quantize_main instead, and skip
straight to post_quantize/export when it happens. The generation-preview
args passed to post_quantize are None in this path, which disables the
before/after generation comparison (matching --skip_generate's existing
None-preview behavior) without needing calib_dataloader, which post_quantize
never actually used.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 examples/hf_ptq/hf_ptq.py                 | 35 ++++++++++++++---
 tests/examples/hf_ptq/test_hf_ptq_args.py | 48 +++++++++++++++++++++++
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index e169b7841dc..5bffd1639ba 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -1173,6 +1173,34 @@ def quantize_main(
     default_pad_token,
     device: torch.device,
 ):
+    # Detect if this is a Nemotron VL model using architecture-based detection. Cheap and
+    # needed on both the restore and calibration paths below.
+    is_nemotron_vl_model = is_nemotron_vl(full_model)
+
+    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. Passing None for the generation-preview args disables the
+        # before/after generation comparison inside post_quantize; export still runs.
+        post_quantize(
+            args,
+            full_model,
+            language_model,
+            model_type,
+            tokenizer,
+            processor,
+            None,
+            None,
+            None,
+            is_nemotron_vl_model,
+            None,
+            default_padding_side,
+            default_pad_token,
+            None,
+        )
+        return
+
     # Load the recipe up front so we can detect layerwise calibration before batch-size probing.
     recipe = None
     if args.recipe is not None:
@@ -1277,9 +1305,6 @@ def _is_layerwise(obj):
         ),
     )
 
-    # 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
     )
@@ -1349,9 +1374,7 @@ def _is_layerwise(obj):
             quant_cfg = copy.deepcopy(quant_cfg)
             force_weight_quantizers_static(quant_cfg["quant_cfg"])
 
-        if _restore_quantized_state_if_requested(args, full_model):
-            pass
-        elif quant_cfg:
+        if quant_cfg:
             mono_quantize(
                 args,
                 quant_cfg,
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index ad36c588bda..38cc7c166cf 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -115,6 +115,54 @@ def test_restore_quantized_state_with_deprecated_auto_quantize_cli_is_rejected(m
         )
 
 
+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)
+    monkeypatch.setattr(hf_ptq, "is_nemotron_vl", lambda model: False)
+
+    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, tokenizer, dps, dpt: 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
+        "left",  # default_padding_side
+        None,  # default_pad_token
+        "cuda",  # device
+    )
+
+    assert export_calls == [full_model]
+
+
 def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs."""
     hf_ptq, args = _parse_hf_ptq_args(

From 2f0dfa001c88c68e1a9767506b6bae3b21c97306 Mon Sep 17 00:00:00 2001
From: Wyatt Neal <wyatt.neal+git@gmail.com>
Date: Tue, 11 Aug 2026 08:35:20 -0400
Subject: [PATCH 9/9] test(hf_ptq): add real mto.save/mto.restore persistence
 round-trip test

The existing save/restore tests mock both mto.save and mto.restore, so they
verify delegation and rejection flow but never that ModelOpt state is
actually written, restored, or usable afterward. Add a focused test that
quantizes a small local Linear layer, saves its state through hf_ptq's
--save_quantized_state helper, restores it into a fresh model via
--restore_quantized_state, and confirms the restored quantizer state matches
and the model runs a forward pass.

Also switches the existing mocked save test off a hardcoded /tmp path onto
tmp_path.

Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
---
 tests/examples/hf_ptq/test_hf_ptq_args.py | 36 +++++++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 38cc7c166cf..51293def8d6 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -20,6 +20,7 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 import yaml
 
 from modelopt.recipe import load_recipe
@@ -55,11 +56,12 @@ 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):
+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", "/tmp/calib.pt"
+        monkeypatch, "--pyt_ckpt_path", "dummy", "--save_quantized_state", state_path
     )
     calls = []
     monkeypatch.setattr(hf_ptq.mto, "save", lambda model, path: calls.append((model, path)))
@@ -67,7 +69,7 @@ def test_save_quantized_state_calls_mto_save_with_model_and_path(monkeypatch):
     dummy_model = object()
     hf_ptq._save_quantized_state_if_requested(args, dummy_model)
 
-    assert calls == [(dummy_model, "/tmp/calib.pt")]
+    assert calls == [(dummy_model, state_path)]
 
 
 def test_save_and_restore_quantized_state_together_is_rejected(monkeypatch):
@@ -115,6 +117,34 @@ def test_restore_quantized_state_with_deprecated_auto_quantize_cli_is_rejected(m
         )
 
 
+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