From a4afd791768d5624697da9b9edae63b33c02d0e1 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 12 Oct 2026 18:42:59 +0000 Subject: [PATCH 1/2] fix(export): drop reverse renames that invent names the source checkpoint lacks A reverse rename maps an in-memory name back to its source-checkpoint name, so it must be a no-op on names that are already in source form -- the invariant _idempotent_pattern already enforces for prefix rules. A model that declares a backwards-compatibility *load* rename as a reversible WeightRenaming breaks it: reversing the rule rewrites the canonical name to a legacy spelling the published checkpoint never used, and a deployment loader keyed on the canonical name then cannot find the tensor. Read the source checkpoint's tensor names and drop any reverse rename that rewrites a name the checkpoint already carries. A rule is dropped only on positive evidence -- it must rewrite at least one source key, and every source key it rewrites must land on a name absent from the checkpoint -- so renames that legitimately restore source names, and rules that match nothing, are kept. When the source checkpoint cannot be read (a config-built model, a hub id never materialized locally, an unrecognized layout) the behavior is unchanged. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- .../torch/export/quant_aware_conversion.py | 78 +++++++++++++++++++ .../export/test_quant_aware_conversion.py | 53 +++++++++++++ 2 files changed, 131 insertions(+) diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 8c750702751..1ac845b79f0 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -62,10 +62,12 @@ (in-memory-name) behavior rather than emit a silently-wrong checkpoint. """ +import json import re from collections.abc import Iterable from dataclasses import dataclass from fnmatch import fnmatchcase +from pathlib import Path from typing import Any import torch @@ -504,6 +506,81 @@ def _scope_prefixes(rev) -> tuple[str, ...]: return tuple(seen) +_SAFETENSORS_INDEX = "model.safetensors.index.json" +_SAFETENSORS_SINGLE = "model.safetensors" + + +def _source_checkpoint_keys(model) -> set[str] | None: + """Tensor names of the checkpoint ``model`` was loaded from, or ``None``. + + ``None`` means "unknown" and leaves the reverse conversion exactly as it was: it is + returned for a model built from a config, a hub id that was never materialized + locally, or a checkpoint layout this does not recognize. + """ + path = getattr(model, "name_or_path", None) or getattr( + getattr(model, "config", None), "_name_or_path", None + ) + if not path: + return None + try: + directory = Path(path) + if not directory.is_dir(): + return None + index = directory / _SAFETENSORS_INDEX + if index.is_file(): + with index.open() as handle: + return set(json.load(handle)["weight_map"]) + single = directory / _SAFETENSORS_SINGLE + if single.is_file(): + from safetensors import safe_open + + with safe_open(single, framework="pt") as handle: + return set(handle.keys()) + except Exception: # unreadable or malformed checkpoint -> treat as unknown + return None + return None + + +def _drop_non_idempotent_renames( + rules: list[RenameRule], source_keys: set[str] | None +) -> list[RenameRule]: + """Drop reverse renames that rewrite names the source checkpoint already carries. + + A reverse rename maps an in-memory name back to its source-checkpoint name, so it + must be a *no-op* on names that are already in source form -- the same invariant + :func:`_idempotent_pattern` enforces for prefix rules. A rule that breaks it is not + restoring a name, it is inventing one. + + That happens when a model declares a backwards-compatibility *load* rename as a + reversible ``WeightRenaming``. Reversing it rewrites the canonical name to a legacy + spelling the published checkpoint never used, and a deployment loader keyed on the + canonical name then cannot find the tensor. + + A rule is dropped only on positive evidence: it must rewrite at least one source key, + and every source key it rewrites must land on a name absent from the checkpoint. A + rule that matches nothing, or whose results are themselves real source names, is kept + -- so renames that legitimately restore source names are never touched. + """ + if not source_keys: + return rules + kept: list[RenameRule] = [] + for rule in rules: + pattern = re.compile(_idempotent_pattern(rule.pattern, rule.repl)) + rewrites_a_source_key = False + every_result_absent = True + for key in source_keys: + mapped = _sub_scoped(pattern, rule.repl, key, rule.scope_prefixes) + if mapped == key: + continue + rewrites_a_source_key = True + if mapped in source_keys: + every_result_absent = False + break + if not (rewrites_a_source_key and every_result_absent): + kept.append(rule) + return kept + + def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: """Drop child-model reverse renames when the child namespace already exists. @@ -660,6 +737,7 @@ def _build_reverse_rules( # substrings and are applied first. weight_renamings = _drop_shadowed_prefix_renames(model, weight_renamings) rename_rules = leaf_renamings + list(reversed(weight_renamings)) + rename_rules = _drop_non_idempotent_renames(rename_rules, _source_checkpoint_keys(model)) return split_rules, merge_rules, rename_rules, expert_fused_leaves diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index e2efe7c7ff6..b2a8f350381 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -38,6 +38,8 @@ RenameRule, SplitRule, _assert_experts_pre_expanded, + _drop_non_idempotent_renames, + _source_checkpoint_keys, apply_reverse_rules, build_reverse_name_mapper, revert_quant_config_names, @@ -187,6 +189,57 @@ def test_idempotency_guard_leaves_disjoint_rules_alone(): assert list(out) == ["vision_model.radio_model.model.blocks.0.attn.weight"] +# A NemotronH-shaped source checkpoint: the canonical embedding name is plural, and the +# projector weights sit at the top level under their fused ``mlp1.`` names. +_SOURCE_KEYS = { + "backbone.embeddings.weight", + "backbone.layers.0.mixer.in_proj.weight", + "mlp1.0.weight", + "mlp1.1.weight", + "vision_projector.vision_final_layernorm.weight", +} + + +def test_rename_inventing_a_name_absent_from_source_is_dropped(): + """A BC *load* rename declared reversible would rewrite the canonical name. + + Reversing ``embedding.weight -> embeddings.weight`` turns the published + ``backbone.embeddings.weight`` into a legacy spelling the checkpoint never used, + which a deployment loader keyed on the canonical name cannot find. + """ + rule = RenameRule(r"embeddings\.weight", "embedding.weight") + assert _drop_non_idempotent_renames([rule], _SOURCE_KEYS) == [] + + +def test_rename_restoring_a_source_name_is_kept(): + """A rename whose output is a real source name is left alone.""" + rule = RenameRule(r"^vision_projector\.mlp1\.norm\.", "mlp1.0.") + assert _drop_non_idempotent_renames([rule], _SOURCE_KEYS) == [rule] + + +def test_rename_matching_no_source_key_is_kept(): + """No evidence means no drop: a rule that never fires on source names survives.""" + rule = RenameRule(r"^sound_projector\.sound_encoder\.", "sound_encoder.") + assert _drop_non_idempotent_renames([rule], _SOURCE_KEYS) == [rule] + + +@pytest.mark.parametrize("source_keys", [None, set()]) +def test_guard_is_inert_when_source_is_unknown(source_keys): + """An unreadable or absent source checkpoint keeps the previous behaviour.""" + rules = [RenameRule(r"embeddings\.weight", "embedding.weight")] + assert _drop_non_idempotent_renames(rules, source_keys) == rules + + +def test_source_checkpoint_keys_reads_sharded_index(tmp_path): + """The guard's input comes from the checkpoint the model was loaded from.""" + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert _source_checkpoint_keys(types.SimpleNamespace(name_or_path="", config=None)) is None + + def test_rename_carries_scale_siblings(): """A module rename rewrites weight + all scale siblings with identical values.""" sd = _nvfp4_linear("model.language_model.layers.10.mlp.experts.40.gate_proj", 8, 16) From 69fda764f5410e7e5a5df6385f9c316bb3dfb8b0 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 12 Oct 2026 19:23:58 +0000 Subject: [PATCH 2/2] docs: add a CHANGELOG entry for the reverse-rename source-name guard Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 78a71f248f6..6874bfa9640 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -131,6 +131,7 @@ Changelog **Bug Fixes** +- Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Reverse renames that rewrite a name the source checkpoint already carries are now skipped, for both the exported tensors and the quantization config's module references. - Fix ONNX quantization reusing existing tensor names for scale and zero-point parameters while preserving INT8 calibration-cache lookup. The boolean ``--target_dla`` / ``target_dla=True`` now preserves scalar multiplication/division operands and padding fill values as constants; use Q/DQ Translator for compatible TensorRT 10.7-and-earlier DLA deployments, or direct Q/DQ with strongly typed DLA enabled in a compatible TensorRT 11.4+ build. - Fix calibration silently dropping samples when an out-of-memory error forces the batch to be split and a smaller working batch size is already known: the slices stopped tiling the batch,