fix(export): drop reverse renames that invent names the source checkpoint lacks

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,

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
 
 

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.<i>`` 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)