[Weight Converter] More fine-grained mappings on classes, scoping for every transforms (including weight converter)

src/transformers/integrations/accelerate.py

@@ -446,15 +446,13 @@ def accelerate_disk_offload(
     renamed) will be mapped to where they already reside on disk. Otherwise, the parameters will be resaved inside
     `disk_offload_folder` during loading.
     """
-    from ..core_model_loading import WeightRenaming, rename_source_key
+    from ..core_model_loading import rename_source_key
 
     if disk_offload_folder is not None:
         os.makedirs(disk_offload_folder, exist_ok=True)
     is_offloaded_safetensors = checkpoint_files is not None and checkpoint_files[0].endswith(".safetensors")
 
-    renamings = []
-    if weight_mapping is not None:
-        renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)]
+    transforms = weight_mapping if weight_mapping is not None else []
 
     # In this case, the offload index is simply the existing safetensors (except if using custom weight loading
     # Operation, e.g. the MoE models, where we need to resave the weights that were changed at loading time)
@@ -470,7 +468,8 @@ def accelerate_disk_offload(
 
         # Update the weight names according to the `weight_mapping`
         weight_renaming_map = {
-            rename_source_key(k, renamings, [], model.base_model_prefix, meta_state_dict)[0]: k for k in weight_map
+            rename_source_key(k, transforms, prefix=model.base_model_prefix, meta_state_dict=meta_state_dict)[0]: k
+            for k in weight_map
         }
 
         # Prepare the index using existing safetensors files

src/transformers/integrations/deepspeed.py

@@ -347,7 +347,7 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
                 "in your DeepSpeed config or convert your checkpoint to the expected format first."
             )
 
-    from ..core_model_loading import WeightConverter, WeightRenaming, dot_natural_key, rename_source_key
+    from ..core_model_loading import WeightConverter, dot_natural_key, rename_source_key
 
     # Preserve metadata from the original state dict
     metadata = getattr(state_dict, "_metadata", None)
@@ -360,14 +360,15 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     for key, param in model.state_dict().items():
         model_state_dict[key] = torch.empty(param.shape, dtype=param.dtype, device="meta")
 
-    renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)]
     converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)]
 
     # Fast path: if we only have simple renamings and no converters, we can skip the expensive collection logic
     if len(converters) == 0:
         new_state_dict = {}
         for original_key, tensor in state_dict.items():
-            renamed_key, _ = rename_source_key(original_key, renamings, [], prefix, model_state_dict)
+            renamed_key, _ = rename_source_key(
+                original_key, weight_mapping, prefix=prefix, meta_state_dict=model_state_dict
+            )
             if renamed_key in model_state_dict:
                 new_state_dict[renamed_key] = tensor
         # Attach metadata to the new state dict
@@ -386,7 +387,9 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     sorted_keys = sorted(state_dict.keys(), key=lambda k: dot_natural_key(k))
     for original_key in sorted_keys:
         tensor = state_dict.pop(original_key)
-        renamed_key, source_pattern = rename_source_key(original_key, renamings, converters, prefix, model_state_dict)
+        renamed_key, source_pattern = rename_source_key(
+            original_key, weight_mapping, prefix=prefix, meta_state_dict=model_state_dict
+        )
 
         # Only process if the renamed key is in the model's state dict
         if renamed_key in model_state_dict:

src/transformers/modeling_utils.py

@@ -1394,6 +1394,12 @@ def post_init(self):
         # Maybe initialize the weights and tie the keys
         self.init_weights()
         self._backward_compatibility_gradient_checkpointing()
+        # Cache the list of (name, submodule) pairs where the submodule is a PreTrainedModel.
+        # This pattern is used in several places across the codebase; computing it once avoids
+        # repeated traversal of the full module tree.
+        self._named_pretrained_submodules: list[tuple[str, PreTrainedModel]] = [
+            (name, module) for name, module in self.named_modules() if isinstance(module, PreTrainedModel)
+        ]
 
     @property
     def tp_plan(self) -> dict[str, str]:

tests/models/aya_vision/test_modeling_aya_vision.py

@@ -206,9 +206,6 @@ def test_sdpa_can_compile_dynamic(self):
     def test_batching_equivalence(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class AyaVisionIntegrationTest(unittest.TestCase):

tests/models/colpali/test_modeling_colpali.py

@@ -226,10 +226,6 @@ def test_sdpa_can_dispatch_on_flash(self):
     def test_sdpa_can_compile_dynamic(self):
         pass
 
-    @unittest.skip(reason="Some weight mappings from paligemma are unreachable here as they use a `^` pattern")
-    def test_reverse_loading_mapping(self):
-        pass
-
 
 @require_torch
 class ColPaliModelIntegrationTest(unittest.TestCase):

tests/models/colqwen2/test_modeling_colqwen2.py

@@ -284,10 +284,6 @@ def test_sdpa_can_compile_dynamic(self):
     def test_load_save_without_tied_weights(self):
         pass
 
-    @unittest.skip(reason="One weight renaming from qwen2 is unreachable here as it uses a `^` pattern")
-    def test_reverse_loading_mapping(self):
-        pass
-
 
 @require_torch
 class ColQwen2ModelIntegrationTest(unittest.TestCase):

tests/models/emu3/test_modeling_emu3.py

@@ -340,9 +340,6 @@ def _image_features_get_expected_num_hidden_states(self, model_tester=None):
         up_down_blocks = len(model_tester.vq_channel_multiplier) * model_tester.vq_num_res_blocks
         return up_down_blocks + 2 + model_tester.vq_num_res_blocks + 1
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class Emu3IntegrationTest(unittest.TestCase):

tests/models/fuyu/test_modeling_fuyu.py

@@ -264,9 +264,6 @@ def test_get_image_features_hidden_states(self):
     def test_get_image_features_attentions(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @slow
 @require_torch_accelerator

tests/models/gemma3/test_modeling_gemma3.py

@@ -427,9 +427,6 @@ def test_flash_attn_3_from_config(self):
     def test_flash_attn_4_from_config(self):
         self.flash_attn_from_config(attn_implementation="flash_attention_4", test_fwd_in_train=False)
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @slow
 @require_torch_accelerator

tests/models/got_ocr2/test_modeling_got_ocr2.py

@@ -161,9 +161,6 @@ def setUp(self):
     def test_config(self):
         self.config_tester.run_common_tests()
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class GotOcr2IntegrationTest(unittest.TestCase):

tests/models/internvl/test_modeling_internvl.py

@@ -214,9 +214,6 @@ def test_sdpa_can_compile_dynamic(self):
     def test_flash_attn_2_fp32_ln(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @slow
 @require_torch_accelerator

tests/models/llava/test_modeling_llava.py

@@ -282,9 +282,6 @@ def test_training_gradient_checkpointing_use_reentrant_true(self):
     def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 @slow

tests/models/llava_next/test_modeling_llava_next.py

@@ -125,9 +125,6 @@ def test_training_gradient_checkpointing_use_reentrant_true(self):
     def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class LlavaNextForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/models/llava_next_video/test_modeling_llava_next_video.py

@@ -340,9 +340,6 @@ def _video_features_prepare_config_and_inputs(self):
         inputs_dict = {"pixel_values": inputs_dict["pixel_values_videos"]}
         return config, inputs_dict
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class LlavaNextVideoForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/models/llava_onevision/test_modeling_llava_onevision.py

@@ -312,9 +312,6 @@ def _video_features_prepare_config_and_inputs(self):
         inputs_dict = {"pixel_values": pixel_values_videos}
         return config, inputs_dict
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class LlavaOnevisionForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/models/mistral3/test_modeling_mistral3.py

@@ -230,9 +230,6 @@ def test_sdpa_can_dispatch_on_flash(self):
     def test_flex_attention_with_grads(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @slow
 @require_torch_accelerator

tests/models/mllama/test_modeling_mllama.py

@@ -452,9 +452,6 @@ def test_left_padding_compatibility(self):
             unpadded_custom_inputs=unpadded_custom_inputs, padded_custom_inputs=padded_custom_inputs
         )
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class MllamaForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/models/paligemma/test_modeling_paligemma.py

@@ -327,9 +327,6 @@ def test_attention_mask_with_token_types(self):
                                     f"Found non-zero attention weights for padding token at batch {batch_idx}, sequence position {seq_idx}",
                                 )
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @slow
 @require_torch

tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py

@@ -441,9 +441,6 @@ def attention_mask_padding_matches_padding_free_with_position_ids(
                 tol = torch.finfo(torch.bfloat16).eps
                 torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol)
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
     @unittest.skip(reason="Feedforward chunking is not yet supported")
     def test_feed_forward_chunking(self):
         pass

tests/models/qwen2_vl/test_modeling_qwen2_vl.py

@@ -411,9 +411,6 @@ def attention_mask_padding_matches_padding_free_with_position_ids(
                 tol = torch.finfo(torch.bfloat16).eps
                 torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol)
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
     @unittest.skip(reason="Feedforward chunking is not yet supported")
     def test_feed_forward_chunking(self):
         pass

tests/models/video_llava/test_modeling_video_llava.py

@@ -397,9 +397,6 @@ def test_vision_feature_layers(self, vision_feature_layer):
             assert base_model.multi_modal_projector.linear_1.in_features == expected_features
             model(**input_dict)
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class VideoLlavaForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/models/vipllava/test_modeling_vipllava.py

@@ -274,9 +274,6 @@ def test_training_gradient_checkpointing_use_reentrant_true(self):
     def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
         pass
 
-    def test_reverse_loading_mapping(self):
-        super().test_reverse_loading_mapping(skip_base_model=True)
-
 
 @require_torch
 class VipLlavaForConditionalGenerationIntegrationTest(unittest.TestCase):

tests/test_modeling_common.py

@@ -4801,7 +4801,20 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_
                     # mess up the prefixes only if the loaded checkpoints were doing so as well)
                     if isinstance(conversion, PrefixChange):
                         continue
-                    for source_pattern in conversion.source_patterns:
+
+                    # Skip if the conversion is scoped to a sub-model, as the conversion is already tested on the sub-model.
+                    # Also, it might be the case that the original/weights model did not include the sub-model (e.g. LlavaForConditionalGeneration)
+                    if conversion.scope_prefix is not None:
+                        continue
+
+                    # Single pass over serialized_keys: the compiled regex already tests all
+                    # pattern branches at once, so one call per key is enough.
+                    matched_groups: set[str] = set()
+                    for key in serialized_keys:
+                        if (match := conversion._scoped_match(key)) is not None:
+                            matched_groups.add(match[2].lastgroup)  # "g0", "g1", ...
+
+                    for pattern_index, source_pattern in enumerate(conversion.source_patterns):
                         # Some patterns are written for gen-model only and won't be applied on base model
                         if "lm_head" in source_pattern and model_class not in [
                             *get_values(MODEL_FOR_CAUSAL_LM_MAPPING_NAMES),
@@ -4818,9 +4831,9 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_
                                 target_pattern_reversed = target_pattern_reversed.replace(r"\1", captured_group)
                             if any(re.search(target_pattern_reversed, k) for k in model.all_tied_weights_keys.keys()):
                                 continue
-                        num_matches = sum(re.search(source_pattern, key) is not None for key in serialized_keys)
+
                         self.assertTrue(
-                            num_matches > 0,
+                            f"g{pattern_index}" in matched_groups,
                             f"`{source_pattern}` in `{conversion}` did not match any of the source keys. "
                             "This indicates whether that the pattern is not properly written, or that it could not be reversed correctly",
                         )