Improve nested base_model_prefix handling in weight conversion and loading

src/transformers/core_model_loading.py

@@ -1030,16 +1030,123 @@ class SkipParameters(Exception):
     pass
 
 
+def _compute_all_prefixes(model) -> list[str]:
+    """
+    Return all base-model prefix paths reachable from `model`, ordered shortest-first (BFS).
+
+    `base_model_prefix` on a class means "when I am stored as a submodule in a parent
+    model, the parent stores me under the attribute named `base_model_prefix`". A child is
+    therefore a "base model" of the current model when its `base_model_prefix` matches the
+    attribute name it is stored under.
+
+    Multiple base-model children are supported (e.g. a multi-modal model that contains
+    both `self.vision_model` and `self.text_model`).
+
+    Examples:
+
+        DetrForObjectDetection  -> ["model"]
+        DetrForSegmentation     -> ["detr", "detr.model"]
+        LlamaForCausalLM        -> ["model"]
+        CLIPModel               -> ["vision_model", "text_model"]
+    """
+    prefixes: list[str] = [getattr(model, "base_model_prefix", "")]
+    queue: list[tuple] = [(model, getattr(model, "base_model_prefix", ""))]
+
+    while queue:
+        current_model, accumulated_prefix = queue.pop(0)
+        for name, child in current_model.named_children():
+            child_prefix = getattr(child, "base_model_prefix", "")
+            if child_prefix and child_prefix == name:
+                next_accumulated = f"{accumulated_prefix}.{name}" if accumulated_prefix else name
+                prefixes.append(next_accumulated)
+                queue.append((child, next_accumulated))
+
+    return prefixes
+
+
+def _strip_model_prefix_for_save(key: str, model) -> str:
+    """
+    Recursively strip all `base_model_prefix` segments from a state-dict key so that
+    reverse conversion rules (written relative to the innermost base model) operate on
+    bare keys regardless of nesting depth.
+
+    We identify each prefix level by finding the direct child whose `base_model_prefix`
+    matches its attribute name (same logic as `_compute_all_prefixes`).
+
+    Examples for `DetrForSegmentation` (prefix chain `detr` -> `model`):
+
+        "detr.model.backbone.x"            ->  "backbone.x"
+        "detr.class_labels_classifier.x"   ->  "class_labels_classifier.x"
+        "mask_head.x"                      ->  "mask_head.x"
+    """
+    for name, child in model.named_children():
+        child_prefix = getattr(child, "base_model_prefix", "")
+        if child_prefix and child_prefix == name and key.startswith(name + "."):
+            stripped_key = key[len(name) + 1 :]
+            return _strip_model_prefix_for_save(stripped_key, child)
+    return key
+
+
+def _resolve_key_for_prefix_nesting(
+    renamed_key: str,
+    valid_prefixes: list[str],
+    meta_state_dict: dict,
+) -> str:
+    """
+    Rewrite `renamed_key` with `valid_prefixes` from `_compute_all_prefixes` (longest prefixes first) so
+    `base_model_prefix` lines up for head and base models (strip wrapper prefixes or add missing inner ones).
+
+    - Per prefix (longest first): strip leading `prefix.`; if `prefix` is dotted, also try prepending the substring
+      after its first `.`.
+    - If still unmatched: `valid_prefixes` only reflects the load target, so keys from a more wrapped checkpoint can
+      still embed `prefix.` in the middle of the path. For each prefix, restart from `renamed_key` and
+      repeatedly replace the string with everything after the first `prefix.` (discarding that segment and anything
+      before it), while the string starts with `prefix.` or contains `.{prefix}.`, until a suffix exists in
+      `meta_state_dict`.
+
+    Args:
+        renamed_key: Key after weight renamings and conversion patterns.
+        valid_prefixes: Candidate `base_model_prefix` paths for the model being loaded.
+        meta_state_dict: Reference key set (e.g. `model.state_dict()`).
+
+    Returns:
+        A matching key in `meta_state_dict`, or `renamed_key`.
+    """
+    for prefix in reversed(valid_prefixes):
+        if renamed_key.startswith(prefix + "."):
+            candidate = renamed_key[len(prefix) + 1 :]
+            if candidate in meta_state_dict:
+                return candidate
+        if "." in prefix:
+            # remove the first prefix (current model's prefix) when adding it to the key
+            add_prefix = prefix.split(".", maxsplit=1)[1]
+            candidate = f"{add_prefix}.{renamed_key}"
+            if candidate in meta_state_dict:
+                return candidate
+    # Checkpoint may wrap the target at 2+ nesting levels (outer prefixes not in valid_prefixes),
+    # so we need to check for the prefix inside the key.
+    for prefix in reversed(valid_prefixes):
+        candidate = renamed_key
+        # avoid matching parts of module names containing the prefix
+        while f".{prefix}." in candidate or candidate.startswith(f"{prefix}."):
+            candidate = candidate.split(prefix + ".", maxsplit=1)[1]
+            if candidate in meta_state_dict:
+                return candidate
+
+    return renamed_key
+
+
 def rename_source_key(
     source_key: str,
     weight_renamings: list[WeightRenaming],
     weight_converters: list[WeightConverter],
-    prefix: str | None = None,
+    valid_prefixes: list[str] | None = None,
     meta_state_dict: dict | None = None,
 ) -> tuple[str, str | None]:
     """
-    Rename a source key given all the renaming and weight conversion patterns we have. Also takes care of adding/removing
-    the base model prefix during loading if necessary.
+    Apply all renaming and conversion patterns to `source_key`, then reconcile the
+    result against the model state dict (step 3) by trying to add or strip each prefix
+    level from `valid_prefixes` until the key is found.
     """
     renamed_key = source_key
     # 1. apply all renamings in turns (if multiple match, it's the responsibility of the mappings to make sure they
