diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 0790f644308..5bffd1639ba 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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 diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index ddae8e2b409..7d1e9fe9a00 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -104,6 +104,65 @@ 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. @@ -111,9 +170,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/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..bdea6805115 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,12 @@ 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. 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): diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 419ec4c3025..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,6 +56,143 @@ 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_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_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) + 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( diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index f3be2564312..6925d846a47 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -261,6 +261,96 @@ 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_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 = { 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..9ddc6d2d41c 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,81 @@ def _collate(samples): assert model.config.use_cache is True +def _tiny_loader(num_batches: int): + class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.forward_calls = 0 + + def forward(self, **kwargs): + self.forward_calls += 1 + + 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. 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(model.forward_calls), + ) + + assert calls == [2, 4] + + +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_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)