From d3998544a193e174ca9c9b3f697045bc446b6945 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 13 Apr 2026 23:32:21 +0000 Subject: [PATCH 1/5] support nested prefixes --- src/transformers/core_model_loading.py | 114 +++++++++++++++--- .../test_modeling_conditional_detr.py | 90 +------------- tests/models/detr/test_modeling_detr.py | 89 -------------- tests/test_modeling_common.py | 18 +++ 4 files changed, 113 insertions(+), 198 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index b43d5354e8ac..e8bea246ebe1 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1030,16 +1030,78 @@ class SkipParameters(Exception): pass +def _compute_all_prefixes(model) -> list[str]: + """ + Return all cumulative `base_model_prefix` paths for the model's nesting hierarchy, + ordered from shortest to longest. + + Examples: + + RfDetrModel / RfDetrForObjectDetection -> ["model"] + RfDetrForInstanceSegmentation -> ["model", "model.model"] + ConditionalDetrForPanopticSegmentation -> ["conditional_detr", "conditional_detr.model"] + LlamaForCausalLM -> ["model"] + """ + prefixes: list[str] = [] + current_model = model + accumulated_prefix = "" + + while True: + prefix = getattr(current_model, "base_model_prefix", "") + if not prefix: + break + + next_accumulated = f"{accumulated_prefix}.{prefix}" if accumulated_prefix else prefix + prefixes.append(next_accumulated) + + inner_model = getattr(current_model, prefix, None) + if inner_model is None: + break # current_model is the leaf base model + + # Stop when the inner model is itself a leaf (no deeper nesting to traverse). + inner_prefix = getattr(inner_model, "base_model_prefix", "") + if not inner_prefix or getattr(inner_model, inner_prefix, None) is None: + break + + accumulated_prefix = next_accumulated + current_model = inner_model + + 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. + + Examples for `RfDetrForInstanceSegmentation` (prefix chain `model` -> `model`): + + "model.model.backbone.backbone.x" -> "backbone.backbone.x" + "model.class_labels_classifier.x" -> "class_labels_classifier.x" + "query_features_block.mlp.fc1.x" -> "query_features_block.mlp.fc1.x" + """ + prefix = getattr(model, "base_model_prefix", "") + if not prefix or not key.startswith(prefix + "."): + return key + stripped_key = key[len(prefix) + 1 :] + inner_model = getattr(model, prefix, None) + if inner_model is not None: + stripped_key = _strip_model_prefix_for_save(stripped_key, inner_model) + return stripped_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 +1117,19 @@ 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: + for prefix in reversed(valid_prefixes): + if renamed_key.startswith(prefix + "."): + candidate = renamed_key[len(prefix) + 1 :] + if candidate in meta_state_dict: + renamed_key = candidate + break + candidate = f"{prefix}.{renamed_key}" + if candidate in meta_state_dict: + renamed_key = candidate + break return renamed_key, source_pattern @@ -1161,7 +1227,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 +1282,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 +1444,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) diff --git a/tests/models/conditional_detr/test_modeling_conditional_detr.py b/tests/models/conditional_detr/test_modeling_conditional_detr.py index e2eeec9bfdfa..eabbe9194fdb 100644 --- a/tests/models/conditional_detr/test_modeling_conditional_detr.py +++ b/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): diff --git a/tests/models/detr/test_modeling_detr.py b/tests/models/detr/test_modeling_detr.py index c4baec276f4f..2943ef755e34 100644 --- a/tests/models/detr/test_modeling_detr.py +++ b/tests/models/detr/test_modeling_detr.py @@ -16,17 +16,12 @@ import copy import inspect import math -import os -import re -import tempfile import unittest from functools import cached_property from parameterized import parameterized from transformers import DetrConfig, 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 Expectations, require_timm, require_torch, require_vision, slow, torch_device from ...test_configuration_common import ConfigTester @@ -34,7 +29,6 @@ TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION, ModelTesterMixin, _test_eager_matches_sdpa_inference, - compare_state_dicts, floats_tensor, ) from ...test_pipeline_mixin import PipelineTesterMixin @@ -42,7 +36,6 @@ if is_torch_available(): import torch - from safetensors.torch import load_file from transformers import DetrForObjectDetection, DetrForSegmentation, DetrModel @@ -206,88 +199,6 @@ class DetrModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): test_missing_keys = False zero_init_hidden_state = True - 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 != DetrForSegmentation: - 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)) - # special case for head models def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 9dbf44c03c12..519468e6ec80 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -4773,6 +4773,11 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_ config_to_set.mlp_only_layers = [0] # same but for qwens config_to_set.num_dense_layers = 1 # lfm2_moe + # Precompute state dict keys for every model class to detect dead conversion + # rules: a rule skipped for the current class must still apply to at least one. + all_classes_model_keys = { + cls: list(cls(copy.deepcopy(config)).state_dict().keys()) for cls in self.all_model_classes + } for model_class in self.all_model_classes: if skip_base_model and "For" not in model_class.__name__: continue @@ -4817,6 +4822,19 @@ 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 + + # Skip rules whose target doesn't appear in this model class (e.g. class-specific head rules), + # but assert the rule still matches at least one class + if not any(re.search(target_pattern_reversed, k) for k in model_keys): + self.assertTrue( + any( + any(re.search(target_pattern_reversed, k) for k in keys) + for keys in all_classes_model_keys.values() + ), + f"`{target_pattern_reversed}` in `{conversion}` does not match any " + "model class — the rule may be dead code or incorrectly written.", + ) + continue num_matches = sum(re.search(source_pattern, key) is not None for key in serialized_keys) self.assertTrue( num_matches > 0, From 67cecfa7b11cfe0ac68fdf8cb86c44f1906543e2 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 13 Apr 2026 23:52:51 +0000 Subject: [PATCH 2/5] use correct base_model_prefix for DetrForSegmentation --- src/transformers/models/detr/modeling_detr.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/models/detr/modeling_detr.py b/src/transformers/models/detr/modeling_detr.py index 384cc388cfd7..274afff09f3b 100644 --- a/src/transformers/models/detr/modeling_detr.py +++ b/src/transformers/models/detr/modeling_detr.py @@ -1435,6 +1435,8 @@ def forward( """ ) class DetrForSegmentation(DetrPreTrainedModel): + base_model_prefix = "detr" + def __init__(self, config: DetrConfig): super().__init__(config) From 98ca41bc7e035077a7870fa219259a408ef56969 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 13 Apr 2026 23:58:31 +0000 Subject: [PATCH 3/5] fix modular --- .../models/conditional_detr/modeling_conditional_detr.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/models/conditional_detr/modeling_conditional_detr.py b/src/transformers/models/conditional_detr/modeling_conditional_detr.py index 2f10c81b38e1..cbf008d1bce7 100644 --- a/src/transformers/models/conditional_detr/modeling_conditional_detr.py +++ b/src/transformers/models/conditional_detr/modeling_conditional_detr.py @@ -1621,6 +1621,8 @@ def _set_aux_loss(self, outputs_class, outputs_coord): """ ) class ConditionalDetrForSegmentation(ConditionalDetrPreTrainedModel): + base_model_prefix = "conditional_detr" + def __init__(self, config: ConditionalDetrConfig): super().__init__(config) From b087c836297b515e59c0585cb3b3a173f1f8fa25 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 15 Apr 2026 15:25:46 +0000 Subject: [PATCH 4/5] improve nested models loading --- src/transformers/core_model_loading.py | 89 ++++++++++--------- .../modeling_conditional_detr.py | 4 +- src/transformers/models/detr/modeling_detr.py | 4 +- 3 files changed, 52 insertions(+), 45 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index e8bea246ebe1..144098667112 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1032,39 +1032,34 @@ class SkipParameters(Exception): def _compute_all_prefixes(model) -> list[str]: """ - Return all cumulative `base_model_prefix` paths for the model's nesting hierarchy, - ordered from shortest to longest. + Return all base-model prefix paths reachable from `model`, ordered shortest-first (BFS). - Examples: - - RfDetrModel / RfDetrForObjectDetection -> ["model"] - RfDetrForInstanceSegmentation -> ["model", "model.model"] - ConditionalDetrForPanopticSegmentation -> ["conditional_detr", "conditional_detr.model"] - LlamaForCausalLM -> ["model"] - """ - prefixes: list[str] = [] - current_model = model - accumulated_prefix = "" + `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. - while True: - prefix = getattr(current_model, "base_model_prefix", "") - if not prefix: - break + Multiple base-model children are supported (e.g. a multi-modal model that contains + both `self.vision_model` and `self.text_model`). - next_accumulated = f"{accumulated_prefix}.{prefix}" if accumulated_prefix else prefix - prefixes.append(next_accumulated) - - inner_model = getattr(current_model, prefix, None) - if inner_model is None: - break # current_model is the leaf base model - - # Stop when the inner model is itself a leaf (no deeper nesting to traverse). - inner_prefix = getattr(inner_model, "base_model_prefix", "") - if not inner_prefix or getattr(inner_model, inner_prefix, None) is None: - break + Examples: - accumulated_prefix = next_accumulated - current_model = inner_model + 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 @@ -1075,20 +1070,21 @@ def _strip_model_prefix_for_save(key: str, model) -> str: reverse conversion rules (written relative to the innermost base model) operate on bare keys regardless of nesting depth. - Examples for `RfDetrForInstanceSegmentation` (prefix chain `model` -> `model`): + We identify each prefix level by finding the direct child whose `base_model_prefix` + matches its attribute name (same logic as `_compute_all_prefixes`). - "model.model.backbone.backbone.x" -> "backbone.backbone.x" - "model.class_labels_classifier.x" -> "class_labels_classifier.x" - "query_features_block.mlp.fc1.x" -> "query_features_block.mlp.fc1.x" + 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" """ - prefix = getattr(model, "base_model_prefix", "") - if not prefix or not key.startswith(prefix + "."): - return key - stripped_key = key[len(prefix) + 1 :] - inner_model = getattr(model, prefix, None) - if inner_model is not None: - stripped_key = _strip_model_prefix_for_save(stripped_key, inner_model) - return stripped_key + 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 rename_source_key( @@ -1130,6 +1126,17 @@ def rename_source_key( if candidate in meta_state_dict: renamed_key = candidate break + # If we still don't have a match, the checkpoint may originate from a model that wraps the + # target model at 2 or more nesting levels (e.g. loading a DetrForSegmentation checkpoint + # into DetrModel), so we search for a valid prefix anywhere within the key. + for prefix in reversed(valid_prefixes): + # remove the prefix from the key until we don't have it anymore (in case of multiple levels of nesting with the same prefix) + candidate = renamed_key + while prefix in candidate: + candidate = "".join(candidate.split(prefix + ".", maxsplit=1)[1:]) + if candidate in meta_state_dict: + renamed_key = candidate + break return renamed_key, source_pattern diff --git a/src/transformers/models/conditional_detr/modeling_conditional_detr.py b/src/transformers/models/conditional_detr/modeling_conditional_detr.py index cbf008d1bce7..a2c203d763a1 100644 --- a/src/transformers/models/conditional_detr/modeling_conditional_detr.py +++ b/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) @@ -1621,8 +1623,6 @@ def _set_aux_loss(self, outputs_class, outputs_coord): """ ) class ConditionalDetrForSegmentation(ConditionalDetrPreTrainedModel): - base_model_prefix = "conditional_detr" - def __init__(self, config: ConditionalDetrConfig): super().__init__(config) diff --git a/src/transformers/models/detr/modeling_detr.py b/src/transformers/models/detr/modeling_detr.py index 274afff09f3b..292917a4f2a1 100644 --- a/src/transformers/models/detr/modeling_detr.py +++ b/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) @@ -1435,8 +1437,6 @@ def forward( """ ) class DetrForSegmentation(DetrPreTrainedModel): - base_model_prefix = "detr" - def __init__(self, config: DetrConfig): super().__init__(config) From 0c6add7be6386c77353fd9f971c37d444f16bd4b Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 15 Apr 2026 16:19:17 +0000 Subject: [PATCH 5/5] edge cases + tests --- src/transformers/core_model_loading.py | 71 +++++++++++++++++-------- tests/models/detr/test_modeling_detr.py | 26 +++++++++ 2 files changed, 76 insertions(+), 21 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 144098667112..c7b599d3c5f7 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1087,6 +1087,55 @@ def _strip_model_prefix_for_save(key: str, model) -> str: 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], @@ -1116,27 +1165,7 @@ def rename_source_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: - for prefix in reversed(valid_prefixes): - if renamed_key.startswith(prefix + "."): - candidate = renamed_key[len(prefix) + 1 :] - if candidate in meta_state_dict: - renamed_key = candidate - break - candidate = f"{prefix}.{renamed_key}" - if candidate in meta_state_dict: - renamed_key = candidate - break - # If we still don't have a match, the checkpoint may originate from a model that wraps the - # target model at 2 or more nesting levels (e.g. loading a DetrForSegmentation checkpoint - # into DetrModel), so we search for a valid prefix anywhere within the key. - for prefix in reversed(valid_prefixes): - # remove the prefix from the key until we don't have it anymore (in case of multiple levels of nesting with the same prefix) - candidate = renamed_key - while prefix in candidate: - candidate = "".join(candidate.split(prefix + ".", maxsplit=1)[1:]) - if candidate in meta_state_dict: - renamed_key = candidate - break + renamed_key = _resolve_key_for_prefix_nesting(renamed_key, valid_prefixes, meta_state_dict) return renamed_key, source_pattern diff --git a/tests/models/detr/test_modeling_detr.py b/tests/models/detr/test_modeling_detr.py index 2943ef755e34..f1a2fdbea70b 100644 --- a/tests/models/detr/test_modeling_detr.py +++ b/tests/models/detr/test_modeling_detr.py @@ -16,6 +16,7 @@ import copy import inspect import math +import tempfile import unittest from functools import cached_property @@ -513,6 +514,31 @@ def test_greyscale_images(self): self.assertTrue(outputs) + def test_nested_base_model_prefix_checkpoint_loading(self): + """Segmentation checkpoints load into Seg / OD / backbone without missing keys; backbone-only checkpoints load + without unexpected keys (nested `base_model_prefix` key resolution).""" + config = self.model_tester.get_config() + + with tempfile.TemporaryDirectory() as seg_ckpt_dir: + DetrForSegmentation(config).save_pretrained(seg_ckpt_dir) + for model_class in (DetrForSegmentation, DetrForObjectDetection, DetrModel): + _, info = model_class.from_pretrained(seg_ckpt_dir, output_loading_info=True) + self.assertEqual( + info["missing_keys"], + set(), + msg=f"Seg checkpoint -> {model_class.__name__}: missing_keys={sorted(info['missing_keys'])}", + ) + + with tempfile.TemporaryDirectory() as base_ckpt_dir: + DetrModel(config).save_pretrained(base_ckpt_dir) + for model_class in (DetrForSegmentation, DetrForObjectDetection, DetrModel): + _, info = model_class.from_pretrained(base_ckpt_dir, output_loading_info=True) + self.assertEqual( + info["unexpected_keys"], + set(), + msg=f"DetrModel checkpoint -> {model_class.__name__}: unexpected_keys={sorted(info['unexpected_keys'])}", + ) + # override test_eager_matches_sdpa_inference to set use_attention_mask to False # as masks used in test are not adapted to the ones used in the model @parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)