@@ -1055,15 +1162,10 @@ def rename_source_key(
         if source_pattern is not None:
             break
 
-    # 3. check if we need to add or remove prefix if necessary (only during loading, not saving)
-    if prefix is not None and meta_state_dict is not None:
-        if (
-            renamed_key.startswith(prefix)
-            and meta_state_dict.get(re.sub(f"^{prefix}.", "", renamed_key, count=1)) is not None
-        ):
-            renamed_key = re.sub(f"^{prefix}.", "", renamed_key, count=1)
-        elif meta_state_dict.get(f"{prefix}.{renamed_key}") is not None:
-            renamed_key = f"{prefix}.{renamed_key}"
+    # 3. If the key is still not in the model state dict, try adding or removing each
+    # prefix level (longest first) until a match is found.  Only active during loading.
+    if valid_prefixes is not None and meta_state_dict is not None and renamed_key not in meta_state_dict:
+        renamed_key = _resolve_key_for_prefix_nesting(renamed_key, valid_prefixes, meta_state_dict)
 
     return renamed_key, source_pattern
 
@@ -1161,7 +1263,9 @@ def convert_and_load_state_dict_in_model(
     ```
 
     """
-    prefix = model.base_model_prefix
+    # All valid base_model_prefix paths for this model (e.g. ["rf_detr", "rf_detr.model"]
+    # for RfDetrForInstanceSegmentation); passed to rename_source_key to resolve keys.
+    valid_prefixes = _compute_all_prefixes(model)
     tp_plan = tp_plan or {}
     device_map = load_config.device_map or {"": "cpu"}
     hf_quantizer = load_config.hf_quantizer
@@ -1214,11 +1318,13 @@ def convert_and_load_state_dict_in_model(
     for original_key, tensor in state_dict:
         # 1. Rename the key according to all renaming pattern and optional weight converter patterns
         renamed_key, source_pattern = rename_source_key(
-            original_key, renamings, converters, prefix, meta_model_state_dict
+            original_key, renamings, converters, valid_prefixes, meta_model_state_dict
         )
         if renamed_key not in meta_model_state_dict and original_key in meta_model_state_dict:
-            # Key should probably not have been renamed but we might need the `prefix` to be added.`
-            renamed_key, source_pattern = rename_source_key(original_key, [], [], prefix, meta_model_state_dict)
+            # Key should probably not have been renamed but we might need the prefix(es) to be added.
+            renamed_key, source_pattern = rename_source_key(
+                original_key, [], [], valid_prefixes, meta_model_state_dict
+            )
 
         # 2. finally, collect the tensor into the proper converter
         if renamed_key in meta_model_state_dict:
@@ -1374,17 +1480,21 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
     pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns}
     conversion_mapping = {}
 
+    # Opt in via `_checkpoint_conversion_prefix_free = True` when the source checkpoint is fully flat,
+    # so that all prefixes should be stripped before saving.
+    strip_prefix = getattr(model, "_checkpoint_conversion_prefix_free", False)
+
     state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
     for original_key, tensor in state_dict:
-        # Rename the key according to all renaming pattern and optional weight converter patterns
-        renamed_key, source_pattern = rename_source_key(original_key, renamings, converters)
+        bare_key = _strip_model_prefix_for_save(original_key, model) if strip_prefix else original_key
+        renamed_key, source_pattern = rename_source_key(bare_key, renamings, converters)
         if source_pattern is not None:
             new_converter = deepcopy(pattern_to_converter[source_pattern])
             # each target key gets its own converter instance
             mapping = conversion_mapping.setdefault(renamed_key, new_converter)
         else:
-            mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
-            source_pattern = original_key
+            mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(bare_key, renamed_key))
+            source_pattern = bare_key
 
         mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
 

src/transformers/models/conditional_detr/modeling_conditional_detr.py

@@ -1477,6 +1477,8 @@ def inverse_sigmoid(x, eps=1e-5):
     """
 )
 class ConditionalDetrForObjectDetection(ConditionalDetrPreTrainedModel):
+    base_model_prefix = "conditional_detr"
+
     def __init__(self, config: ConditionalDetrConfig):
         super().__init__(config)
 

src/transformers/models/detr/modeling_detr.py

@@ -1298,6 +1298,8 @@ def forward(self, x):
     """
 )
 class DetrForObjectDetection(DetrPreTrainedModel):
