From 7ad712a1ac5e082676504114168a67075c382baa Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 27 Apr 2026 18:51:04 +0000 Subject: [PATCH 01/15] mappings on classes, scoping for every transforms --- src/transformers/conversion_mapping.py | 156 +++++++++++++++---------- src/transformers/core_model_loading.py | 32 +++-- 2 files changed, 114 insertions(+), 74 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index a6e7b3734f9f..313e9f197146 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -348,24 +348,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming("out_proj", "o_proj"), WeightRenaming(r"layers.(\d+).fc1", r"layers.\1.mlp.fc1"), WeightRenaming(r"layers.(\d+).fc2", r"layers.\1.mlp.fc2"), - # `DetrForSegmentation` - WeightRenaming("bbox_attention.q_linear", "bbox_attention.q_proj"), - WeightRenaming("bbox_attention.k_linear", "bbox_attention.k_proj"), - # Mask head refactor - WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), - WeightRenaming("mask_head.gn1", "mask_head.conv1.norm"), - WeightRenaming("mask_head.lay2", "mask_head.conv2.conv"), - WeightRenaming("mask_head.gn2", "mask_head.conv2.norm"), - WeightRenaming("mask_head.adapter1", "mask_head.fpn_stages.0.fpn_adapter"), - WeightRenaming("mask_head.lay3", "mask_head.fpn_stages.0.refine.conv"), - WeightRenaming("mask_head.gn3", "mask_head.fpn_stages.0.refine.norm"), - WeightRenaming("mask_head.adapter2", "mask_head.fpn_stages.1.fpn_adapter"), - WeightRenaming("mask_head.lay4", "mask_head.fpn_stages.1.refine.conv"), - WeightRenaming("mask_head.gn4", "mask_head.fpn_stages.1.refine.norm"), - WeightRenaming("mask_head.adapter3", "mask_head.fpn_stages.2.fpn_adapter"), - WeightRenaming("mask_head.lay5", "mask_head.fpn_stages.2.refine.conv"), - WeightRenaming("mask_head.gn5", "mask_head.fpn_stages.2.refine.norm"), - WeightRenaming("mask_head.out_lay", "mask_head.output_conv"), ], "rt_detr": [ WeightRenaming("out_proj", "o_proj"), @@ -394,24 +376,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming( r"decoder.layers.(\d+).ca_qpos_sine_proj", r"decoder.layers.\1.encoder_attn.q_pos_sine_proj" ), - # The rest of patterns are used only in `ConditionalDetrForSegmentation` - WeightRenaming("bbox_attention.q_linear", "bbox_attention.q_proj"), - WeightRenaming("bbox_attention.k_linear", "bbox_attention.k_proj"), - # Mask head refactor - WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), - WeightRenaming("mask_head.gn1", "mask_head.conv1.norm"), - WeightRenaming("mask_head.lay2", "mask_head.conv2.conv"), - WeightRenaming("mask_head.gn2", "mask_head.conv2.norm"), - WeightRenaming("mask_head.adapter1", "mask_head.fpn_stages.0.fpn_adapter"), - WeightRenaming("mask_head.lay3", "mask_head.fpn_stages.0.refine.conv"), - WeightRenaming("mask_head.gn3", "mask_head.fpn_stages.0.refine.norm"), - WeightRenaming("mask_head.adapter2", "mask_head.fpn_stages.1.fpn_adapter"), - WeightRenaming("mask_head.lay4", "mask_head.fpn_stages.1.refine.conv"), - WeightRenaming("mask_head.gn4", "mask_head.fpn_stages.1.refine.norm"), - WeightRenaming("mask_head.adapter3", "mask_head.fpn_stages.2.fpn_adapter"), - WeightRenaming("mask_head.lay5", "mask_head.fpn_stages.2.refine.conv"), - WeightRenaming("mask_head.gn5", "mask_head.fpn_stages.2.refine.norm"), - WeightRenaming("mask_head.out_lay", "mask_head.output_conv"), ], "deformable_detr": [ WeightRenaming("backbone.conv_encoder", "backbone"), @@ -580,6 +544,28 @@ def _build_checkpoint_conversion_mapping(): ), ] + mapping["DetrForSegmentation"] = [ + WeightRenaming("bbox_attention.q_linear", "bbox_attention.q_proj"), + WeightRenaming("bbox_attention.k_linear", "bbox_attention.k_proj"), + WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), + WeightRenaming("mask_head.gn1", "mask_head.conv1.norm"), + WeightRenaming("mask_head.lay2", "mask_head.conv2.conv"), + WeightRenaming("mask_head.gn2", "mask_head.conv2.norm"), + WeightRenaming("mask_head.adapter1", "mask_head.fpn_stages.0.fpn_adapter"), + WeightRenaming("mask_head.lay3", "mask_head.fpn_stages.0.refine.conv"), + WeightRenaming("mask_head.gn3", "mask_head.fpn_stages.0.refine.norm"), + WeightRenaming("mask_head.adapter2", "mask_head.fpn_stages.1.fpn_adapter"), + WeightRenaming("mask_head.lay4", "mask_head.fpn_stages.1.refine.conv"), + WeightRenaming("mask_head.gn4", "mask_head.fpn_stages.1.refine.norm"), + WeightRenaming("mask_head.adapter3", "mask_head.fpn_stages.2.fpn_adapter"), + WeightRenaming("mask_head.lay5", "mask_head.fpn_stages.2.refine.conv"), + WeightRenaming("mask_head.gn5", "mask_head.fpn_stages.2.refine.norm"), + WeightRenaming("mask_head.out_lay", "mask_head.output_conv"), + ] + mapping["ConditionalDetrForSegmentation"] = mapping["DetrForSegmentation"].copy() + mapping["DetrForSegmentation"] = mapping["detr"].copy() + mapping["DetrForSegmentation"] + mapping["ConditionalDetrForSegmentation"] = mapping["conditional_detr"].copy() + mapping["ConditionalDetrForSegmentation"] + mapping["ernie4_5_moe"] = mapping["qwen2_moe"].copy() mapping["ernie4_5_moe"] += [ WeightRenaming("mlp.moe_statics.e_score_correction_bias", "mlp.gate.moe_statics.e_score_correction_bias") @@ -621,30 +607,47 @@ def get_checkpoint_conversion_mapping(model_type): def register_checkpoint_conversion_mapping( - model_type: str, + model_type_or_class_name: str, mapping: list[WeightConverter | WeightRenaming], overwrite: bool = False, ) -> None: + """ + Register a conversion mapping for a model type string or a class name. + + Class names take priority over ``model_type`` strings during lookup (see + :func:`extract_weight_conversions_for_model`), making it possible to define + task-head-specific or class-specific conversions that differ from the shared + ``model_type`` baseline. + """ global _checkpoint_conversion_mapping_cache if _checkpoint_conversion_mapping_cache is None: _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping() - if model_type in _checkpoint_conversion_mapping_cache and not overwrite: - raise ValueError(f"Model type {model_type} already exists in the checkpoint conversion mapping.") - _checkpoint_conversion_mapping_cache[model_type] = mapping + if model_type_or_class_name in _checkpoint_conversion_mapping_cache and not overwrite: + raise ValueError( + f"Conversion mapping for '{model_type_or_class_name}' already exists. " + f"Pass overwrite=True to replace it." + ) + _checkpoint_conversion_mapping_cache[model_type_or_class_name] = mapping -def extract_weight_conversions_for_model(model: PreTrainedModel, model_prefix: str) -> list[WeightTransform] | None: +def extract_weight_conversions_for_model( + model: PreTrainedModel, +) -> list[WeightTransform] | None: + """ + Return the registered conversion list for ``model``, or ``None`` if none exists. + + Looks up by class name first (enables task-head-specific overrides), then + falls back to ``model.config.model_type``. Transforms are returned + unmodified; the caller sets ``scope_prefix`` on each transform for sub-module isolation. + """ + class_name = type(model).__name__ model_type = getattr(model.config, "model_type", None) - if model_type is not None: - model_specific_conversions = get_checkpoint_conversion_mapping(model_type) - # In this case, add the prefix to `PrefixChange` instances, in order to know where to add/remove the prefix - if model_specific_conversions is not None and model_prefix != "": - for i, conversion in enumerate(model_specific_conversions): - # In this case, add the prefix, as otherwise we don't know where we need to re-add it exactly in the module name chain - if isinstance(conversion, PrefixChange): - model_specific_conversions[i] = conversion.with_submodel_prefix(model_prefix) - return model_specific_conversions - return None + + # Class name takes priority — allows ForXxx-specific overrides + conversions = get_checkpoint_conversion_mapping(class_name) + if conversions is None and model_type is not None: + conversions = get_checkpoint_conversion_mapping(model_type) + return conversions def get_model_conversion_mapping( @@ -654,8 +657,15 @@ def get_model_conversion_mapping( add_legacy: bool = True, ) -> list[WeightTransform]: """ - For a given `model`, obtain the weight conversion mapping if any are registered either as a simple renaming - `_checkpoint_conversion_mapping` class argument, or in the general WeightConverter mapping. + Collect the ordered list of weight transforms for ``model`` (used during + loading and, when reversed, during saving). + + Each ``PreTrainedModel`` sub-module is looked up by class name then + ``model_type``. Root transforms are applied globally; sub-module transforms + have their ``scope_prefix`` set so they only match keys under that prefix. After any + sub-module is processed, both its class name and ``model_type`` are marked + seen to prevent ``XForY`` / ``XModel`` pairs from applying the same mapping + twice via different lookup paths. """ # Lazy import to avoid circular import issues from .modeling_utils import PreTrainedModel @@ -667,16 +677,36 @@ def get_model_conversion_mapping( if key_mapping is not None: weight_conversions = [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()] - # Model have several `PreTrainedModel` within with the same model type, for example: XForConditionalGeneration -> XModel - # We don't want to apply the same conversion pattern twice because of that - seen_model_types = set() - # Recurse over submodules and collect all conversions - for name, submodule in model.named_modules(): - if isinstance(submodule, PreTrainedModel) and submodule.config.model_type not in seen_model_types: - conversions = extract_weight_conversions_for_model(submodule, name) - if conversions is not None: - weight_conversions.extend(conversions) - seen_model_types.add(submodule.config.model_type) + seen_identifiers: set[str] = set() + + for module_name, submodule in model.named_modules(): + if not isinstance(submodule, PreTrainedModel): + continue + + class_name = type(submodule).__name__ + model_type = getattr(submodule.config, "model_type", None) + + # Skip if this architecture was already processed via either lookup path. + if class_name in seen_identifiers or (model_type and model_type in seen_identifiers): + continue + + conversions = extract_weight_conversions_for_model(submodule) + if conversions is None: + continue + + is_root_model = module_name == "" + if not is_root_model: + # Scope each transform so it only matches keys under this sub-module's prefix. + for transform in conversions: + transform.scope_prefix = module_name + weight_conversions.extend(conversions) + + # Only mark seen when a mapping was actually applied, so that a root model + # with no mapping does not prematurely block sub-modules with the same + # model_type from getting their own scoped transforms entry. + seen_identifiers.add(class_name) + if model_type: + seen_identifiers.add(model_type) if add_legacy: weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index cd0710649c91..84feded52866 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -591,6 +591,7 @@ class WeightTransform: "_original_source_patterns", "_original_target_patterns", "_was_used", + "scope_prefix", ) def __init__(self, source_patterns: str | list[str], target_patterns: str | list[str]): @@ -608,6 +609,9 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list # Flag to notice if the Transform was used self._was_used = False + # Optional prefix scope: when set, this transform only applies to keys starting with + # ``scope_prefix + "."``, stripping / re-attaching the prefix around the pattern match. + self.scope_prefix: str | None = None # We need to process a few exceptions here when instantiating the reverse mapping (i.e. the targets become # sources, and sources become targets). The issues lie in the sources usually, so here we need to check the @@ -680,8 +684,17 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: In case of a one-to-many transform, i.e. we have several target patterns, the matching source pattern will be replaced by the first of all the target patterns (they are then correctly expanded in the Operations). """ + # When scoped, only process keys under the prefix; patterns operate on the bare suffix. + prefix_dot = None + key_to_match = source_key + if self.scope_prefix is not None: + prefix_dot = self.scope_prefix + "." + if not source_key.startswith(prefix_dot): + return source_key, None + key_to_match = source_key[len(prefix_dot):] + # Try matching one of the alternation branches - match_object = self.compiled_sources.search(source_key) + match_object = self.compiled_sources.search(key_to_match) if match_object is None: return source_key, None @@ -699,7 +712,9 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: # inside that matched named group replaced_group_idx = self.compiled_sources.groupindex[matching_group_name] + 1 replacement = replacement.replace(r"\1", match_object.group(replaced_group_idx)) - renamed_key = source_key.replace(match_object.group(0), replacement, 1) + renamed_key = key_to_match.replace(match_object.group(0), replacement, 1) + if prefix_dot is not None: + renamed_key = prefix_dot + renamed_key return renamed_key, source_pattern_that_matched def reverse_transform(self) -> WeightTransform: @@ -717,7 +732,7 @@ def reverse_transform(self) -> WeightTransform: reverse_transform = self.__class__( source_patterns=self._original_target_patterns, target_patterns=self._original_source_patterns, **kwargs ) - + reverse_transform.scope_prefix = self.scope_prefix return reverse_transform def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: @@ -836,16 +851,11 @@ def reverse_transform(self) -> WeightTransform: raise ValueError("Cannot reverse the transform with TP or quantization") # Only one of the 2 can ever be used, so 1 is always None - return PrefixChange( + result = PrefixChange( prefix_to_add=self.prefix_to_remove, prefix_to_remove=self.prefix_to_add, model_prefix=self.model_prefix ) - - def with_submodel_prefix(self, prefix: str) -> PrefixChange: - new_prefix = f"{prefix}.{self.model_prefix}" if self.model_prefix != "" else prefix - return PrefixChange( - prefix_to_add=self.prefix_to_add, prefix_to_remove=self.prefix_to_remove, model_prefix=new_prefix - ) - + result.scope_prefix = self.scope_prefix + return result # List of classes that are known to be able to use m:n _INTERNAL_MANY_TO_MANY_CONVERSIONS = ( From c63a7d8cbfa9172f45d9ad6c4f763cf24dcbcc4d Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 27 Apr 2026 19:09:35 +0000 Subject: [PATCH 02/15] fix style --- src/transformers/conversion_mapping.py | 7 ++++--- src/transformers/core_model_loading.py | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 313e9f197146..8aac9008f07f 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -564,7 +564,9 @@ def _build_checkpoint_conversion_mapping(): ] mapping["ConditionalDetrForSegmentation"] = mapping["DetrForSegmentation"].copy() mapping["DetrForSegmentation"] = mapping["detr"].copy() + mapping["DetrForSegmentation"] - mapping["ConditionalDetrForSegmentation"] = mapping["conditional_detr"].copy() + mapping["ConditionalDetrForSegmentation"] + mapping["ConditionalDetrForSegmentation"] = ( + mapping["conditional_detr"].copy() + mapping["ConditionalDetrForSegmentation"] + ) mapping["ernie4_5_moe"] = mapping["qwen2_moe"].copy() mapping["ernie4_5_moe"] += [ @@ -624,8 +626,7 @@ def register_checkpoint_conversion_mapping( _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping() if model_type_or_class_name in _checkpoint_conversion_mapping_cache and not overwrite: raise ValueError( - f"Conversion mapping for '{model_type_or_class_name}' already exists. " - f"Pass overwrite=True to replace it." + f"Conversion mapping for '{model_type_or_class_name}' already exists. Pass overwrite=True to replace it." ) _checkpoint_conversion_mapping_cache[model_type_or_class_name] = mapping diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 84feded52866..b7313515c4b9 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -691,7 +691,7 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: prefix_dot = self.scope_prefix + "." if not source_key.startswith(prefix_dot): return source_key, None - key_to_match = source_key[len(prefix_dot):] + key_to_match = source_key[len(prefix_dot) :] # Try matching one of the alternation branches match_object = self.compiled_sources.search(key_to_match) @@ -857,6 +857,7 @@ def reverse_transform(self) -> WeightTransform: result.scope_prefix = self.scope_prefix return result + # List of classes that are known to be able to use m:n _INTERNAL_MANY_TO_MANY_CONVERSIONS = ( ErnieFuseAndSplitTextVisionExperts, From 8f726c7f807f25cdedfac347d81515c0eddb3c71 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 27 Apr 2026 20:22:27 +0000 Subject: [PATCH 03/15] Fix deduplication removes submodel mappings of the same type --- src/transformers/conversion_mapping.py | 28 +++++++++++++++----------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 8aac9008f07f..a6eadf46f97f 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -343,7 +343,7 @@ def _build_checkpoint_conversion_mapping(): operations=[ErnieFuseAndSplitTextVisionExperts(stack_dim=0, concat_dim=1)], ), ], - "detr": [ + "DetrModel": [ WeightRenaming("backbone.conv_encoder", "backbone"), WeightRenaming("out_proj", "o_proj"), WeightRenaming(r"layers.(\d+).fc1", r"layers.\1.mlp.fc1"), @@ -355,7 +355,7 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"layers.(\d+).fc2", r"layers.\1.mlp.fc2"), WeightRenaming(r"encoder.encoder.(\d+).layers", r"encoder.aifi.\1.layers"), ], - "conditional_detr": [ + "ConditionalDetrModel": [ WeightRenaming("backbone.conv_encoder", "backbone"), WeightRenaming("self_attn.out_proj", "self_attn.o_proj"), WeightRenaming("encoder_attn.out_proj", "encoder_attn.o_proj"), @@ -543,7 +543,8 @@ def _build_checkpoint_conversion_mapping(): target_patterns=".parametrizations.weight.original1", ), ] - + # Base DetrModel/ConditionalDetrModel transforms are picked up automatically as + # scoped sub-module transforms; only the segmentation-specific patterns are needed here. mapping["DetrForSegmentation"] = [ WeightRenaming("bbox_attention.q_linear", "bbox_attention.q_proj"), WeightRenaming("bbox_attention.k_linear", "bbox_attention.k_proj"), @@ -563,10 +564,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming("mask_head.out_lay", "mask_head.output_conv"), ] mapping["ConditionalDetrForSegmentation"] = mapping["DetrForSegmentation"].copy() - mapping["DetrForSegmentation"] = mapping["detr"].copy() + mapping["DetrForSegmentation"] - mapping["ConditionalDetrForSegmentation"] = ( - mapping["conditional_detr"].copy() + mapping["ConditionalDetrForSegmentation"] - ) mapping["ernie4_5_moe"] = mapping["qwen2_moe"].copy() mapping["ernie4_5_moe"] += [ @@ -691,7 +688,13 @@ def get_model_conversion_mapping( if class_name in seen_identifiers or (model_type and model_type in seen_identifiers): continue - conversions = extract_weight_conversions_for_model(submodule) + # Try class name first, then model_type. Track which path produced the hit so + # we know whether to block model_type for subsequent sub-modules (see below). + conversions = get_checkpoint_conversion_mapping(class_name) + found_via_class = conversions is not None + if not found_via_class and model_type is not None: + conversions = get_checkpoint_conversion_mapping(model_type) + if conversions is None: continue @@ -702,11 +705,12 @@ def get_model_conversion_mapping( transform.scope_prefix = module_name weight_conversions.extend(conversions) - # Only mark seen when a mapping was actually applied, so that a root model - # with no mapping does not prematurely block sub-modules with the same - # model_type from getting their own scoped transforms entry. seen_identifiers.add(class_name) - if model_type: + # Only block model_type when the hit was via model_type. When the hit was via + # class name, sub-modules that share the same model_type but have no class-specific + # mapping of their own (e.g. DetrModel under DetrForSegmentation) must still be + # reachable so their base transforms are picked up and scoped automatically. + if not found_via_class and model_type: seen_identifiers.add(model_type) if add_legacy: From fd2c613bd03352355af7d38b6c00c9b75573a32c Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 27 Apr 2026 21:39:38 +0000 Subject: [PATCH 04/15] Fix scoped WeightConverter not applied in the correct order, now interleaving renames and converts ops --- src/transformers/core_model_loading.py | 50 ++++++++++----------- src/transformers/integrations/accelerate.py | 8 ++-- src/transformers/integrations/deepspeed.py | 7 ++- tests/utils/test_core_model_loading.py | 8 ++-- 4 files changed, 35 insertions(+), 38 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index b7313515c4b9..1afb5de91478 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1123,30 +1123,34 @@ class SkipParameters(Exception): def rename_source_key( source_key: str, - weight_renamings: list[WeightRenaming], - weight_converters: list[WeightConverter], + weight_transforms: list[WeightTransform], prefix: 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 + Rename a source key given all the weight transforms we have. Also takes care of adding/removing the base model prefix during loading if necessary. + + Transforms are applied in their natural interleaved order (the order they appear in the list). + When a ``WeightConverter`` matches, it is recorded as the source pattern and remaining + ``WeightRenaming`` transforms continue to run, which is required when a scoped + ``WeightConverter`` must fire *before* a renaming that strips the scope prefix. """ renamed_key = source_key - # 1. apply all renamings in turns (if multiple match, it's the responsibility of the mappings to make sure they - # are coherent) - for renaming in weight_renamings: - renamed_key, _ = renaming.rename_source_key(renamed_key) - - # 2. apply renaming through weight conversions on the key if we have any WeightConverter (here we stop after - # the first match, as we assume only 1 converter can match any source key) source_pattern = None - for converter in weight_converters: - renamed_key, source_pattern = converter.rename_source_key(renamed_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) + for transform in weight_transforms: + if isinstance(transform, WeightConverter): + if source_pattern is not None: + # Already matched a converter; skip subsequent converters. + continue + renamed_key, sp = transform.rename_source_key(renamed_key) + if sp is not None: + source_pattern = sp + else: + renamed_key, _ = transform.rename_source_key(renamed_key) + + # 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) @@ -1288,7 +1292,6 @@ def convert_and_load_state_dict_in_model( else: thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS) - renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)] converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)] param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {} @@ -1303,13 +1306,11 @@ def convert_and_load_state_dict_in_model( state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) 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 - ) + # 1. Rename the key according to all renaming and weight conversion patterns. + renamed_key, source_pattern = rename_source_key(original_key, weight_mapping, prefix, 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` to be added. + renamed_key, source_pattern = rename_source_key(original_key, [], prefix, meta_model_state_dict) # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: @@ -1471,15 +1472,14 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch # Reverse all Transform to correctly match keys reverse_weight_conversion = [conversion.reverse_transform() for conversion in weight_conversions] # If we are still here, we need to create the (reverse) conversion mapping from scratch - renamings = [entry for entry in reverse_weight_conversion if isinstance(entry, WeightRenaming)] converters = [entry for entry in reverse_weight_conversion if isinstance(entry, WeightConverter)] pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} conversion_mapping = {} 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) + renamed_key, source_pattern = rename_source_key(original_key, reverse_weight_conversion) + if source_pattern is not None: new_converter = deepcopy(pattern_to_converter[source_pattern]) # each target key gets its own converter instance diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index c2b7fa603570..0b278ebb4d40 100644 --- a/src/transformers/integrations/accelerate.py +++ b/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,7 @@ 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, model.base_model_prefix, meta_state_dict)[0]: k for k in weight_map } # Prepare the index using existing safetensors files diff --git a/src/transformers/integrations/deepspeed.py b/src/transformers/integrations/deepspeed.py index 9703f642f8bc..79f3896cb48c 100644 --- a/src/transformers/integrations/deepspeed.py +++ b/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,13 @@ 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, 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 +385,7 @@ 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, model_state_dict) # Only process if the renamed key is in the model's state dict if renamed_key in model_state_dict: diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index a358822d19f8..85e77b834d3e 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -147,14 +147,14 @@ def test_sub_key_rewrites_targets(self): ] self.assertEqual( - rename_source_key("foo.block_sparse_moe.experts.3.w1.weight", renamings, [])[0], + rename_source_key("foo.block_sparse_moe.experts.3.w1.weight", renamings)[0], "foo.mlp.experts.gate_up_proj", ) self.assertEqual( - rename_source_key("foo.block_sparse_moe.experts.3.w2.weight", renamings, [])[0], + rename_source_key("foo.block_sparse_moe.experts.3.w2.weight", renamings)[0], "foo.mlp.experts.down_proj", ) - self.assertEqual(rename_source_key("model.language_model.lm_head.weight", renamings, [])[0], "language_model") + self.assertEqual(rename_source_key("model.language_model.lm_head.weight", renamings)[0], "language_model") def test_sub_key_no_match_returns_original(self): renamings = [ @@ -162,7 +162,7 @@ def test_sub_key_no_match_returns_original(self): ] key = "unrelated.key" - renamed_key, _ = rename_source_key(key, renamings, []) + renamed_key, _ = rename_source_key(key, renamings) self.assertEqual(renamed_key, key) From 28ed270290cfed18c1c64adefb7b9e645fe38eeb Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Mon, 27 Apr 2026 22:31:15 +0000 Subject: [PATCH 05/15] temp fix paligemma --- src/transformers/conversion_mapping.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index a6eadf46f97f..c0d2c18b2b2d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -103,6 +103,7 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"^vision_tower", target_patterns="model.vision_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], + "PaliGemmaModel": [], "llava_next": [ WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), From c75244e84711b141f07978e6689dff58745b17c3 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Tue, 28 Apr 2026 02:34:22 +0000 Subject: [PATCH 06/15] Fix incompatible mappings between head and base model for VLMs --- src/transformers/conversion_mapping.py | 70 +++++++++++++++++--------- 1 file changed, 46 insertions(+), 24 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index c0d2c18b2b2d..25e298ca37c0 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -62,19 +62,8 @@ "rt_detr_v2": "rt_detr", "pp_doclayout_v2": "rt_detr", "pp_doclayout_v3": "rt_detr", - "paligemma": "llava", - "aya_vision": "llava", - "got_ocr2": "llava", - "shieldgemma2": "llava", - "gemma3": "llava", - "internvl": "llava", - "llava_next_video": "llava_next", - "llava_onevision": "llava_next", - "vipllava": "llava", - "mistral3": "llava", "qwen2_5_vl": "qwen2_vl", "sam3_tracker_video": "sam3_tracker", - "pp_chart2table": "llava", "altclip_vision_model": "clip_vision_model", "chinese_clip_vision_model": "clip_vision_model", "clipseg_vision_model": "clip_vision_model", @@ -89,6 +78,31 @@ "siglip_text_model": "clip_text_model", "siglip2_text_model": "clip_text_model", "xclip_text_model": "clip_text_model", + # class-based mappings + "PaliGemmaModel": "LlavaModel", + "AyaVisionModel": "LlavaModel", + "GotOcr2Model": "LlavaModel", + "Gemma3Model": "LlavaModel", + "InternVLModel": "LlavaModel", + "VipLlavaModel": "LlavaModel", + "Mistral3Model": "LlavaModel", + "PPChart2TableModel": "LlavaModel", + "LlavaNextModel": "LlavaModel", + "LlavaNextVideoModel": "LlavaModel", + "LlavaOnevisionModel": "LlavaModel", + "FuyuModel": "LlavaModel", + "MllamaModel": "LlavaModel", + "ShieldGemma2ForImageClassification": "LlavaForConditionalGeneration", + "PaliGemmaForConditionalGeneration": "LlavaForConditionalGeneration", + "AyaVisionForConditionalGeneration": "LlavaForConditionalGeneration", + "GotOcr2ForConditionalGeneration": "LlavaForConditionalGeneration", + "Gemma3ForConditionalGeneration": "LlavaForConditionalGeneration", + "InternVLForConditionalGeneration": "LlavaForConditionalGeneration", + "VipLlavaForConditionalGeneration": "LlavaForConditionalGeneration", + "Mistral3ForConditionalGeneration": "LlavaForConditionalGeneration", + "PPChart2TableForConditionalGeneration": "LlavaForConditionalGeneration", + "LlavaNextVideoForConditionalGeneration": "LlavaNextForConditionalGeneration", + "LlavaOnevisionForConditionalGeneration": "LlavaNextForConditionalGeneration", } @@ -97,43 +111,51 @@ def _build_checkpoint_conversion_mapping(): "altclip": [ WeightRenaming(source_patterns=r"layer\.", target_patterns="layers."), ], - "llava": [ - WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), + "LlavaModel": [ + WeightRenaming(source_patterns=r"^language_model.model", target_patterns="language_model"), + ], + "LlavaForConditionalGeneration": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_tower", target_patterns="model.vision_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], - "PaliGemmaModel": [], - "llava_next": [ - WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), + "LlavaNextForConditionalGeneration": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_tower", target_patterns="model.vision_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), WeightRenaming(source_patterns=r"^image_newline", target_patterns="model.image_newline"), ], "clip_vision_model": [PrefixChange(prefix_to_remove="vision_model")], "clip_text_model": [PrefixChange(prefix_to_remove="text_model")], - "video_llava": [ - WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), + "VideoLlavaModel": [ + WeightRenaming(source_patterns=r"^language_model.model", target_patterns="language_model"), + ], + "VideoLlavaForConditionalGeneration": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^image_tower", target_patterns="model.image_tower"), WeightRenaming(source_patterns=r"^video_tower", target_patterns="model.video_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], - "fuyu": [ - WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), + "FuyuForCausalLM": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_embed_tokens", target_patterns="model.vision_embed_tokens"), ], - "mllama": [ - WeightRenaming(source_patterns=r"^language_model.model", target_patterns="model.language_model"), + "MllamaForConditionalGeneration": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_model", target_patterns="model.vision_model"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], - "emu3": [ - WeightRenaming(source_patterns=r"^text_model.model", target_patterns="model.text_model"), + "Emu3Model": [ + WeightRenaming(source_patterns=r"^text_model.model", target_patterns="text_model"), + ], + "Emu3ForConditionalGeneration": [ WeightRenaming(source_patterns=r"^text_model.lm_head", target_patterns="lm_head"), + WeightRenaming(source_patterns=r"^text_model", target_patterns="model.text_model"), WeightRenaming(source_patterns=r"^vqmodel", target_patterns="model.vqmodel"), ], "paddleocr_vl": [ From 1dbd3da751c0566991e6f6699f417c50c2166b97 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Tue, 28 Apr 2026 14:47:31 +0000 Subject: [PATCH 07/15] fix gemma3 mapping --- src/transformers/conversion_mapping.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 25e298ca37c0..5f86dc3b4a82 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -97,6 +97,7 @@ "AyaVisionForConditionalGeneration": "LlavaForConditionalGeneration", "GotOcr2ForConditionalGeneration": "LlavaForConditionalGeneration", "Gemma3ForConditionalGeneration": "LlavaForConditionalGeneration", + "Gemma3ForSequenceClassification": "LlavaForConditionalGeneration", "InternVLForConditionalGeneration": "LlavaForConditionalGeneration", "VipLlavaForConditionalGeneration": "LlavaForConditionalGeneration", "Mistral3ForConditionalGeneration": "LlavaForConditionalGeneration", From b799f324327521ed964b1bd089318a6df1818e15 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Tue, 28 Apr 2026 19:16:42 +0000 Subject: [PATCH 08/15] Fix more issues, address reviews --- src/transformers/conversion_mapping.py | 52 +++--- src/transformers/core_model_loading.py | 59 ++++-- src/transformers/modeling_utils.py | 6 + .../aya_vision/test_modeling_aya_vision.py | 3 - tests/models/colpali/test_modeling_colpali.py | 4 - .../models/colqwen2/test_modeling_colqwen2.py | 4 - tests/models/emu3/test_modeling_emu3.py | 3 - tests/models/fuyu/test_modeling_fuyu.py | 3 - tests/models/gemma3/test_modeling_gemma3.py | 3 - .../models/got_ocr2/test_modeling_got_ocr2.py | 3 - .../models/internvl/test_modeling_internvl.py | 3 - tests/models/llava/test_modeling_llava.py | 3 - .../llava_next/test_modeling_llava_next.py | 3 - .../test_modeling_llava_next_video.py | 3 - .../test_modeling_llava_onevision.py | 3 - .../models/mistral3/test_modeling_mistral3.py | 3 - tests/models/mllama/test_modeling_mllama.py | 3 - .../paligemma/test_modeling_paligemma.py | 3 - .../qwen2_5_vl/test_modeling_qwen2_5_vl.py | 3 - .../models/qwen2_vl/test_modeling_qwen2_vl.py | 3 - .../video_llava/test_modeling_video_llava.py | 3 - .../models/vipllava/test_modeling_vipllava.py | 3 - tests/test_modeling_common.py | 14 +- tests/utils/test_core_model_loading.py | 174 +++++++++++++++++- 24 files changed, 254 insertions(+), 110 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 5f86dc3b4a82..3d7b2bd8bd13 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -78,6 +78,17 @@ "siglip_text_model": "clip_text_model", "siglip2_text_model": "clip_text_model", "xclip_text_model": "clip_text_model", + "shield_gemma2": "llava", + "paligemma": "llava", + "aya_vision": "llava", + "got_ocr2": "llava", + "gemma3": "llava", + "internvl": "llava", + "vipllava": "llava", + "mistral3": "llava", + "pp_chart2table": "llava", + "llava_next_video": "llava_next", + "llava_onevision": "llava_next", # class-based mappings "PaliGemmaModel": "LlavaModel", "AyaVisionModel": "LlavaModel", @@ -92,18 +103,7 @@ "LlavaOnevisionModel": "LlavaModel", "FuyuModel": "LlavaModel", "MllamaModel": "LlavaModel", - "ShieldGemma2ForImageClassification": "LlavaForConditionalGeneration", - "PaliGemmaForConditionalGeneration": "LlavaForConditionalGeneration", - "AyaVisionForConditionalGeneration": "LlavaForConditionalGeneration", - "GotOcr2ForConditionalGeneration": "LlavaForConditionalGeneration", - "Gemma3ForConditionalGeneration": "LlavaForConditionalGeneration", - "Gemma3ForSequenceClassification": "LlavaForConditionalGeneration", - "InternVLForConditionalGeneration": "LlavaForConditionalGeneration", - "VipLlavaForConditionalGeneration": "LlavaForConditionalGeneration", - "Mistral3ForConditionalGeneration": "LlavaForConditionalGeneration", - "PPChart2TableForConditionalGeneration": "LlavaForConditionalGeneration", - "LlavaNextVideoForConditionalGeneration": "LlavaNextForConditionalGeneration", - "LlavaOnevisionForConditionalGeneration": "LlavaNextForConditionalGeneration", + "Qwen2_5_VLModel": "Qwen2VLModel", } @@ -115,13 +115,13 @@ def _build_checkpoint_conversion_mapping(): "LlavaModel": [ WeightRenaming(source_patterns=r"^language_model.model", target_patterns="language_model"), ], - "LlavaForConditionalGeneration": [ + "llava": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_tower", target_patterns="model.vision_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], - "LlavaNextForConditionalGeneration": [ + "llava_next": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_tower", target_patterns="model.vision_tower"), @@ -133,19 +133,19 @@ def _build_checkpoint_conversion_mapping(): "VideoLlavaModel": [ WeightRenaming(source_patterns=r"^language_model.model", target_patterns="language_model"), ], - "VideoLlavaForConditionalGeneration": [ + "video_llava": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^image_tower", target_patterns="model.image_tower"), WeightRenaming(source_patterns=r"^video_tower", target_patterns="model.video_tower"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], - "FuyuForCausalLM": [ + "fuyu": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_embed_tokens", target_patterns="model.vision_embed_tokens"), ], - "MllamaForConditionalGeneration": [ + "mllama": [ WeightRenaming(source_patterns=r"^language_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^language_model", target_patterns="model.language_model"), WeightRenaming(source_patterns=r"^vision_model", target_patterns="model.vision_model"), @@ -154,7 +154,7 @@ def _build_checkpoint_conversion_mapping(): "Emu3Model": [ WeightRenaming(source_patterns=r"^text_model.model", target_patterns="text_model"), ], - "Emu3ForConditionalGeneration": [ + "emu3": [ WeightRenaming(source_patterns=r"^text_model.lm_head", target_patterns="lm_head"), WeightRenaming(source_patterns=r"^text_model", target_patterns="model.text_model"), WeightRenaming(source_patterns=r"^vqmodel", target_patterns="model.vqmodel"), @@ -167,15 +167,12 @@ def _build_checkpoint_conversion_mapping(): target_patterns="model.language_model", ), ], + "Qwen2VLModel": [WeightRenaming(source_patterns=r"^model.", target_patterns="")], "qwen2_vl": [ + WeightRenaming(source_patterns=r"^visual", target_patterns="model.visual"), WeightRenaming( source_patterns=r"(? tuple[str, str | None]: + def _scoped_match(self, source_key: str) -> tuple[str | None, str, re.Match[str]] | None: """ - Return a tuple (renamed_key, source_pattern_producing_the_match). - Try renaming `source_key` according to the source and target patterns of the current WeightTransform. - In case of a one-to-many transform, i.e. we have several target patterns, the matching source pattern - will be replaced by the first of all the target patterns (they are then correctly expanded in the Operations). + Apply ``scope_prefix`` stripping (if any), then match ``compiled_sources`` against the suffix. + + Returns ``(prefix_dot, key_to_match, match_object)`` when a branch matches, where ``prefix_dot`` is ``None`` + if ``scope_prefix`` is unset, else ``f"{scope_prefix}."``. Returns ``None`` when out of scope or unmatched. + Does not set ``_was_used``. """ - # When scoped, only process keys under the prefix; patterns operate on the bare suffix. prefix_dot = None key_to_match = source_key if self.scope_prefix is not None: prefix_dot = self.scope_prefix + "." if not source_key.startswith(prefix_dot): - return source_key, None + return None key_to_match = source_key[len(prefix_dot) :] - # Try matching one of the alternation branches match_object = self.compiled_sources.search(key_to_match) if match_object is None: + return None + return (prefix_dot, key_to_match, match_object) + + def rename_source_key(self, source_key: str) -> tuple[str, str | None]: + """ + Return a tuple (renamed_key, source_pattern_producing_the_match). + Try renaming `source_key` according to the source and target patterns of the current WeightTransform. + In case of a one-to-many transform, i.e. we have several target patterns, the matching source pattern + will be replaced by the first of all the target patterns (they are then correctly expanded in the Operations). + """ + matched = self._scoped_match(source_key) + if matched is None: return source_key, None + prefix_dot, key_to_match, match_object = matched + # We have a match, so the Transform was used self._was_used = True @@ -1128,13 +1141,29 @@ def rename_source_key( meta_state_dict: dict | None = None, ) -> tuple[str, str | None]: """ - Rename a source key given all the weight transforms we have. Also takes care of adding/removing - the base model prefix during loading if necessary. - - Transforms are applied in their natural interleaved order (the order they appear in the list). - When a ``WeightConverter`` matches, it is recorded as the source pattern and remaining - ``WeightRenaming`` transforms continue to run, which is required when a scoped - ``WeightConverter`` must fire *before* a renaming that strips the scope prefix. + Rename a source key according to ``weight_transforms``, also handling the base model prefix. + + Transforms are applied in list order, interleaving ``WeightRenaming`` and ``WeightConverter`` + instances as they appear. The same list, reversed and with each transform individually + inverted, is used on the save path, so relative ordering is preserved in both directions. + + At most one ``WeightConverter`` fires per key; subsequent converters are skipped. + ``WeightRenaming`` always runs, even after a converter has already fired. + + Example (root rename followed by a scoped sub-model converter):: + + transforms = [ + WeightRenaming("^old_prefix", "model.vlm"), + WeightConverter("^q_proj", "qkv_proj", ...), # scope_prefix="model.vlm" + ] + # Load: "old_prefix.q_proj" + # → WeightRenaming → "model.vlm.q_proj" + # → WeightConverter → "model.vlm.qkv_proj" + # + # Save (inverted list, each transform reversed): + # "model.vlm.q_proj" + # → rev(WeightConverter) → "model.vlm.q_proj" + # → rev(WeightRenaming) → "old_prefix.q_proj" """ renamed_key = source_key source_pattern = None diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index b041964bbdfc..24526af755c8 100644 --- a/src/transformers/modeling_utils.py +++ b/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]: diff --git a/tests/models/aya_vision/test_modeling_aya_vision.py b/tests/models/aya_vision/test_modeling_aya_vision.py index c88f6889d123..b5e659826a21 100644 --- a/tests/models/aya_vision/test_modeling_aya_vision.py +++ b/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): diff --git a/tests/models/colpali/test_modeling_colpali.py b/tests/models/colpali/test_modeling_colpali.py index 88df8ade9a49..699d5019fbe1 100644 --- a/tests/models/colpali/test_modeling_colpali.py +++ b/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): diff --git a/tests/models/colqwen2/test_modeling_colqwen2.py b/tests/models/colqwen2/test_modeling_colqwen2.py index 110576ebe5c6..fb213177fb8b 100644 --- a/tests/models/colqwen2/test_modeling_colqwen2.py +++ b/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): diff --git a/tests/models/emu3/test_modeling_emu3.py b/tests/models/emu3/test_modeling_emu3.py index 4fa55987c70f..6b356bb659c9 100644 --- a/tests/models/emu3/test_modeling_emu3.py +++ b/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): diff --git a/tests/models/fuyu/test_modeling_fuyu.py b/tests/models/fuyu/test_modeling_fuyu.py index 6c0a2657c4f7..6f06c8d5f68d 100644 --- a/tests/models/fuyu/test_modeling_fuyu.py +++ b/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 diff --git a/tests/models/gemma3/test_modeling_gemma3.py b/tests/models/gemma3/test_modeling_gemma3.py index 913d6b9cf5ff..288c41eed6fb 100644 --- a/tests/models/gemma3/test_modeling_gemma3.py +++ b/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 diff --git a/tests/models/got_ocr2/test_modeling_got_ocr2.py b/tests/models/got_ocr2/test_modeling_got_ocr2.py index 404b88d08dde..0c5dbde9b68b 100644 --- a/tests/models/got_ocr2/test_modeling_got_ocr2.py +++ b/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): diff --git a/tests/models/internvl/test_modeling_internvl.py b/tests/models/internvl/test_modeling_internvl.py index c9b8d06ba9fa..190f2f02a99e 100644 --- a/tests/models/internvl/test_modeling_internvl.py +++ b/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 diff --git a/tests/models/llava/test_modeling_llava.py b/tests/models/llava/test_modeling_llava.py index 4e50c56eb55b..b0bcf5afbbbd 100644 --- a/tests/models/llava/test_modeling_llava.py +++ b/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 diff --git a/tests/models/llava_next/test_modeling_llava_next.py b/tests/models/llava_next/test_modeling_llava_next.py index a5bd146fcc6d..7e7b40e4eaba 100644 --- a/tests/models/llava_next/test_modeling_llava_next.py +++ b/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): diff --git a/tests/models/llava_next_video/test_modeling_llava_next_video.py b/tests/models/llava_next_video/test_modeling_llava_next_video.py index 33f3efa69a64..f506d9685bb1 100644 --- a/tests/models/llava_next_video/test_modeling_llava_next_video.py +++ b/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): diff --git a/tests/models/llava_onevision/test_modeling_llava_onevision.py b/tests/models/llava_onevision/test_modeling_llava_onevision.py index 6de0193f03a9..bf955f6e0816 100644 --- a/tests/models/llava_onevision/test_modeling_llava_onevision.py +++ b/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): diff --git a/tests/models/mistral3/test_modeling_mistral3.py b/tests/models/mistral3/test_modeling_mistral3.py index f33b3e0e1f8a..da6b21733c20 100644 --- a/tests/models/mistral3/test_modeling_mistral3.py +++ b/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 diff --git a/tests/models/mllama/test_modeling_mllama.py b/tests/models/mllama/test_modeling_mllama.py index dbcf88869deb..a898244254d0 100644 --- a/tests/models/mllama/test_modeling_mllama.py +++ b/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): diff --git a/tests/models/paligemma/test_modeling_paligemma.py b/tests/models/paligemma/test_modeling_paligemma.py index 6399608b29bc..6fda37e5effa 100644 --- a/tests/models/paligemma/test_modeling_paligemma.py +++ b/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 diff --git a/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py b/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py index 5a425b434e7d..374f3fc4ed27 100644 --- a/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py +++ b/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 diff --git a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py b/tests/models/qwen2_vl/test_modeling_qwen2_vl.py index 6027feac66fe..0ac8cd1bd385 100644 --- a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py +++ b/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 diff --git a/tests/models/video_llava/test_modeling_video_llava.py b/tests/models/video_llava/test_modeling_video_llava.py index 8d5995ae8a30..533e3c10a00a 100644 --- a/tests/models/video_llava/test_modeling_video_llava.py +++ b/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): diff --git a/tests/models/vipllava/test_modeling_vipllava.py b/tests/models/vipllava/test_modeling_vipllava.py index 5b06d8659145..24bb47deda74 100644 --- a/tests/models/vipllava/test_modeling_vipllava.py +++ b/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): diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index bc8f65891445..798fe3a15b43 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -4799,7 +4799,15 @@ 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: + + # 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), @@ -4816,9 +4824,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", ) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 85e77b834d3e..b23a4c088c50 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -19,7 +19,11 @@ import torch.nn as nn from transformers import PretrainedConfig -from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping +from transformers.conversion_mapping import ( + get_checkpoint_conversion_mapping, + get_model_conversion_mapping, + register_checkpoint_conversion_mapping, +) from transformers.core_model_loading import ( Chunk, Concatenate, @@ -406,7 +410,7 @@ def test_moe_and_qkv_conversion_reversed(self): def test_qkv_chunk_rope_permute_with_fp8_quantization(self): if is_triton_available(): - from transformers.integrations.finegrained_fp8 import Fp8Dequantize + from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize else: self.skipTest("Fine-grained FP8 integration tests require Triton to be installed.") n_heads = 2 @@ -472,6 +476,7 @@ def __init__(self): self, "quantization_config", SimpleNamespace(weight_block_size=bs) ), "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), + "get_quantize_ops": lambda self: Fp8Quantize(self), "pre_quantized": False, }, ) @@ -479,11 +484,11 @@ def __init__(self): weight_mapping = [ WeightConverter( - "model.layers.*.self_attn.qkv_proj.weight", + "self_attn.qkv_proj.weight", [ - "model.layers.*.self_attn.q_proj.weight", - "model.layers.*.self_attn.k_proj.weight", - "model.layers.*.self_attn.v_proj.weight", + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", ], operations=[Chunk(dim=0), PermuteForRope()], ) @@ -526,6 +531,138 @@ def __init__(self): ) torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) + def test_scoped_renaming_does_not_leak_to_sibling_keys(self): + """scope_prefix restricts a WeightRenaming to keys under that sub-prefix only. + + A "^"-anchored pattern must match the suffix after stripping the prefix, and + must not fire at all on keys that do not start with the scope prefix. + + Without scope_prefix, "^old_q" would rename *any* key beginning with "old_q" + at any nesting level — including root-level ones that belong to a different + part of the model. + """ + + class _Attn(nn.Module): + def __init__(self): + super().__init__() + self.q = DummyParamModule((1, 2)) + + class _Encoder(nn.Module): + def __init__(self): + super().__init__() + self.attn = _Attn() + + class _ScopedModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.encoder = _Encoder() + self.q = DummyParamModule((1, 2)) # root-level q — must not be touched + + model = _ScopedModel() + model.config = PretrainedConfig() + + enc_val = torch.tensor([[1.0, 2.0]]) + checkpoint = { + "encoder.attn.old_q.weight": enc_val.clone(), + "old_q.weight": torch.tensor([[9.0, 9.0]]), # outside scope + } + + scoped_rename = WeightRenaming("^old_q", "q") + scoped_rename.scope_prefix = "encoder.attn" + + loading_info, _ = convert_and_load_state_dict_in_model( + model, + checkpoint, + LoadStateDictConfig(weight_mapping=[scoped_rename]), + tp_plan=None, + ) + + # The root-level "old_q.weight" must be unmatched (unexpected), not silently + # loaded into "q.weight". + self.assertEqual(loading_info.unexpected_keys, {"old_q.weight"}) + self.assertEqual(loading_info.missing_keys, {"q.weight"}) + self.assertEqual(loading_info.mismatched_keys, set()) + self.assertEqual(loading_info.conversion_errors, {}) + + torch.testing.assert_close(model.encoder.attn.q.weight, enc_val) + # Root q.weight must still be its initialised zero value. + torch.testing.assert_close(model.q.weight, torch.zeros(1, 2)) + + def test_interleaved_renaming_and_converter_round_trip(self): + """A WeightRenaming preceding a WeightConverter in the list must fire in the + reverse (save) direction even after the converter has already set source_pattern. + + Forward [WeightRenaming, WeightConverter]: + "decoder.attn.qkv_proj.weight" + → WeightRenaming : "encoder.attn.qkv_proj.weight" + → WeightConverter: "encoder.attn.{q,k,v}_proj.weight" (source_pattern set) + + Reverse [rev(WeightConverter), rev(WeightRenaming)]: + "encoder.attn.{q,k,v}_proj.weight" + → rev(WeightConverter): "encoder.attn.qkv_proj.weight" (source_pattern set) + → rev(WeightRenaming) : "decoder.attn.qkv_proj.weight" ← must still run! + """ + + class _Attn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = DummyParamModule((2, 4)) + self.k_proj = DummyParamModule((2, 4)) + self.v_proj = DummyParamModule((2, 4)) + + class _Encoder(nn.Module): + def __init__(self): + super().__init__() + self.attn = _Attn() + + class _InterleavedModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.encoder = _Encoder() + + qkv = torch.arange(24, dtype=torch.float32).reshape(6, 4) + model = _InterleavedModel() + model.config = PretrainedConfig() + + # Checkpoint uses a "decoder" prefix and stores QKV packed together. + checkpoint = {"decoder.attn.qkv_proj.weight": qkv.clone()} + + weight_mapping = [ + WeightRenaming("^decoder", "encoder"), # step 1: fix prefix + WeightConverter( # step 2: unpack QKV (fires after rename) + "attn.qkv_proj.weight", + ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ), + ] + + loading_info, _ = convert_and_load_state_dict_in_model( + model, + checkpoint, + LoadStateDictConfig(weight_mapping=weight_mapping), + tp_plan=None, + ) + + self.assertEqual(loading_info.missing_keys, set()) + self.assertEqual(loading_info.unexpected_keys, set()) + self.assertEqual(loading_info.mismatched_keys, set()) + self.assertEqual(loading_info.conversion_errors, {}) + + q, k, v = torch.chunk(qkv, 3, dim=0) + torch.testing.assert_close(model.encoder.attn.q_proj.weight, q) + torch.testing.assert_close(model.encoder.attn.k_proj.weight, k) + torch.testing.assert_close(model.encoder.attn.v_proj.weight, v) + + # Round-trip: saving must reconstruct the original "decoder.*" checkpoint. + # This relies on rev(WeightRenaming) firing after rev(WeightConverter) has set + # source_pattern — if it were skipped the prefix would remain "encoder". + saved = revert_weight_conversion(model, model.state_dict()) + self.assertTrue(compare_state_dicts(saved, checkpoint)) + def test_ernie4_5_vl_moe_conversion(self): model = DummyRoot(add_extra_moe=True) model.config = PretrainedConfig() @@ -1020,6 +1157,31 @@ def test_can_add_prefix_submodule(self): for k, v in saved_state_dict.items(): self.assertTrue((v == model_state_dict[k]).all()) + def test_class_name_wins_over_model_type(self): + """Class-name registry entry takes priority over model_type for the same model.""" + register_checkpoint_conversion_mapping("_TstCls", [WeightRenaming(r"^cls_key", "cls_renamed")], overwrite=True) + register_checkpoint_conversion_mapping( + "_tst_mtype", [WeightRenaming(r"^type_key", "type_renamed")], overwrite=True + ) + + def make_mock(class_name): + m = type(class_name, (), {})() + m.config = SimpleNamespace(model_type="_tst_mtype") + m._named_pretrained_submodules = [("", m)] + return m + + # A module whose class name has a registry entry → class entry wins. + transforms = get_model_conversion_mapping(make_mock("_TstCls"), add_legacy=False) + patterns = [t.source_patterns for t in transforms] + self.assertIn(["^cls_key"], patterns) + self.assertNotIn(["^type_key"], patterns) + + # A module with no class entry falls through to the model_type entry. + transforms = get_model_conversion_mapping(make_mock("_TstOther"), add_legacy=False) + patterns = [t.source_patterns for t in transforms] + self.assertIn(["^type_key"], patterns) + self.assertNotIn(["^cls_key"], patterns) + if __name__ == "__main__": unittest.main() From 832d11f0a594d6249235f799431f2fc14ed8b3de Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 29 Apr 2026 15:12:56 +0000 Subject: [PATCH 09/15] remove class mapping from model mapping deduplication --- src/transformers/conversion_mapping.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 3d7b2bd8bd13..81da8d461ff2 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -706,16 +706,20 @@ def get_model_conversion_mapping( class_name = type(submodule).__name__ model_type = getattr(submodule.config, "model_type", None) - # Skip if this architecture was already processed via either lookup path. - if class_name in seen_identifiers or (model_type and model_type in seen_identifiers): + if class_name in seen_identifiers: continue - # Try class name first, then model_type. Track which path produced the hit so - # we know whether to block model_type for subsequent sub-modules (see below). + # Class name takes priority — a class-specific mapping bypasses the model_type + # deduplication check (e.g. LlavaModel nested inside LlavaForConditionalGeneration + # must still get its own scoped mapping even after "llava" is marked seen). conversions = get_checkpoint_conversion_mapping(class_name) found_via_class = conversions is not None - if not found_via_class and model_type is not None: - conversions = get_checkpoint_conversion_mapping(model_type) + + if not found_via_class: + if model_type and model_type in seen_identifiers: + continue + if model_type is not None: + conversions = get_checkpoint_conversion_mapping(model_type) if conversions is None: continue @@ -729,9 +733,9 @@ def get_model_conversion_mapping( seen_identifiers.add(class_name) # Only block model_type when the hit was via model_type. When the hit was via - # class name, sub-modules that share the same model_type but have no class-specific - # mapping of their own (e.g. DetrModel under DetrForSegmentation) must still be - # reachable so their base transforms are picked up and scoped automatically. + # class name, other sub-modules sharing the same model_type but without a + # class-specific mapping (e.g. DetrModel under DetrForSegmentation) must still + # be reachable so their base transforms are picked up and scoped. if not found_via_class and model_type: seen_identifiers.add(model_type) From 9da04fc06b9d2a0f0a5642bbbdc392db80eac615 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 29 Apr 2026 21:46:24 +0000 Subject: [PATCH 10/15] Fix conversion mappings --- src/transformers/conversion_mapping.py | 6 ++++-- tests/test_modeling_common.py | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 55526d0e4ff2..8266b4b5066d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -104,6 +104,7 @@ "FuyuModel": "LlavaModel", "MllamaModel": "LlavaModel", "Qwen2_5_VLModel": "Qwen2VLModel", + "MaskFormerDetrDecoder": "DetrModel", } @@ -167,13 +168,14 @@ def _build_checkpoint_conversion_mapping(): target_patterns="model.language_model", ), ], - "Qwen2VLModel": [WeightRenaming(source_patterns=r"^model.", target_patterns="")], + "Qwen2VLModel": [], "qwen2_vl": [ WeightRenaming(source_patterns=r"^visual", target_patterns="model.visual"), WeightRenaming( - source_patterns=r"(? Date: Tue, 5 May 2026 18:46:15 +0000 Subject: [PATCH 11/15] nit --- src/transformers/core_model_loading.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 3a8b767dd611..7041b5c2b42f 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -733,7 +733,7 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: lambda m: match_object.group(group_start + int(m.group(1))), replacement, ) - renamed_key = source_key.replace(match_object.group(0), replacement, 1) + renamed_key = key_to_match.replace(match_object.group(0), replacement, 1) if prefix_dot is not None: renamed_key = prefix_dot + renamed_key return renamed_key, source_pattern_that_matched From d96c84fa998e3f571c3bb04ede600debdf231a93 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 6 May 2026 19:48:47 +0000 Subject: [PATCH 12/15] fix after review --- src/transformers/conversion_mapping.py | 49 ++++--- src/transformers/core_model_loading.py | 50 ++++---- tests/utils/test_core_model_loading.py | 171 ++++++++++++++++++------- 3 files changed, 181 insertions(+), 89 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index ca321b3d47cb..0bde1810691b 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -14,6 +14,7 @@ from __future__ import annotations +from collections import defaultdict from copy import deepcopy from typing import TYPE_CHECKING @@ -78,7 +79,6 @@ "siglip_text_model": "clip_text_model", "siglip2_text_model": "clip_text_model", "xclip_text_model": "clip_text_model", - "shield_gemma2": "llava", "paligemma": "llava", "aya_vision": "llava", "got_ocr2": "llava", @@ -359,6 +359,7 @@ def _build_checkpoint_conversion_mapping(): ), ], "colqwen2": [PrefixChange(prefix_to_remove="model", model_prefix="vlm")], + "shieldgemma2": [PrefixChange(prefix_to_add="model", model_prefix="model")], "timm_wrapper": [PrefixChange(prefix_to_add="timm_model")], "pi0": [ WeightRenaming(source_patterns=r"state_proj", target_patterns="embed_action_time.state_proj"), @@ -851,11 +852,11 @@ def extract_weight_conversions_for_model( unmodified; the caller sets ``scope_prefix`` on each transform for sub-module isolation. """ class_name = type(model).__name__ - model_type = getattr(model.config, "model_type", None) + model_type = model.config.model_type # Class name takes priority — allows ForXxx-specific overrides conversions = get_checkpoint_conversion_mapping(class_name) - if conversions is None and model_type is not None: + if conversions is None and model_type: conversions = get_checkpoint_conversion_mapping(model_type) return conversions @@ -867,14 +868,14 @@ def get_model_conversion_mapping( add_legacy: bool = True, ) -> list[WeightTransform]: """ - Collect the ordered list of weight transforms for ``model`` (used during + Collect the ordered list of weight transforms for `model` (used during loading and, when reversed, during saving). - Each ``PreTrainedModel`` sub-module is looked up by class name then - ``model_type``. Root transforms are applied globally; sub-module transforms - have their ``scope_prefix`` set so they only match keys under that prefix. After any - sub-module is processed, both its class name and ``model_type`` are marked - seen to prevent ``XForY`` / ``XModel`` pairs from applying the same mapping + Each `PreTrainedModel` sub-module is looked up by class name then + `model_type`. Root transforms are applied globally; sub-module transforms + have their `scope_prefix` set so they only match keys under that prefix. After any + sub-module is processed, both its class name and `model_type` are marked + seen to prevent `XForY` / `XModel` pairs from applying the same mapping twice via different lookup paths. """ # Lazy import to avoid circular import issues @@ -886,18 +887,21 @@ def get_model_conversion_mapping( if key_mapping is not None: weight_conversions = [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()] - seen_identifiers: set[str] = set() + # Maps each identifier (class name or model_type) to the module paths that have + # already claimed it. A later module is skipped only when one of those paths is + # an ancestor of the current module path — siblings are never ancestors of each + # other, so two sibling sub-models with the same model_type both get their own + # scoped transforms. A child is an ancestor of everything nested under it, which + # prevents a parent's transforms from being duplicated with a scoped copy for the child. + seen_identifiers: defaultdict[str, list[str]] = defaultdict(list) - named_pretrained = getattr(model, "_named_pretrained_submodules", None) - if named_pretrained is None: - from .modeling_utils import PreTrainedModel - - named_pretrained = [(name, m) for name, m in model.named_modules() if isinstance(m, PreTrainedModel)] + named_pretrained = model._named_pretrained_submodules for module_name, submodule in named_pretrained: class_name = type(submodule).__name__ - model_type = getattr(submodule.config, "model_type", None) + model_type = submodule.config.model_type - if class_name in seen_identifiers: + # Skip if an ancestor already claimed this class (its unscoped transforms already cover this subtree). + if any(seen == "" or module_name.startswith(seen + ".") for seen in seen_identifiers[class_name]): continue # Class name takes priority — a class-specific mapping bypasses the model_type @@ -907,7 +911,10 @@ def get_model_conversion_mapping( found_via_class = conversions is not None if not found_via_class: - if model_type and model_type in seen_identifiers: + # Same ancestor check as above, but via model_type for modules without a class-specific mapping. + if model_type and any( + seen == "" or module_name.startswith(seen + ".") for seen in seen_identifiers[model_type] + ): continue if model_type is not None: conversions = get_checkpoint_conversion_mapping(model_type) @@ -922,13 +929,13 @@ def get_model_conversion_mapping( transform.scope_prefix = module_name weight_conversions.extend(conversions) - seen_identifiers.add(class_name) - # Only block model_type when the hit was via model_type. When the hit was via + seen_identifiers[class_name].append(module_name) + # Only record model_type when the hit was via model_type. When the hit was via # class name, other sub-modules sharing the same model_type but without a # class-specific mapping (e.g. DetrModel under DetrForSegmentation) must still # be reachable so their base transforms are picked up and scoped. if not found_via_class and model_type: - seen_identifiers.add(model_type) + seen_identifiers[model_type].append(module_name) if add_legacy: weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7041b5c2b42f..9509b1032953 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -112,7 +112,7 @@ def convert(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]: class Chunk(ConversionOps): - """Split a tensor along ``dim`` into equally sized chunks.""" + """Split a tensor along `dim` into equally sized chunks.""" def __init__(self, dim: int = 0): self.dim = dim @@ -156,8 +156,8 @@ def convert( target_pattern = self.get_target_pattern(target_patterns) all_tensors = [] # Very important to keep the relative order of the source patterns here, so we iterate over them not the - # input directly as it's unordered! Skip patterns that prior ops in the chain (e.g. ``Fp8Dequantize``) - # have already consumed and dropped from ``input_dict``. + # input directly as it's unordered! Skip patterns that prior ops in the chain (e.g. `Fp8Dequantize`) + # have already consumed and dropped from `input_dict`. for source_pattern in source_patterns: if source_pattern not in input_dict: continue @@ -613,7 +613,7 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list # Flag to notice if the Transform was used self._was_used = False # Optional prefix scope: when set, this transform only applies to keys starting with - # ``scope_prefix + "."``, stripping / re-attaching the prefix around the pattern match. + # `scope_prefix + "."`, stripping / re-attaching the prefix around the pattern match. self.scope_prefix: str | None = None # We need to process a few exceptions here when instantiating the reverse mapping (i.e. the targets become @@ -682,11 +682,11 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu def _scoped_match(self, source_key: str) -> tuple[str | None, str, re.Match[str]] | None: """ - Apply ``scope_prefix`` stripping (if any), then match ``compiled_sources`` against the suffix. + Apply `scope_prefix` stripping (if any), then match `compiled_sources` against the suffix. - Returns ``(prefix_dot, key_to_match, match_object)`` when a branch matches, where ``prefix_dot`` is ``None`` - if ``scope_prefix`` is unset, else ``f"{scope_prefix}."``. Returns ``None`` when out of scope or unmatched. - Does not set ``_was_used``. + Returns `(prefix_dot, key_to_match, match_object)` when a branch matches, where `prefix_dot` is `None` + if `scope_prefix` is unset, else `f"{scope_prefix}."`. Returns `None` when out of scope or unmatched. + Does not set `_was_used`. """ prefix_dot = None key_to_match = source_key @@ -1011,17 +1011,17 @@ def _job(): def dot_natural_key(s: str): - """Sort key for state-dict names: split on ``"."`` and sort digits numerically + """Sort key for state-dict names: split on `"."` and sort digits numerically and strings alphabetically. We emit a tuple at each point to sort ints first and strings second to avoid int-string comparison failures. """ - result = [] - for p in s.split("."): - if p.isdigit(): - result.append((0, int(p))) + parts = [] + for part in s.split("."): + if part.isdigit(): + parts.append((0, int(part))) else: - result.append((1, p)) - return result + parts.append((1, part)) + return parts @contextmanager @@ -1149,14 +1149,14 @@ def rename_source_key( meta_state_dict: dict | None = None, ) -> tuple[str, str | None]: """ - Rename a source key according to ``weight_transforms``, also handling the base model prefix. + Rename a source key according to `weight_transforms`, also handling the base model prefix. - Transforms are applied in list order, interleaving ``WeightRenaming`` and ``WeightConverter`` + Transforms are applied in list order, interleaving `WeightRenaming` and `WeightConverter` instances as they appear. The same list, reversed and with each transform individually inverted, is used on the save path, so relative ordering is preserved in both directions. - At most one ``WeightConverter`` fires per key; subsequent converters are skipped. - ``WeightRenaming`` always runs, even after a converter has already fired. + At most one `WeightConverter` fires per key; subsequent converters are skipped. + `WeightRenaming` always runs, even after a converter has already fired. Example (root rename followed by a scoped sub-model converter):: @@ -1178,13 +1178,17 @@ def rename_source_key( for transform in weight_transforms: if isinstance(transform, WeightConverter): + # At most one WeightConverter may claim a key. The returned source_pattern tells + # the caller which converter to feed this tensor to; a second match would redirect + # the tensor to the wrong converter, leaving the first one incomplete. + # TODO: `continue` here can hide misconfigured mapping (two converters with overlapping patterns). Raise an error instead. if source_pattern is not None: - # Already matched a converter; skip subsequent converters. continue - renamed_key, sp = transform.rename_source_key(renamed_key) - if sp is not None: - source_pattern = sp + renamed_key, matched_pattern = transform.rename_source_key(renamed_key) + if matched_pattern is not None: + source_pattern = matched_pattern else: + # WeightRenaming always runs. Multiple renames can safely chain before (and after) a converter fires. renamed_key, _ = transform.rename_source_key(renamed_key) # check if we need to add or remove prefix if necessary (only during loading, not saving) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index b23a4c088c50..d9c700a9e588 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -221,6 +221,10 @@ def __init__(self, add_extra_moe=False, with_mlp=True): self.model = DummyTopModel(add_extra_moe) if with_mlp: self.mlp = DummyMLP() + self.config = PretrainedConfig() + # Mirror what PreTrainedModel.post_init() does so that + # get_model_conversion_mapping() can be called on DummyRoot. + self._named_pretrained_submodules = [("", self)] class TestConvertAndLoadStateDict(unittest.TestCase): @@ -531,46 +535,42 @@ def __init__(self): ) torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) - def test_scoped_renaming_does_not_leak_to_sibling_keys(self): - """scope_prefix restricts a WeightRenaming to keys under that sub-prefix only. + def test_scoped_renaming_does_not_leak_to_sibling_or_parent(self): + """scope_prefix gates a WeightRenaming to keys under one submodel only — + neither the sibling submodel nor the parent's own keys must be affected. - A "^"-anchored pattern must match the suffix after stripping the prefix, and - must not fire at all on keys that do not start with the scope prefix. - - Without scope_prefix, "^old_q" would rename *any* key beginning with "old_q" - at any nesting level — including root-level ones that belong to a different - part of the model. + This mirrors the real use-case: a composite model (e.g. vision-language) has + two top-level submodels and a root-level weight, each with an ``old_q`` key. + The rename is registered only for ``vision_model``; the sibling + ``language_model.old_q`` and the root-level ``old_q`` must remain unmatched. """ - class _Attn(nn.Module): + class _Submodel(nn.Module): def __init__(self): super().__init__() self.q = DummyParamModule((1, 2)) - class _Encoder(nn.Module): - def __init__(self): - super().__init__() - self.attn = _Attn() - - class _ScopedModel(nn.Module): + class _CompositeModel(nn.Module): base_model_prefix = "" def __init__(self): super().__init__() - self.encoder = _Encoder() - self.q = DummyParamModule((1, 2)) # root-level q — must not be touched + self.vision_model = _Submodel() + self.language_model = _Submodel() + self.q = DummyParamModule((1, 2)) # root-level weight with the same name - model = _ScopedModel() + model = _CompositeModel() model.config = PretrainedConfig() - enc_val = torch.tensor([[1.0, 2.0]]) + vision_val = torch.tensor([[1.0, 2.0]]) checkpoint = { - "encoder.attn.old_q.weight": enc_val.clone(), - "old_q.weight": torch.tensor([[9.0, 9.0]]), # outside scope + "vision_model.old_q.weight": vision_val.clone(), + "language_model.old_q.weight": torch.tensor([[9.0, 9.0]]), + "old_q.weight": torch.tensor([[7.0, 7.0]]), # root-level, must not be renamed } scoped_rename = WeightRenaming("^old_q", "q") - scoped_rename.scope_prefix = "encoder.attn" + scoped_rename.scope_prefix = "vision_model" loading_info, _ = convert_and_load_state_dict_in_model( model, @@ -579,30 +579,33 @@ def __init__(self): tp_plan=None, ) - # The root-level "old_q.weight" must be unmatched (unexpected), not silently - # loaded into "q.weight". - self.assertEqual(loading_info.unexpected_keys, {"old_q.weight"}) - self.assertEqual(loading_info.missing_keys, {"q.weight"}) + # Sibling and parent keys must be unmatched. + self.assertEqual(loading_info.unexpected_keys, {"language_model.old_q.weight", "old_q.weight"}) + self.assertEqual(loading_info.missing_keys, {"language_model.q.weight", "q.weight"}) self.assertEqual(loading_info.mismatched_keys, set()) self.assertEqual(loading_info.conversion_errors, {}) - torch.testing.assert_close(model.encoder.attn.q.weight, enc_val) - # Root q.weight must still be its initialised zero value. + torch.testing.assert_close(model.vision_model.q.weight, vision_val) + torch.testing.assert_close(model.language_model.q.weight, torch.zeros(1, 2)) torch.testing.assert_close(model.q.weight, torch.zeros(1, 2)) def test_interleaved_renaming_and_converter_round_trip(self): - """A WeightRenaming preceding a WeightConverter in the list must fire in the - reverse (save) direction even after the converter has already set source_pattern. - - Forward [WeightRenaming, WeightConverter]: - "decoder.attn.qkv_proj.weight" - → WeightRenaming : "encoder.attn.qkv_proj.weight" - → WeightConverter: "encoder.attn.{q,k,v}_proj.weight" (source_pattern set) - - Reverse [rev(WeightConverter), rev(WeightRenaming)]: - "encoder.attn.{q,k,v}_proj.weight" - → rev(WeightConverter): "encoder.attn.qkv_proj.weight" (source_pattern set) - → rev(WeightRenaming) : "decoder.attn.qkv_proj.weight" ← must still run! + """A WeightRenaming preceding a WeightConverter must also fire on the save path + even after the converter has already set source_pattern. + + Loading (checkpoint key → model key): + 1. WeightRenaming (source "^decoder", target "encoder"): + "decoder.attn.qkv_proj.weight" → "encoder.attn.qkv_proj.weight" + 2. WeightConverter (source "^attn.qkv_proj.weight", targets "attn.{q,k,v}_proj.weight", scope "encoder"): + strips "encoder.", matches source, unpacks QKV, re-attaches "encoder." + "encoder.attn.qkv_proj.weight" → "encoder.attn.{q,k,v}_proj.weight" + + Saving (model key → checkpoint key, transforms applied in reverse): + 1. rev(WeightConverter) (source "attn.{q,k,v}_proj.weight", target "^attn.qkv_proj.weight", scope "encoder"): + strips "encoder.", matches source, repacks QKV, re-attaches "encoder." + "encoder.attn.{q,k,v}_proj.weight" → "encoder.attn.qkv_proj.weight" + 2. rev(WeightRenaming) (source "^encoder", target "decoder"): + "encoder.attn.qkv_proj.weight" → "decoder.attn.qkv_proj.weight" """ class _Attn(nn.Module): @@ -631,13 +634,17 @@ def __init__(self): # Checkpoint uses a "decoder" prefix and stores QKV packed together. checkpoint = {"decoder.attn.qkv_proj.weight": qkv.clone()} + qkv_converter = WeightConverter( + "^attn.qkv_proj.weight", + ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + # scope_prefix mirrors what get_model_conversion_mapping sets for a submodel + qkv_converter.scope_prefix = "encoder" + weight_mapping = [ WeightRenaming("^decoder", "encoder"), # step 1: fix prefix - WeightConverter( # step 2: unpack QKV (fires after rename) - "attn.qkv_proj.weight", - ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ), + qkv_converter, # step 2: unpack QKV (fires after rename, scoped to encoder) ] loading_info, _ = convert_and_load_state_dict_in_model( @@ -1182,6 +1189,80 @@ def make_mock(class_name): self.assertIn(["^type_key"], patterns) self.assertNotIn(["^cls_key"], patterns) + def test_sibling_submodels_same_model_type_both_get_transforms(self): + """Two sibling sub-models sharing a model_type must each get their own scoped transforms. + + Old flat-set deduplication would mark the model_type seen after the first + sibling and silently skip the second one. The ancestor-based check must + allow both because neither sibling is an ancestor of the other. + """ + register_checkpoint_conversion_mapping( + "_tst_shared_type", [WeightRenaming(r"^w", "renamed_w")], overwrite=True + ) + + def make_root_with_siblings(cls_a, cls_b): + """Root model with two children of different classes but the same model_type.""" + child_a = type(cls_a, (), {})() + child_a.config = SimpleNamespace(model_type="_tst_shared_type") + + child_b = type(cls_b, (), {})() + child_b.config = SimpleNamespace(model_type="_tst_shared_type") + + root = type("_TstRoot", (), {})() + root.config = SimpleNamespace(model_type="_tst_root_only") + root._named_pretrained_submodules = [("encoder", child_a), ("decoder", child_b)] + return root + + transforms = get_model_conversion_mapping( + make_root_with_siblings("_TstEncCls", "_TstDecCls"), add_legacy=False + ) + scope_prefixes = [t.scope_prefix for t in transforms] + + # Both siblings must be represented with their own scoped transforms. + self.assertIn("encoder", scope_prefixes) + self.assertIn("decoder", scope_prefixes) + + def test_sibling_submodels_same_class_both_get_transforms(self): + """Two sibling sub-models of the *same* class must each get their own scoped transforms.""" + register_checkpoint_conversion_mapping("_TstSharedCls", [WeightRenaming(r"^w", "renamed_w")], overwrite=True) + + child_a = type("_TstSharedCls", (), {})() + child_a.config = SimpleNamespace(model_type="_tst_shared_cls_mtype") + + child_b = type("_TstSharedCls", (), {})() + child_b.config = SimpleNamespace(model_type="_tst_shared_cls_mtype") + + root = type("_TstRootSharedCls", (), {})() + root.config = SimpleNamespace(model_type="_tst_root_only2") + root._named_pretrained_submodules = [("encoder", child_a), ("decoder", child_b)] + + transforms = get_model_conversion_mapping(root, add_legacy=False) + scope_prefixes = [t.scope_prefix for t in transforms] + + self.assertIn("encoder", scope_prefixes) + self.assertIn("decoder", scope_prefixes) + + def test_child_with_same_model_type_as_root_is_skipped(self): + """When the root model claims a model_type unscoped, a nested child with the + same model_type must NOT produce a second (incorrectly scoped) copy of those + transforms — the root's unscoped transforms already cover all keys.""" + register_checkpoint_conversion_mapping( + "_tst_root_child_shared", [WeightRenaming(r"^w", "renamed_w")], overwrite=True + ) + + child = type("_TstChildSame", (), {})() + child.config = SimpleNamespace(model_type="_tst_root_child_shared") + + root = type("_TstRootSame", (), {})() + root.config = SimpleNamespace(model_type="_tst_root_child_shared") + # Root ("") appears before child ("submodel") — mirrors DFS order. + root._named_pretrained_submodules = [("", root), ("submodel", child)] + + transforms = get_model_conversion_mapping(root, add_legacy=False) + # Only one unscoped transform (from the root); child must be suppressed. + self.assertEqual(len(transforms), 1) + self.assertIsNone(transforms[0].scope_prefix) + if __name__ == "__main__": unittest.main() From 2a0b56b1b5f8fdc0a9d356f8ca3c8be17833ef97 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 6 May 2026 19:55:46 +0000 Subject: [PATCH 13/15] fix bacticks --- src/transformers/conversion_mapping.py | 54 +++++++++++++------------- src/transformers/core_model_loading.py | 2 +- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 0bde1810691b..c494582dd179 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -114,29 +114,29 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"layer\.", target_patterns="layers."), ], "deepseek_v4": [ - # Upstream V4-Flash checkpoint uses a flatter V3-style namespace: ``attn`` / - # ``ffn`` instead of ``self_attn`` / ``mlp``, ``attn_norm`` / ``ffn_norm`` - # instead of ``input_layernorm`` / ``post_attention_layernorm``, ``hc_attn_*`` - # / ``hc_ffn_*`` for the Hyper-Connection params (wrapped here as - # ``attn_hc`` / ``ffn_hc`` submodules), ``embed`` / ``head`` / bare ``norm`` - # for the model head, ``hc_head_*`` for the final HC collapse, and indexer - # weights nested under ``attn.indexer.compressor.*`` upstream but flattened + # Upstream V4-Flash checkpoint uses a flatter V3-style namespace: `attn` / + # `ffn` instead of `self_attn` / `mlp`, `attn_norm` / `ffn_norm` + # instead of `input_layernorm` / `post_attention_layernorm`, `hc_attn_*` + # / `hc_ffn_*` for the Hyper-Connection params (wrapped here as + # `attn_hc` / `ffn_hc` submodules), `embed` / `head` / bare `norm` + # for the model head, `hc_head_*` for the final HC collapse, and indexer + # weights nested under `attn.indexer.compressor.*` upstream but flattened # onto the Indexer module here. # - # All targets stay in the bare base-model namespace (no ``model.`` prefix). - # ``convert_and_load_state_dict_in_model`` consults + # All targets stay in the bare base-model namespace (no `model.` prefix). + # `convert_and_load_state_dict_in_model` consults # :attr:`DeepseekV4PreTrainedModel.base_model_prefix = "model"` and adds / - # strips the ``model.`` prefix automatically based on whether the loader + # strips the `model.` prefix automatically based on whether the loader # target is the base model or a head model. # # Ordering matters for save round-tripping: :func:`revert_weight_conversion` # reverses the order *and* each transform, so a structural prefix-only rule # placed before a specific in-prefix rename would steal the reverse match - # and emit ``layers.X.attn.sinks`` instead of ``layers.X.attn.attn_sink``. + # and emit `layers.X.attn.sinks` instead of `layers.X.attn.attn_sink`. # We split into two passes: structural prefix renames first (so they apply # last on save / first on load), then specific in-prefix renames that - # operate on the already-prefixed keys. FP8 ``.scale`` → ``.weight_scale_inv`` - # rename lives in the FP8 quantizer's ``update_weight_conversions`` (only + # operate on the already-prefixed keys. FP8 `.scale` → `.weight_scale_inv` + # rename lives in the FP8 quantizer's `update_weight_conversions` (only # active under FP8 dequant), so the V4 static mapping below stays free of # FP8-only rules. # ---- Pass 1: top-level + structural prefix renames ---- @@ -176,7 +176,7 @@ def _build_checkpoint_conversion_mapping(): ), # ---- Pass 2: in-prefix specific renames (operate on already-prefixed keys) ---- # These can safely run after the structural prefix renames because their - # source patterns include the ``layers.X.self_attn.`` / ``layers.X.mlp.`` + # source patterns include the `layers.X.self_attn.` / `layers.X.mlp.` # prefix. On reverse the order flips so these undo first, restoring the # specific upstream names *before* the structural rules strip the prefix. WeightRenaming( @@ -208,10 +208,10 @@ def _build_checkpoint_conversion_mapping(): target_patterns=r"layers.\1.self_attn.compressor.position_bias", ), # Attention / compressor / indexer leaf weights: upstream uses paper notation - # (``wq_a`` / ``wq_b`` / ``wkv`` / ``wo_a`` / ``wo_b`` / ``wgate``); we - # rename to the standard transformers ``*_proj`` form. Compressor / Indexer - # ``wkv`` / ``wgate`` are caught by the same patterns since they sit under - # ``self_attn.`` after the Pass 1 prefix rewrite. + # (`wq_a` / `wq_b` / `wkv` / `wo_a` / `wo_b` / `wgate`); we + # rename to the standard transformers `*_proj` form. Compressor / Indexer + # `wkv` / `wgate` are caught by the same patterns since they sit under + # `self_attn.` after the Pass 1 prefix rewrite. WeightRenaming( source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wq_a\.", target_patterns=r"layers.\1.self_attn.\2.q_a_proj.", @@ -264,8 +264,8 @@ def _build_checkpoint_conversion_mapping(): source_patterns=r"^layers\.(\d+)\.self_attn\.q_norm\.", target_patterns=r"layers.\1.self_attn.q_a_norm.", ), - # Aux-loss-free routing bias: upstream ships ``gate.bias`` (V3 convention); - # we register it as ``e_score_correction_bias`` (cross-model standard name). + # Aux-loss-free routing bias: upstream ships `gate.bias` (V3 convention); + # we register it as `e_score_correction_bias` (cross-model standard name). WeightRenaming( source_patterns=r"^layers\.(\d+)\.mlp\.gate\.bias$", target_patterns=r"layers.\1.mlp.gate.e_score_correction_bias", @@ -826,10 +826,10 @@ def register_checkpoint_conversion_mapping( """ Register a conversion mapping for a model type string or a class name. - Class names take priority over ``model_type`` strings during lookup (see - :func:`extract_weight_conversions_for_model`), making it possible to define + Class names take priority over `model_type` strings during lookup (see + `extract_weight_conversions_for_model`), making it possible to define task-head-specific or class-specific conversions that differ from the shared - ``model_type`` baseline. + `model_type` baseline. """ global _checkpoint_conversion_mapping_cache if _checkpoint_conversion_mapping_cache is None: @@ -845,11 +845,11 @@ def extract_weight_conversions_for_model( model: PreTrainedModel, ) -> list[WeightTransform] | None: """ - Return the registered conversion list for ``model``, or ``None`` if none exists. + Return the registered conversion list for `model`, or `None` if none exists. Looks up by class name first (enables task-head-specific overrides), then - falls back to ``model.config.model_type``. Transforms are returned - unmodified; the caller sets ``scope_prefix`` on each transform for sub-module isolation. + falls back to `model.config.model_type`. Transforms are returned + unmodified; the caller sets `scope_prefix` on each transform for sub-module isolation. """ class_name = type(model).__name__ model_type = model.config.model_type @@ -941,7 +941,7 @@ def get_model_conversion_mapping( weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) # Let the quantizer rewrite / augment the conversion pipeline. This is where the - # FP8 dequantizer (when ``dequantize=True``) prepends a ``Fp8Dequantize`` op to + # FP8 dequantizer (when `dequantize=True`) prepends a `Fp8Dequantize` op to # every existing converter so that per-block scales are applied *before* any # expert-merge / concat ops flatten the per-expert structure away. if hf_quantizer is not None: diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 9509b1032953..ada64cf75e66 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -220,7 +220,7 @@ def reverse_op(self) -> ConversionOps: class SplitModulelist(ConversionOps): - """Inverse of :class:`MergeModulelist` using explicit split sizes per group.""" + """Inverse of `MergeModulelist` using explicit split sizes per group.""" def __init__(self, dim: int = 0): self.dim = dim From dee8826f1472862ec8e1c2551ed7635506e16ef7 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Wed, 6 May 2026 22:24:47 +0000 Subject: [PATCH 14/15] update docs + BC for rename_source_key --- docs/source/en/weightconverter.md | 264 +++++++++++++++----- src/transformers/conversion_mapping.py | 8 +- src/transformers/core_model_loading.py | 92 ++++--- src/transformers/integrations/accelerate.py | 3 +- src/transformers/integrations/deepspeed.py | 8 +- 5 files changed, 279 insertions(+), 96 deletions(-) diff --git a/docs/source/en/weightconverter.md b/docs/source/en/weightconverter.md index 6ad9e426bffa..a9269c8492bf 100644 --- a/docs/source/en/weightconverter.md +++ b/docs/source/en/weightconverter.md @@ -21,15 +21,16 @@ Checkpoints are often serialized in a format that does not match what a model ex 1. **Fused weights**: Checkpoints store separate `gate_proj` and `up_proj` weights, but the model uses a fused `gate_up_proj` for efficiency. 2. **MoE expert consolidation**: Individual expert weights (`experts.0.weight`, `experts.1.weight`, ...) need to be stacked into a single 3D tensor. 3. **Legacy naming**: Old checkpoints use different naming conventions (e.g., `LayerNorm.gamma` vs `LayerNorm.weight`). -4. **Quantization**: Weights may be stored in quantized formats that need deserialization. +4. **Composite models**: A vision-language model contains two `PreTrainedModel` sub-modules, each with its own checkpoint convention. +5. **Quantization**: Weights may be stored in quantized formats that need deserialization. -Dynamic weight loading addresses this by applying scheduled, reversible operations to checkpoint tensors as they are loaded. Transformers makes this available through [`WeightConverter`], which maps one or more source keys to target keys by running a list of composable conversion operations. This approach adapts to new weight layouts, and supports loading quantized mixture-of-experts (MoEs) or enabling tensor parallelism and MoEs. +Dynamic weight loading addresses this by applying scheduled, reversible operations to checkpoint tensors as they are loaded. Transformers exposes this through [`WeightConverter`] and [`WeightRenaming`], which describe how one or more checkpoint keys map to one or more model parameters and which composable [`ConversionOps`] should run on the matched tensors. This approach adapts to new weight layouts, supports quantized mixture-of-experts (MoEs), and integrates with tensor parallelism. -This guide demonstrates how to use the [`WeightConverter`] to convert tensors. Your [`WeightConverter`] should be added inside [_build_checkpoint_conversion_mapping()](https://github.com/huggingface/transformers/blob/4c9fde2a2a3aece0bcf1be93f696e88297da9397/src/transformers/conversion_mapping.py#L34) in the [conversion_mapping.py](https://github.com/huggingface/transformers/blob/main/src/transformers/conversion_mapping.py) file. +This guide demonstrates how to use [`WeightConverter`] to convert tensors. Conversion mappings live in [`conversion_mapping.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/conversion_mapping.py); a registered mapping is keyed by either a `model_type` string (e.g. `"mixtral"`) or a class name (e.g. `"LlavaModel"`). ## Full loading pipeline -All models go through the dynamic weight loading system. Conversion mapping is an **optional step within that system** that only activates when the model has entries in `_MODEL_TO_CONVERSION_PATTERN`. +All models go through the dynamic weight loading system. Conversion mapping is an **optional step within that system** that only activates when the model has entries registered for its class name or `model_type`. ``` Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model() @@ -50,7 +51,7 @@ Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model() | Step | When it activates | |------|-------------------| | Dynamic loading | Always, for all models | -| Conversion mapping | Only when `model_type` is in `_MODEL_TO_CONVERSION_PATTERN` | +| Conversion mapping | Only when the model's class or `model_type` is registered in `_MODEL_TO_CONVERSION_PATTERN` | | TP sharding | Only when `tp_plan="auto"` and model has `base_model_tp_plan` | | Dequantization/deserialization | Only when loading a pre-quantized checkpoint | | Quantization | Only when a quantization config is provided and weights are not pre-quantized | @@ -60,7 +61,7 @@ Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model() For most dense models, the checkpoint format matches the model format directly, so no conversion mapping is needed. Some models may still require renaming (e.g., legacy naming conventions). TP sharding still applies when enabled. ``` -Checkpoint: Model: +Checkpoint: Model: model.layers.0.self_attn.q_proj.weight → model.layers.0.self_attn.q_proj.weight model.layers.0.self_attn.k_proj.weight → model.layers.0.self_attn.k_proj.weight model.layers.0.mlp.gate_proj.weight → model.layers.0.mlp.gate_proj.weight @@ -89,48 +90,51 @@ experts.0.w3.weight │ + Concatenate experts.1.w3.weight ─┘ ``` +### Composite models (e.g., vision-language) + +A `PreTrainedModel` may contain other `PreTrainedModel` sub-modules. Each sub-model can have its own conversion mapping, registered against either its class name or its `model_type`. When the parent model is loaded, [`get_model_conversion_mapping`] walks the sub-models in depth-first order, collects their mappings, and **automatically scopes** each transform to the path of its sub-module via `scope_prefix`. + +``` +Composite model: Per-submodel mappings (auto-scoped): +LlavaForConditionalGeneration +├── vision_model: SiglipVisionModel → SiglipVisionModel mapping (scope="vision_model") +└── language_model: LlamaForCausalLM → LlamaForCausalLM mapping (scope="language_model") +``` + +`scope_prefix` is the dotted path of the sub-module (`"vision_model"`, `"language_model.model"`, etc.). A scoped transform only fires on keys that start with `f"{scope_prefix}."`; the prefix is stripped before pattern matching and re-attached after substitution, so each sub-model's mapping is written *relative to the sub-model*, exactly as if it were the root. + ## Architecture -The system is built around several key components defined in `src/transformers/core_model_loading.py`: +The system is built around several key components defined in [`core_model_loading.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/core_model_loading.py): **Phase 1 — Per-key processing** (iterates over checkpoint keys): -1. **Rename key** via `WeightRenaming` (e.g. `block_sparse_moe` -> `mlp`) -2. **Match pattern** via `WeightConverter` (e.g. `experts.*.w1.weight`) -3. **Shard (TP) and send to device** asynchronously via `ThreadPoolExecutor` -4. **Collect** tensors with the same `source_pattern` together (e.g. all MoE expert weights, gate + up projections) +1. **Walk the transform list** once. Every [`WeightRenaming`] that matches fires (e.g. `block_sparse_moe` → `mlp`), and **at most one** [`WeightConverter`] may claim the key (e.g. `experts.*.w1.weight`). +2. **Shard (TP) and send to device** asynchronously via `ThreadPoolExecutor`. +3. **Collect** tensors with the same `source_pattern` together (e.g. all MoE expert weights, gate + up projections). **Phase 2 — Per-mapping processing** (iterates over collected mappings): -1. **Dequantize/deserialize** (pre-quantized checkpoints only) -2. **Apply `ConversionOps` chain**: `Chunk`, `Concatenate`, `MergeModulelist`, `Transpose`, etc. -3. **Quantize** on-the-fly (if not pre-quantized) -4. **Set parameter** on model +1. **Dequantize/deserialize** (pre-quantized checkpoints only). +2. **Apply [`ConversionOps`] chain**: `Chunk`, `Concatenate`, `MergeModulelist`, `Transpose`, etc. +3. **Quantize** on-the-fly (if not pre-quantized). +4. **Set parameter** on model. ### WeightTransform The base class that handles pattern matching and tensor collection: -- **Pattern compilation**: Converts glob-style patterns (`*.weight`) to regex. -- **Key renaming**: `rename_source_key()` transforms checkpoint keys to model keys. -- **Tensor collection**: `add_tensor()` gathers related tensors for batch processing. -- **Reversibility**: `reverse_transform()` creates the inverse operation for saving. - -```python -@dataclass(slots=True) -class WeightTransform: - source_patterns: str | list[str] # Checkpoint key patterns - target_patterns: str | list[str] # Model key patterns - compiled_sources: re.Pattern # Compiled regex for matching - distributed_operation: TensorParallelLayer | None - quantization_operation: ConversionOps | None - collected_tensors: dict[str, list[Future]] # Gathered tensors - layer_targets: dict[str, set[str]] # Target key tracking -``` +- **Pattern compilation**: Source patterns are full regular expressions matched with `re.search()`. The `*` wildcard matches any indexable component and groups all matches together for batch operations. +- **Capturing groups & backreferences**: Capturing groups in source patterns can be referenced as `\1`, `\2`, ... in target patterns to preserve substrings (e.g. layer indices) across the rename. +- **Scoping**: `scope_prefix` (set automatically per sub-model by [`get_model_conversion_mapping`]) restricts the transform to keys under that path. The prefix is stripped before matching and re-attached after substitution. +- **Key renaming**: `rename_source_key()` applies the regex (with scope handling) and returns `(renamed_key, source_pattern)` so the loader knows which converter, if any, claimed the key. +- **Tensor collection**: `add_tensor()` accumulates resolved tensors (or `Future`s) under their `source_pattern` so that all tensors needed by a single conversion (e.g. all MoE expert weights) are batched together before the operation chain runs. +- **Reversibility**: `reverse_transform()` swaps source ↔ target patterns and inverts each operation, so the same list reversed drives saving. +- **Tracking**: `was_used()` reports whether the transform actually matched any key during loading; this is required so that non-bijective renames (e.g. [`PrefixChange`] adding a prefix that may already be present) can be re-applied symmetrically on save. ### WeightRenaming -[`WeightRenaming`] is a specialized [`WeightTransform`] for simple 1:1 key renaming without tensor operations: +[`WeightRenaming`] is a specialized [`WeightTransform`] for pure key renames without tensor operations. Unlike [`WeightConverter`], a `WeightRenaming` does not **claim** the key (it does not occupy the "at most one converter per key" slot), so multiple renames may chain freely both before and after a [`WeightConverter`] has fired. ```py # Legacy checkpoint compatibility @@ -138,22 +142,64 @@ WeightRenaming("LayerNorm.gamma", "LayerNorm.weight") # Module path changes WeightRenaming(".block_sparse_moe.", ".mlp.") +``` + +[`PrefixChange`] is a higher-level wrapper around [`WeightRenaming`] for the common case of stripping or adding an entire path component. The optional `model_prefix` scopes the operation to keys under that namespace: -# Adding prefixes -WeightRenaming("(.+)", "timm_model.\\1") +```py +# "model.layers.bad_prefix.weight" → "model.layers.weight" +PrefixChange(prefix_to_remove="bad_prefix", model_prefix="model.layers") +# "layers.0.weight" → "model.layers.0.weight" +PrefixChange(prefix_to_add="model") ``` ### WeightConverter -[`WeightConverter`] extends [`WeightTransform`] with a list of [`ConversionOps`]: +[`WeightConverter`] extends [`WeightTransform`] with a chain of [`ConversionOps`] that act on the collected tensors. The four supported cardinalities are: + +| Cardinality | Source patterns | Target patterns | Typical operation | +|-------------|-----------------|-----------------|-------------------| +| one-to-one | 1 | 1 | [`Transpose`], [`PermuteForRope`] | +| one-to-many | 1 | >1 | [`Chunk`] (e.g. unpack `qkv_proj`) | +| many-to-one | >1 | 1 | [`Concatenate`], [`MergeModulelist`] (e.g. fuse experts) | +| many-to-many | >1 | >1 | only with operations that explicitly support it (e.g. `ErnieFuseAndSplitTextVisionExperts`) | ```python -@dataclass(slots=True) -class WeightConverter(WeightTransform): - operations: list[ConversionOps] # Chain of operations +WeightConverter( + source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"], + target_patterns=".experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], +) ``` -It supports many-to-one (e.g., concatenating `gate` + `up` → `gate_up`), one-to-many (e.g., splitting `qkv` → `q`, `k`, `v`), and chained operations applied sequentially. +A `WeightConverter` is also reversible: `reverse_transform()` swaps source ↔ target and replaces each operation with its `reverse_op`, so a registered conversion mapping is bidirectional by construction (loading uses the list as-is, saving uses it reversed). + +### Interleaving renames and converters + +The transform list is walked in order, once per checkpoint key. For each key: + +- Every `WeightRenaming` that matches fires; multiple renames can chain. +- The first `WeightConverter` that matches **claims** the key. Subsequent converters are skipped — this guarantees the tensor is routed to a single converter for the merge/split step, and is the reason two converters with overlapping intent would be a misconfiguration. + +The typical pattern is therefore: one or more `WeightRenaming`s that put the key into the canonical form, followed by exactly one `WeightConverter` that performs the real transformation. + +```python +weight_mapping = [ + WeightRenaming("^old_prefix", "encoder"), # rename runs always + WeightConverter( # converter claims the key + "attn.qkv_proj.weight", + ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ), +] +# Load: "old_prefix.attn.qkv_proj.weight" +# → WeightRenaming → "encoder.attn.qkv_proj.weight" +# → WeightConverter → "encoder.attn.{q,k,v}_proj.weight" +# +# Save: list reversed, each transform inverted: +# → rev(WeightConverter) repacks QKV → "encoder.attn.qkv_proj.weight" +# → rev(WeightRenaming) fixes prefix → "old_prefix.attn.qkv_proj.weight" +``` ## Conversion operations @@ -169,6 +215,7 @@ Operations are fully reversible. Saving reverses the conversions and returns the | [`SplitModulelist(dim)`] | [`MergeModulelist(dim)`] | | [`Transpose(d0, d1)`] | [`Transpose(d1, d0)`] | | [`PermuteForRope()`] | [`PermuteForRope()`] | +| [`Conv3dToLinear(...)`] | [`LinearToConv3d(...)`] | ### Chunk @@ -200,12 +247,9 @@ WeightConverter( ```py WeightConverter( - ["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight",], + ["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight"], "mlp.experts.gate_up_proj", - operations=[ - MergeModulelist(dim=0), - Concatenate(dim=1), - ], + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ) ``` @@ -228,11 +272,12 @@ WeightConverter( ```py WeightConverter( ["model.layers.*.self_attn.qkv_proj.weight"], - ["model.layers.*.self_attn.q_proj.weight", "model.layers.*.self_attn.k_proj.weight", "model.layers.*.self_attn.v_proj.weight",], - operations=[ - Chunk(dim=0), - PermuteForRope(), + [ + "model.layers.*.self_attn.q_proj.weight", + "model.layers.*.self_attn.k_proj.weight", + "model.layers.*.self_attn.v_proj.weight", ], + operations=[Chunk(dim=0), PermuteForRope()], ) ``` @@ -282,12 +327,63 @@ After Concatenate(dim=1): ".experts.gate_up_proj": (8, 4096, 28672) # fused gate_up ``` -### Pattern matching +### Pattern matching rules + +- Source patterns are full regexes evaluated with `re.search()`. `^` anchors to the start of the **scoped** key (after `scope_prefix` stripping). +- `*` is a per-index wildcard that collects all matching tensors under the same source pattern (preserving checkpoint order for correct concatenation). +- Capturing groups in source patterns can be referenced as `\1`, `\2`, ... in target patterns to keep parts of the original key (e.g. layer indices). +- Scoped transforms (`scope_prefix` set) only match keys that start with `f"{scope_prefix}."`. + +## Registering a conversion mapping + +A conversion list is registered against a string key — either a `model_type` (e.g. `"mixtral"`) or a class name (e.g. `"LlavaModel"`): + +```python +from transformers.conversion_mapping import register_checkpoint_conversion_mapping + +register_checkpoint_conversion_mapping( + "my_model_type", + [ + WeightRenaming(".old.", ".new."), + WeightConverter( + ".experts.*.w1.weight", + ".experts.gate_proj", + operations=[MergeModulelist(dim=0)], + ), + ], +) +``` + +### Lookup rules -The `*` in patterns acts as a wildcard: -- During loading, it matches any numeric index (`experts.0.`, `experts.1.`, etc.). -- Tensors with the same pattern (differing only in index) are grouped together. -- The order of collection is preserved for correct concatenation. +When [`get_model_conversion_mapping`] processes a `PreTrainedModel`, every sub-`PreTrainedModel` is visited in DFS order (`nn.Module.named_modules()` filtered to `PreTrainedModel` instances). For each one: + +1. **Class-name lookup is tried first**, then `model_type`. If both are registered for the same module, the **class-name mapping wins** and the `model_type` one is ignored for that module — this lets a task head (e.g. `LlavaForConditionalGeneration`) override the shared `model_type` baseline (`"llava"`). +2. The selected mapping has `scope_prefix` set to the sub-module's dotted path (`""` for the root). +3. **Ancestor-based deduplication** decides whether to keep the mapping: + - If an **ancestor** path has already claimed the same identifier (class name or `model_type`), the sub-module is **skipped** — the ancestor's unscoped or higher-scoped mapping already covers this subtree. + - If only a **sibling** has claimed it, the sub-module is **kept** with its own `scope_prefix`. Each sibling gets its own scoped mapping. + +The class-name and `model_type` seen-lists are tracked separately, with one subtlety: when a module is matched **via its class name**, its `model_type` is *not* added to the seen-list. This is so that other modules sharing the same `model_type` but without a class-specific mapping (e.g. `DetrModel` under `DetrForSegmentation`) remain reachable through the `model_type` lookup. + +### Class-based mappings vs `model_type` aliases + +Both styles coexist in `_MODEL_TO_CONVERSION_PATTERN`: + +```python +_MODEL_TO_CONVERSION_PATTERN = { + # model_type aliases (lookup by config.model_type) + "minimax": "mixtral", + "qwen3_moe": "qwen2_moe", + "mistral3": "llava", + # class-name aliases (lookup by type(submodule).__name__) + "PaliGemmaModel": "LlavaModel", + "MaskFormerDetrDecoder": "DetrModel", + ... +} +``` + +Class-name keys are preferred when the `model_type` is shared but a specific class needs different behaviour. ## Tensor parallelism integration @@ -350,6 +446,8 @@ Quantization hooks into the loading pipeline in two ways, depending on whether t - **Pre-quantized checkpoints**: The quantizer provides [`WeightConverter`] instances (via `get_weight_conversions()`) that deserialize quantized tensors. Checkpoint dtypes are preserved to avoid unwanted casts. - **On-the-fly quantization**: The quantizer provides a quantization operation that is applied after conversion ops, quantizing weights as they are loaded. +The quantizer can also rewrite the entire conversion list at the end of [`get_model_conversion_mapping`] via `update_weight_conversions(...)` — for example, the FP8 dequantizer prepends a `Fp8Dequantize` op to every existing converter so per-block scales are applied *before* any expert-merge / concat ops flatten the per-expert structure. + ## Fast and efficient model loading Loading a model is faster and uses less memory because the loader knows which tensors are required for operations and schedules their materialization lazily. @@ -377,14 +475,16 @@ def spawn_materialize(thread_pool, tensor, device, dtype) -> Future | Callable: Sync loading is used when: - `HF_DEACTIVATE_ASYNC_LOAD=1` environment variable is set. - Disk offloading is enabled (memory constraints require sequential loading). +- On-the-fly quantization is enabled (avoids worker threads racing ahead of the quantization step). ### Materialization flow ``` 1. Checkpoint iteration (Phase 1): - - For each key, submit materialization job to ThreadPoolExecutor + - For each key, walk the transform list once + - Submit materialization job to ThreadPoolExecutor - Job returns Future (async) or Callable (sync) - - Collect into the matching WeightConverter/WeightRenaming + - Collect into the matching WeightConverter / WeightRenaming 2. Per-mapping processing (Phase 2, one mapping at a time): - materialize_tensors() waits for this mapping's Futures only @@ -431,6 +531,8 @@ def revert_weight_conversion(model, state_dict): realized_value = reversed_converter.convert(first_param_name, model=model) ``` +The list of transforms used at load time is cached on the model as `_weight_conversions` (only entries that actually fired are kept, so non-bijective renames such as [`PrefixChange`] are correctly re-applied symmetrically). When the model was instantiated without `from_pretrained` (and hence has no `_weight_conversions`), `revert_weight_conversion` falls back to recomputing the mapping via [`get_model_conversion_mapping`] and drops any [`PrefixChange`] from it (we cannot tell whether the original checkpoint had the prefix). + Target patterns may contain regex elements that need processing for the reverse direction: ```python @@ -478,6 +580,54 @@ model.layers.0.mlp.experts.down_proj # (8, 14336, 4096) ], ``` +### Composite vision-language model (Gemma3) + +Gemma3 has the canonical "different prefix logic for the head model than for the base model" setup. Two entries are registered (here via aliases pointing at the shared `"llava"` / `"LlavaModel"` lists): + +```python +# model_type: applies to Gemma3ForConditionalGeneration / Gemma3ForSequenceClassification +"gemma3": "llava" # → adds "model." in front of language_model / vision_tower / ... + +# class: applies to the inner Gemma3Model only +"Gemma3Model": "LlavaModel" # → minimal rename inside the already-prefixed namespace +``` + +DFS walk for `Gemma3ForConditionalGeneration`: + +1. Root — class lookup misses (`Gemma3ForConditionalGeneration` is not registered), model_type lookup hits (`"gemma3"`) → the `"llava"` prefix-rewriting transforms are added **unscoped**, putting `language_model`, `vision_tower`, `multi_modal_projector` under `model.*`. +2. Inner `model: Gemma3Model` — **class lookup hits** (`Gemma3Model` → `LlavaModel`) → only the `LlavaModel` mapping is applied, scoped to `"model"`. The much broader `"gemma3"` (=`"llava"`) prefix renames are **not** re-applied here, which is exactly what you want: the inner model lives in a namespace where the prefix is already correct. + +The same pair shape is used by every Llava-family VLM (`PaliGemma`, `InternVL`, `Mistral3`, ...). The model_type entry handles the head model's prefix surgery; the class entry keeps the inner base model's mapping minimal. + +### Deeper nesting with a head-specific override (DETR) + +`DetrForSegmentation` shows a head-specific override layered on top of two nested levels: + +``` +DetrForSegmentation (class registered: segmentation-only renames) +├── detr: DetrForObjectDetection (no mapping; just walked through) +│ └── model: DetrModel (class registered: shared base transforms) +│ └── backbone, encoder, ... +└── mask_head, bbox_attention (head-specific weights) +``` + +Two mappings are involved (one per registered key): + +```python +"DetrModel": [WeightRenaming("backbone.conv_encoder", "backbone"), ...] # shared base +"DetrForSegmentation": [WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), ...] # head-specific +``` + +DFS walk: + +1. Root `DetrForSegmentation` matches by class → segmentation renames added **unscoped**. +2. `detr: DetrForObjectDetection` is not registered → no transforms added; DFS continues into it. +3. `detr.model: DetrModel` matches by class → base transforms added with `scope_prefix="detr.model"`. + +This is why `DetrForObjectDetection` does not need its own mapping: the only registered mapping in its subtree (`DetrModel`) is automatically scoped to the right path. + +Class-keyed aliases reuse the base mapping without extra registration: `"MaskFormerDetrDecoder": "DetrModel"` makes a `MaskFormer` decoder pick up the same transforms under its own class name. + ### Custom operations (ERNIE 4.5 VL MoE) When the built-in operations aren't sufficient, you can create a custom [`ConversionOps`] subclass. For example, ERNIE 4.5 VL MoE needs to split a shared expert list between text and vision modalities — something no single built-in op handles. The custom `ErnieFuseAndSplitTextVisionExperts` operation splits and re-stacks experts across two target keys: @@ -544,7 +694,7 @@ These APIs are exposed to allow you to handle custom code, custom weight formats | File | Purpose | |------|---------| -| `src/transformers/core_model_loading.py` | Core loading logic, WeightConverter, ConversionOps | -| `src/transformers/conversion_mapping.py` | Built-in conversion patterns for all models | +| `src/transformers/core_model_loading.py` | Core loading logic, `WeightConverter`, `WeightRenaming`, `ConversionOps` | +| `src/transformers/conversion_mapping.py` | Built-in mappings and per-submodel composition (`get_model_conversion_mapping`) | | `src/transformers/integrations/tensor_parallel.py` | TP sharding classes and utilities | | `src/transformers/quantizers/base.py` | Quantization hooks and base class | diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index c494582dd179..f4184e346a34 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -63,7 +63,6 @@ "rt_detr_v2": "rt_detr", "pp_doclayout_v2": "rt_detr", "pp_doclayout_v3": "rt_detr", - "qwen2_5_vl": "qwen2_vl", "sam3_tracker_video": "sam3_tracker", "altclip_vision_model": "clip_vision_model", "chinese_clip_vision_model": "clip_vision_model", @@ -103,8 +102,8 @@ "LlavaOnevisionModel": "LlavaModel", "FuyuModel": "LlavaModel", "MllamaModel": "LlavaModel", - "Qwen2_5_VLModel": "Qwen2VLModel", "MaskFormerDetrDecoder": "DetrModel", + "Qwen2_5_VLForConditionalGeneration": "Qwen2VLForConditionalGeneration", } @@ -351,8 +350,7 @@ def _build_checkpoint_conversion_mapping(): target_patterns="model.language_model", ), ], - "Qwen2VLModel": [], - "qwen2_vl": [ + "Qwen2VLForConditionalGeneration": [ WeightRenaming(source_patterns=r"^visual", target_patterns="model.visual"), WeightRenaming( source_patterns=r"^model(?!\.(language_model|visual))", target_patterns="model.language_model" @@ -878,8 +876,6 @@ def get_model_conversion_mapping( seen to prevent `XForY` / `XModel` pairs from applying the same mapping twice via different lookup paths. """ - # Lazy import to avoid circular import issues - # note: this function is used in PEFT, so changing the API requires coordination weight_conversions = [] diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index ada64cf75e66..b0a154944ed5 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -686,7 +686,6 @@ def _scoped_match(self, source_key: str) -> tuple[str | None, str, re.Match[str] Returns `(prefix_dot, key_to_match, match_object)` when a branch matches, where `prefix_dot` is `None` if `scope_prefix` is unset, else `f"{scope_prefix}."`. Returns `None` when out of scope or unmatched. - Does not set `_was_used`. """ prefix_dot = None key_to_match = source_key @@ -1145,36 +1144,65 @@ class SkipParameters(Exception): def rename_source_key( source_key: str, weight_transforms: list[WeightTransform], + weight_converters: list[WeightConverter] | None = None, prefix: str | None = None, meta_state_dict: dict | None = None, ) -> tuple[str, str | None]: """ - Rename a source key according to `weight_transforms`, also handling the base model prefix. - - Transforms are applied in list order, interleaving `WeightRenaming` and `WeightConverter` - instances as they appear. The same list, reversed and with each transform individually - inverted, is used on the save path, so relative ordering is preserved in both directions. - - At most one `WeightConverter` fires per key; subsequent converters are skipped. - `WeightRenaming` always runs, even after a converter has already fired. - - Example (root rename followed by a scoped sub-model converter):: - - transforms = [ - WeightRenaming("^old_prefix", "model.vlm"), - WeightConverter("^q_proj", "qkv_proj", ...), # scope_prefix="model.vlm" - ] - # Load: "old_prefix.q_proj" - # → WeightRenaming → "model.vlm.q_proj" - # → WeightConverter → "model.vlm.qkv_proj" - # - # Save (inverted list, each transform reversed): - # "model.vlm.q_proj" - # → rev(WeightConverter) → "model.vlm.q_proj" - # → rev(WeightRenaming) → "old_prefix.q_proj" + Resolve a checkpoint key into the model-side key it should load into. + + Walks `weight_transforms` in list order, applying every matching `WeightRenaming` and at + most one matching `WeightConverter`. The same list, reversed and with each transform + individually inverted, is used on the save path, so relative ordering is preserved in + both directions. Optionally adds or strips `prefix` so the returned key matches an entry + in `meta_state_dict`. + + Args: + source_key (`str`): + The original key from the checkpoint state dict. + weight_transforms (`list[WeightTransform]`): + Ordered list of transforms (a mix of `WeightRenaming` and `WeightConverter`). + `WeightRenaming` always runs every time it matches; multiple renames may chain. + At most one `WeightConverter` may claim a key — subsequent converters are + skipped, and the returned `source_pattern` identifies which one matched. + weight_converters (`list[WeightConverter]`, *optional*): + Backward-compat shim for the previous `(source_key, weight_renamings, + weight_converters, ...)` signature, used by downstream libraries (e.g. peft). + When provided, it is appended to `weight_transforms` so renamings run before converters. + prefix (`str`, *optional*): + Base-model prefix (e.g. `"model"`) to add to or strip from the renamed key when + needed for it to match an entry in `meta_state_dict`. Only applied on the load + path (i.e. when both `prefix` and `meta_state_dict` are provided). + meta_state_dict (`dict`, *optional*): + Meta state dict of the target model, used to decide whether `prefix` should be + added or stripped. + + Returns: + `tuple[str, str | None]`: The renamed key, and the source pattern of the matching + `WeightConverter` (or `None` if no converter matched). + + Example (root rename followed by a scoped sub-model converter): + + ```python + transforms = [ + WeightRenaming("^old_prefix", "model.vlm"), + WeightConverter("^q_proj", "qkv_proj", ...), # scope_prefix="model.vlm" + ] + # Load: "old_prefix.q_proj" + # → WeightRenaming → "model.vlm.q_proj" + # → WeightConverter → "model.vlm.qkv_proj" + # + # Save (inverted list, each transform reversed): + # "model.vlm.q_proj" + # → rev(WeightConverter) → "model.vlm.q_proj" + # → rev(WeightRenaming) → "old_prefix.q_proj" + ``` """ + if weight_converters is not None: + weight_transforms = [*weight_transforms, *weight_converters] + renamed_key = source_key - source_pattern = None + source_pattern_that_matched = None for transform in weight_transforms: if isinstance(transform, WeightConverter): @@ -1182,11 +1210,11 @@ def rename_source_key( # the caller which converter to feed this tensor to; a second match would redirect # the tensor to the wrong converter, leaving the first one incomplete. # TODO: `continue` here can hide misconfigured mapping (two converters with overlapping patterns). Raise an error instead. - if source_pattern is not None: + if source_pattern_that_matched is not None: continue renamed_key, matched_pattern = transform.rename_source_key(renamed_key) if matched_pattern is not None: - source_pattern = matched_pattern + source_pattern_that_matched = matched_pattern else: # WeightRenaming always runs. Multiple renames can safely chain before (and after) a converter fires. renamed_key, _ = transform.rename_source_key(renamed_key) @@ -1201,7 +1229,7 @@ def rename_source_key( elif meta_state_dict.get(f"{prefix}.{renamed_key}") is not None: renamed_key = f"{prefix}.{renamed_key}" - return renamed_key, source_pattern + return renamed_key, source_pattern_that_matched def convert_and_load_state_dict_in_model( @@ -1348,10 +1376,14 @@ def convert_and_load_state_dict_in_model( state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) for original_key, tensor in state_dict: # 1. Rename the key according to all renaming and weight conversion patterns. - renamed_key, source_pattern = rename_source_key(original_key, weight_mapping, prefix, meta_model_state_dict) + renamed_key, source_pattern = rename_source_key( + original_key, weight_mapping, prefix=prefix, meta_state_dict=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) + renamed_key, source_pattern = rename_source_key( + original_key, [], prefix=prefix, meta_state_dict=meta_model_state_dict + ) # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 0b278ebb4d40..aa14cadbd702 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -468,7 +468,8 @@ def accelerate_disk_offload( # Update the weight names according to the `weight_mapping` weight_renaming_map = { - rename_source_key(k, transforms, 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 diff --git a/src/transformers/integrations/deepspeed.py b/src/transformers/integrations/deepspeed.py index 79f3896cb48c..a525c32615ce 100644 --- a/src/transformers/integrations/deepspeed.py +++ b/src/transformers/integrations/deepspeed.py @@ -366,7 +366,9 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping): if len(converters) == 0: new_state_dict = {} for original_key, tensor in state_dict.items(): - renamed_key, _ = rename_source_key(original_key, weight_mapping, 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 @@ -385,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, weight_mapping, 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: From 670ac73d287741d77d7b90a7ad29813524f10994 Mon Sep 17 00:00:00 2001 From: yonigozlan Date: Thu, 7 May 2026 00:46:34 +0000 Subject: [PATCH 15/15] Fix edge case rename -> convert (fix deepseekv4) --- docs/source/en/weightconverter.md | 8 ++--- src/transformers/conversion_mapping.py | 36 ++++++++++----------- src/transformers/core_model_loading.py | 45 +++++++++++++++----------- 3 files changed, 48 insertions(+), 41 deletions(-) diff --git a/docs/source/en/weightconverter.md b/docs/source/en/weightconverter.md index a9269c8492bf..d8eea6d6bf82 100644 --- a/docs/source/en/weightconverter.md +++ b/docs/source/en/weightconverter.md @@ -174,15 +174,15 @@ WeightConverter( A `WeightConverter` is also reversible: `reverse_transform()` swaps source ↔ target and replaces each operation with its `reverse_op`, so a registered conversion mapping is bidirectional by construction (loading uses the list as-is, saving uses it reversed). -### Interleaving renames and converters +### Ordering renames and converters + +List `WeightRenaming` entries before `WeightConverter` entries, and keep their leaves disjoint: renames normalise the keys converters consume, but never target a leaf a converter produces. The save path relies on this to invert the mapping in two phases (reverse converters, then reverse renames). The transform list is walked in order, once per checkpoint key. For each key: - Every `WeightRenaming` that matches fires; multiple renames can chain. - The first `WeightConverter` that matches **claims** the key. Subsequent converters are skipped — this guarantees the tensor is routed to a single converter for the merge/split step, and is the reason two converters with overlapping intent would be a misconfiguration. -The typical pattern is therefore: one or more `WeightRenaming`s that put the key into the canonical form, followed by exactly one `WeightConverter` that performs the real transformation. - ```python weight_mapping = [ WeightRenaming("^old_prefix", "encoder"), # rename runs always @@ -514,7 +514,7 @@ These worst-case scenarios are uncommon. The actual memory peak tends to stay cl ## Reversibility -The system supports saving models with the inverse transformations, enabling round-trip save/load: +The system supports saving models with the inverse transformations, enabling round-trip save/load. Saving runs in two phases: reversed converters first (each tensor matches at most one), then reversed renames, per the ordering rule above. ```python def revert_weight_conversion(model, state_dict): diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 2f4df46a9283..5af787f5a977 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -283,15 +283,15 @@ def _build_checkpoint_conversion_mapping(): ), WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.w1\.weight", - r"mlp\.experts\..*\.w3\.weight", + "mlp.experts.*.w1.weight", + "mlp.experts.*.w3.weight", ], - target_patterns=r"mlp.experts.gate_up_proj", + target_patterns="mlp.experts.gate_up_proj", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( - source_patterns=r"mlp\.experts\..*\.w2\.weight", - target_patterns=r"mlp.experts.down_proj", + source_patterns="mlp.experts.*.w2.weight", + target_patterns="mlp.experts.down_proj", operations=[MergeModulelist(dim=0)], ), ], @@ -652,6 +652,9 @@ def _build_checkpoint_conversion_mapping(): "jina_embeddings_v3": [ WeightRenaming(source_patterns="emb_ln", target_patterns="embeddings.LayerNorm"), WeightRenaming(source_patterns="encoder.layers", target_patterns="layers"), + WeightRenaming(source_patterns="mixer.out_proj", target_patterns="self_attn.o_proj"), + WeightRenaming(source_patterns="norm1", target_patterns="post_attention_layernorm"), + WeightRenaming(source_patterns="norm2", target_patterns="post_mlp_layernorm"), WeightConverter( source_patterns="mixer.Wqkv", target_patterns=[ @@ -661,9 +664,6 @@ def _build_checkpoint_conversion_mapping(): ], operations=[Chunk(dim=0)], ), - WeightRenaming(source_patterns="mixer.out_proj", target_patterns="self_attn.o_proj"), - WeightRenaming(source_patterns="norm1", target_patterns="post_attention_layernorm"), - WeightRenaming(source_patterns="norm2", target_patterns="post_mlp_layernorm"), ], "cohere_asr": [ WeightRenaming(r"encoder\.pre_encode\.conv\.", r"encoder.subsampling.layers."), @@ -703,16 +703,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"\.ls1", r"\.lambda_1"), WeightRenaming(r"\.ls2", r"\.lambda_2"), WeightRenaming(r"(layers\.\d+)\.attn\.proj\.", r"\1.attention.projection_layer."), - WeightConverter( - source_patterns=["attn.qkv.weight"], - target_patterns=["attention.q_proj.weight", "attention.k_proj.weight", "attention.v_proj.weight"], - operations=[Chunk(dim=0)], - ), - WeightConverter( - source_patterns=["attn.qkv.bias"], - target_patterns=["attention.q_proj.bias", "attention.k_proj.bias", "attention.v_proj.bias"], - operations=[Chunk(dim=0)], - ), WeightRenaming(r"\.norm1\.", r"\.layernorm_before\."), WeightRenaming(r"\.norm2\.", r"\.layernorm_after\."), WeightRenaming(r"\.embeddings\.class_embedding", r"\.embeddings\.cls_token"), @@ -723,6 +713,16 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"^mlp1\.0\.", r"model\.multi_modal_projector\.layer_norm\."), WeightRenaming(r"^mlp1\.1\.", r"model\.multi_modal_projector\.linear_1\."), WeightRenaming(r"^mlp1\.3\.", r"model\.multi_modal_projector\.linear_2\."), + WeightConverter( + source_patterns=["attn.qkv.weight"], + target_patterns=["attention.q_proj.weight", "attention.k_proj.weight", "attention.v_proj.weight"], + operations=[Chunk(dim=0)], + ), + WeightConverter( + source_patterns=["attn.qkv.bias"], + target_patterns=["attention.q_proj.bias", "attention.k_proj.bias", "attention.v_proj.bias"], + operations=[Chunk(dim=0)], + ), ], "legacy": [ WeightRenaming( diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index b0a154944ed5..67d5d54f4b1b 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1542,33 +1542,40 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch # Important: we need to revert the order here, so that potential conversions from submodels are performed first weight_conversions = weight_conversions[::-1] - # Reverse all Transform to correctly match keys - reverse_weight_conversion = [conversion.reverse_transform() for conversion in weight_conversions] - # If we are still here, we need to create the (reverse) conversion mapping from scratch - converters = [entry for entry in reverse_weight_conversion if isinstance(entry, WeightConverter)] - pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} - conversion_mapping = {} + # Two-phase save: first reverse converters, then reverse renamings. Relies on the rule that + # WeightRenamings never operate on WeightConverter outputs (see WeightTransform docstring). + inverted_transforms = [transform.reverse_transform() for transform in weight_conversions] + inverted_converters = [transform for transform in inverted_transforms if isinstance(transform, WeightConverter)] + inverted_renamings = [transform for transform in inverted_transforms if not isinstance(transform, WeightConverter)] + converter_by_pattern = { + pattern: converter for converter in inverted_converters for pattern in converter.source_patterns + } + conversion_mapping: dict[str, WeightTransform] = {} state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) for original_key, tensor in state_dict: - renamed_key, source_pattern = rename_source_key(original_key, reverse_weight_conversion) - - 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) + # `converter_key`: key after phase-1 (converter namespace, used as layer_name by convert()). + # `checkpoint_key`: key after phase-2 (final saved name, layer_name for plain renamings). + converter_key, matched_pattern = rename_source_key(original_key, inverted_converters) + checkpoint_key, _ = rename_source_key(converter_key, inverted_renamings) + + if matched_pattern is not None: + # Bucket under converter_key so all sibling inputs land in the same converter instance. + mapping = conversion_mapping.setdefault(converter_key, deepcopy(converter_by_pattern[matched_pattern])) else: - mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) - source_pattern = original_key + mapping = conversion_mapping.setdefault(checkpoint_key, WeightRenaming(original_key, checkpoint_key)) + matched_pattern = original_key - mapping.add_tensor(renamed_key, original_key, source_pattern, tensor) + mapping.add_tensor(checkpoint_key, original_key, matched_pattern, tensor) new_state_dict = {} - for first_param_name, reversed_converter in conversion_mapping.items(): - # Apply the reverse converter - realized_value = reversed_converter.convert(first_param_name, model=model, config=model.config) - for target_name, param in realized_value.items(): + for layer_name, mapping in conversion_mapping.items(): + realized = mapping.convert(layer_name, model=model, config=model.config) + for target_name, param in realized.items(): param = param[0] if isinstance(param, list) else param + if isinstance(mapping, WeightConverter): + # Bring converter outputs from converter namespace into checkpoint namespace. + target_name, _ = rename_source_key(target_name, inverted_renamings) new_state_dict[target_name] = param return new_state_dict