From 900f8b32a538df9d4ca707ef15f9df1fe902b177 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:08:21 +0000 Subject: [PATCH 1/2] Let a reverse conversion record a merge rule no key matched revert_weight_conversion_quant_aware refuses a merge rule that matches no key of the state dict, so a later rename cannot leave a checkpoint that mixes the two namings. A rank that reverts its own share of a model can hold none of a rule's keys while another rank holds them all, so the conversion now takes an unmatched_merges list: given one, a rule's refusal is recorded there instead of raised, and the caller refuses a rule that no rank matched. Without it nothing changes. [1/3] of #2751's split, at its review's request. Signed-off-by: Shengliang Xu --- .../torch/export/quant_aware_conversion.py | 40 ++++++++++++++++--- .../export/test_quant_aware_conversion.py | 13 ++++++ 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index ad186df9300..5bff5fa7ff8 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -188,8 +188,14 @@ def _apply_split_rule(state_dict: dict[str, torch.Tensor], rule: SplitRule) -> N state_dict[target_key] = _split_leaf_tensor(leaf, tensor, n, idx, rule.dim) -def _apply_merge_rule(state_dict: dict[str, torch.Tensor], rule: _MergeRule) -> None: - """Merge all complete tensor groups matched by a reversed ``WeightConverter``.""" +def _apply_merge_rule( + state_dict: dict[str, torch.Tensor], rule: _MergeRule, unmatched: list[str] | None = None +) -> None: + """Merge all complete tensor groups matched by a reversed ``WeightConverter``. + + A rule matching no key is refused, unless ``unmatched`` is given: it then gets the refusal's + message instead (see :func:`revert_weight_conversion_quant_aware`). + """ converter = rule.converter source_patterns = tuple(_as_list(converter.source_patterns)) groups: dict[str, dict[str, str]] = {} @@ -216,11 +222,15 @@ def _apply_merge_rule(state_dict: dict[str, torch.Tensor], rule: _MergeRule) -> sources[source_pattern] = key if not groups: - raise QuantConversionUnsupportedError( + message = ( f"merge rule for {source_patterns} matched no state-dict key " f"(scope_prefix={getattr(converter, 'scope_prefix', None)!r}, " f"base_model_prefix={getattr(converter, 'base_model_prefix', None)!r})" ) + if unmatched is None: + raise QuantConversionUnsupportedError(message) + unmatched.append(message) + return for target_key, sources in groups.items(): missing = [pattern for pattern in source_patterns if pattern not in sources] @@ -330,6 +340,8 @@ def apply_reverse_rules( rename_rules: list[RenameRule], merge_rules: list[_MergeRule] | None = None, source_keys: set[str] | None = None, + *, + unmatched_merges: list[str] | None = None, ) -> dict[str, torch.Tensor]: """Apply quant-aware reverse conversion: tensor transforms first, then renames. @@ -342,12 +354,14 @@ def apply_reverse_rules( ``source_keys`` are the tensor names of the checkpoint the model was loaded from. They let a rename that would otherwise invent a name resolve to the spelling the checkpoint actually carries; omitting them applies the rules unchanged. + + ``unmatched_merges`` is :func:`revert_weight_conversion_quant_aware`'s. """ out = dict(state_dict) for rule in split_rules: _apply_split_rule(out, rule) for merge_rule in merge_rules or []: - _apply_merge_rule(out, merge_rule) + _apply_merge_rule(out, merge_rule, unmatched_merges) rename = _SourceAwareRenamer(rename_rules, source_keys) renamed: dict[str, torch.Tensor] = {} @@ -359,13 +373,20 @@ def apply_reverse_rules( return renamed -def revert_weight_conversion_quant_aware(model, state_dict: dict[str, torch.Tensor]): +def revert_weight_conversion_quant_aware( + model, state_dict: dict[str, torch.Tensor], *, unmatched_merges: list[str] | None = None +): """Reverse a transformers conversion_mapping on a quantized state dict. Builds reverse rules from the model's conversion mapping and applies them carrying companion scale tensors. Raises :class:`QuantConversionUnsupportedError` when the mapping uses an op that cannot be reversed quant-aware yet, so the caller can fall back to the legacy behavior. + + ``unmatched_merges`` is for a ``state_dict`` that holds one rank's share of the model, + with each merge group whole on one rank: a merge rule matching none of its keys appends + the refusal's message there instead of raising, since other ranks may hold the rule's + groups. The caller refuses a rule that no rank matched. """ split_rules, merge_rules, rename_rules, expert_fused_leaves, source_keys = _build_reverse_rules( model @@ -373,7 +394,14 @@ def revert_weight_conversion_quant_aware(model, state_dict: dict[str, torch.Tens if not split_rules and not merge_rules and not rename_rules: return state_dict _assert_experts_pre_expanded(state_dict, expert_fused_leaves) - return apply_reverse_rules(state_dict, split_rules, rename_rules, merge_rules, source_keys) + return apply_reverse_rules( + state_dict, + split_rules, + rename_rules, + merge_rules, + source_keys, + unmatched_merges=unmatched_merges, + ) def build_reverse_name_mapper(model, *, tensor_keys: bool = False): diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 1a798b2c1fd..1954ee60486 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -987,6 +987,19 @@ def test_radio_merge_that_matches_no_keys_raises(): revert_weight_conversion_quant_aware(model, state_dict) +def test_radio_merge_that_matches_no_keys_is_recorded_for_one_ranks_share(): + """Given ``unmatched_merges``, as for one rank's share of a model, whose other ranks may hold + the rule's keys, the refusal is recorded there instead of raised, and the rest still applies.""" + model = types.SimpleNamespace(_weight_conversions=[_radio_qkv_conversion()]) + state_dict = {"language_model.layers.0.weight": torch.randn(2, 2)} + unmatched: list = [] + + out = revert_weight_conversion_quant_aware(model, state_dict, unmatched_merges=unmatched) + + assert len(unmatched) == 1 and "matched no state-dict key" in unmatched[0] + assert out.keys() == state_dict.keys() + + @pytest.mark.parametrize( ("leaf", "tensors", "message"), [ From b9fbb24230043d07bb504d1142b813a72da239b9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:42:04 +0000 Subject: [PATCH 2/2] Say what unmatched_merges does where apply_reverse_rules takes it, and test the rules after one The parameter's description pointed at another function. The recorded-refusal test now has a second, matching merge rule after the unmatched one, so it checks that the reverse goes on to apply it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/quant_aware_conversion.py | 3 ++- .../torch/export/test_quant_aware_conversion.py | 17 ++++++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 5bff5fa7ff8..1005a3a904f 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -355,7 +355,8 @@ def apply_reverse_rules( They let a rename that would otherwise invent a name resolve to the spelling the checkpoint actually carries; omitting them applies the rules unchanged. - ``unmatched_merges`` is :func:`revert_weight_conversion_quant_aware`'s. + ``unmatched_merges``, when given, collects the refusal message of each merge rule that + matched no key instead of raising; see :func:`revert_weight_conversion_quant_aware`. """ out = dict(state_dict) for rule in split_rules: diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 1954ee60486..0e9e9104529 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -990,14 +990,25 @@ def test_radio_merge_that_matches_no_keys_raises(): def test_radio_merge_that_matches_no_keys_is_recorded_for_one_ranks_share(): """Given ``unmatched_merges``, as for one rank's share of a model, whose other ranks may hold the rule's keys, the refusal is recorded there instead of raised, and the rest still applies.""" - model = types.SimpleNamespace(_weight_conversions=[_radio_qkv_conversion()]) - state_dict = {"language_model.layers.0.weight": torch.randn(2, 2)} + unmatched_rule = _radio_qkv_conversion() + _set_scope_attr(unmatched_rule, "scope_prefix", "unused_model") + model = types.SimpleNamespace(_weight_conversions=[unmatched_rule, _radio_qkv_conversion()]) + parts = [torch.full((2, 2), float(i)) for i in range(3)] + state_dict = { + "language_model.layers.0.weight": torch.randn(2, 2), + **{ + f"vision_model.encoder.layer.0.attention.{part}_proj.weight": tensor + for part, tensor in zip("qkv", parts) + }, + } unmatched: list = [] out = revert_weight_conversion_quant_aware(model, state_dict, unmatched_merges=unmatched) assert len(unmatched) == 1 and "matched no state-dict key" in unmatched[0] - assert out.keys() == state_dict.keys() + fused = "vision_model.encoder.layer.0.attn.qkv.weight" + assert out.keys() == {"language_model.layers.0.weight", fused} + assert torch.equal(out[fused], torch.cat(parts)) @pytest.mark.parametrize(