PTQ reliability fixes: checkpoint resume + exclude_modules sentinel corruption

examples/hf_ptq/hf_ptq.py

@@ -752,6 +752,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],
@@ -1145,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:
@@ -1249,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
     )
@@ -1332,6 +1385,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")
@@ -1560,6 +1614,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.
@@ -1682,6 +1757,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

modelopt/torch/export/plugins/hf_checkpoint_utils.py

@@ -104,16 +104,79 @@ 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, 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 "
+            "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.
 
     Fix conservative deployment-only config incompatibilities:
 
     * 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")

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

modelopt/torch/utils/dataset_utils.py

@@ -1158,13 +1158,21 @@ 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.
 
     Args:
         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. 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)
@@ -1173,11 +1181,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 +1201,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 +1224,13 @@ 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. 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:
 
@@ -1235,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.
@@ -1251,7 +1275,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):