+    base_model_prefix = "detr"
+
     def __init__(self, config: DetrConfig):
         super().__init__(config)
 

tests/models/conditional_detr/test_modeling_conditional_detr.py

@@ -16,25 +16,19 @@
 import copy
 import inspect
 import math
-import os
-import re
-import tempfile
 import unittest
 from functools import cached_property
 
 from transformers import ConditionalDetrConfig, ResNetConfig, is_torch_available, is_vision_available
-from transformers.conversion_mapping import get_model_conversion_mapping
-from transformers.core_model_loading import WeightRenaming, process_target_pattern
 from transformers.testing_utils import require_timm, require_torch, require_vision, slow, torch_device
 
 from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin, compare_state_dicts, floats_tensor
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
 from ...test_pipeline_mixin import PipelineTesterMixin
 
 
 if is_torch_available():
     import torch
-    from safetensors.torch import load_file
 
     from transformers import (
         ConditionalDetrForObjectDetection,
@@ -240,88 +234,6 @@ def test_conditional_detr_object_detection_head_model(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_conditional_detr_object_detection_head_model(*config_and_inputs)
 
-    def test_reverse_loading_mapping(self, check_keys_were_modified=True):
-        # Some conversions from the mapping are specific to `DetrForSegmentation` model only
-        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
-
-        #  Some MoE models alternate between a classic MLP and a MoE layer, in which case we want to have at
-        # lest one MoE layer here to check the mapping
-        config_to_set = config.get_text_config(decoder=True)
-        config_to_set.first_k_dense_replace = 1  # means that the first layer (idx 0) will be MLP, then MoE
-        config_to_set.moe_layer_start_index = 1  # same as above but for Ernie 4.5...
-        config_to_set.mlp_only_layers = [0]  # same but for qwens
-        config_to_set.num_dense_layers = 1  # lfm2_moe
-
-        for model_class in self.all_model_classes:
-            # Each individual model is a subtest
-            with self.subTest(model_class.__name__):
-                model = model_class(copy.deepcopy(config))
-                # Skip if no conversions
-                conversions = get_model_conversion_mapping(model, add_legacy=False)
-                if len(conversions) == 0:
-                    # No conversion mapping for this model only, needs to test other classes
-                    continue
-
-                # Find the model keys, so the targets according to the conversions
-                model_keys = list(model.state_dict().keys())
-
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    # Serialize with reverse mapping
-                    model.save_pretrained(tmpdirname)
-                    state_dict = load_file(os.path.join(tmpdirname, "model.safetensors"))
-                    # Get all the serialized keys that we just saved according to the reverse mapping
-                    serialized_keys = list(state_dict.keys())
-
-                if check_keys_were_modified:
-                    # They should be different, otherwise we did not perform any mapping
-                    self.assertNotEqual(sorted(serialized_keys), sorted(model_keys), "No key mapping was performed!")
-
-                # Check that for each conversion entry, we at least map to one key
-                for conversion in conversions:
-                    for source_pattern in conversion.source_patterns:
-                        # Sometimes the mappings specify keys that are tied, so absent from the saved state dict
-                        if isinstance(conversion, WeightRenaming):
-                            # We need to revert the target pattern to make it compatible with regex search
-                            target_pattern_reversed = conversion.target_patterns[0]
-                            captured_group = process_target_pattern(source_pattern)[1]
-                            if captured_group:
-                                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)
-
-                        # DIFF FROM MIXIN IS HERE
-                        if (
-                            "bbox" in source_pattern or "mask_head" in source_pattern
-                        ) and model_class != ConditionalDetrForSegmentation:
-                            pass
-                        else:
-                            self.assertTrue(
-                                num_matches > 0,
-                                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",
-                            )
-
-                # If everything is still good at this point, let's test that we perform the same operations both when
-                # reverting ops from `from_pretrained` and from `__init__`
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    # The model was instantiated from __init__ before being saved
-                    model.save_pretrained(tmpdirname)
-                    state_dict_saved_from_init = load_file(os.path.join(tmpdirname, "model.safetensors"))
-
-                    # Now reload it
-                    model_reloaded = model_class.from_pretrained(tmpdirname)
-
-                    # Make sure both loaded state_dict are identical
-                    self.assertTrue(compare_state_dicts(model_reloaded.state_dict(), model.state_dict()))
-
-                    # The model was instantiated from `from_pretrained` before being saved
-                    model_reloaded.save_pretrained(tmpdirname)
-                    state_dict_saved_from_pretrained = load_file(os.path.join(tmpdirname, "model.safetensors"))
-
-                    # Make sure both saved state_dict are identical
-                    self.assertTrue(compare_state_dicts(state_dict_saved_from_init, state_dict_saved_from_pretrained))
-
     # TODO: check if this works again for PyTorch 2.x.y
     @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0.")
     def test_multi_gpu_data_parallel_forward(self):