From 47570d77adfb01fae6ca9f5b6fc4c37fd0844a52 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 16:42:04 +0100
Subject: [PATCH 01/36] first shot

---
 src/transformers/conversion_mapping.py |  40 ++--
 src/transformers/core_model_loading.py | 253 +++++++++++++++----------
 2 files changed, 169 insertions(+), 124 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index b6aad7e94650..02c837425bfe 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -28,11 +28,11 @@ def _build_checkpoint_conversion_mapping():
         "mixtral": [
             WeightRenaming(".block_sparse_moe.gate", ".mlp.gate"),
             WeightConverter(
-                source_keys=[
+                source_patterns=[
                     "block_sparse_moe.experts.*.w1.weight",
                     "block_sparse_moe.experts.*.w3.weight",
                 ],  # you give me a list of 2 keys, I collect a list of a list of tensors
-                target_keys="mlp.experts.gate_up_proj",  # target key gets the list of two tensors
+                target_patterns="mlp.experts.gate_up_proj",  # target key gets the list of two tensors
                 operations=[
                     MergeModulelist(
                         dim=0
@@ -41,10 +41,10 @@ def _build_checkpoint_conversion_mapping():
                 ],  # we want the loading to add this shard operation here. Though we can't shard after concats and merge, needs to be first
             ),
             WeightConverter(
-                source_keys=[
+                source_patterns=[
                     "block_sparse_moe.experts.*.w2.weight",
                 ],
-                target_keys="mlp.experts.down_proj",  # target key gets the list of two tensors
+                target_patterns="mlp.experts.down_proj",  # target key gets the list of two tensors
                 operations=[
                     MergeModulelist(
                         dim=0
@@ -54,50 +54,50 @@ def _build_checkpoint_conversion_mapping():
         ],
         "qwen2_moe": [
             WeightConverter(
-                source_keys=[
+                source_patterns=[
                     "mlp.experts.*.gate_proj.weight",
                     "mlp.experts.*.up_proj.weight",
                 ],
-                target_keys="mlp.experts.gate_up_proj",
+                target_patterns="mlp.experts.gate_up_proj",
                 operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
             ),
             WeightConverter(
-                source_keys=["mlp.experts.*.down_proj.weight"],
-                target_keys="mlp.experts.down_proj",
+                source_patterns=["mlp.experts.*.down_proj.weight"],
+                target_patterns="mlp.experts.down_proj",
                 operations=[MergeModulelist(dim=0)],
             ),
         ],
         "legacy": [
             WeightRenaming(
-                source_keys="LayerNorm.gamma",
-                target_keys="LayerNorm.weight",
+                source_patterns="LayerNorm.gamma",
+                target_patterns="LayerNorm.weight",
             ),
             WeightRenaming(
-                source_keys="LayerNorm.beta",
-                target_keys="LayerNorm.bias",
+                source_patterns="LayerNorm.beta",
+                target_patterns="LayerNorm.bias",
             ),
         ],
     }
     if hasattr(torch.nn.utils.parametrizations, "weight_norm"):
         mapping["legacy"] += [
             WeightRenaming(
-                source_keys="weight_g",
-                target_keys="parametrizations.weight.original0",
+                source_patterns="weight_g",
+                target_patterns="parametrizations.weight.original0",
             ),
             WeightRenaming(
-                source_keys="weight_v",
-                target_keys="parametrizations.weight.original1",
+                source_patterns="weight_v",
+                target_patterns="parametrizations.weight.original1",
             ),
         ]
     else:
         mapping["legacy"] += [
             WeightRenaming(
-                source_keys="parametrizations.weight.original0",
-                target_keys="weight_g",
+                source_patterns="parametrizations.weight.original0",
+                target_patterns="weight_g",
             ),
             WeightRenaming(
-                source_keys="parametrizations.weight.original1",
-                target_keys="weight_v",
+                source_patterns="parametrizations.weight.original1",
+                target_patterns="weight_v",
             ),
         ]
 
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 30ce9d6b6732..84eda0c9f8f8 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -19,7 +19,8 @@
 import os
 import re
 from abc import abstractmethod
-from collections.abc import MutableMapping, MutableSet, Sequence
+from collections import defaultdict
+from collections.abc import MutableMapping, MutableSet
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import contextmanager
 from copy import deepcopy
@@ -46,8 +47,6 @@
 
 logger = logging.get_logger(__name__)
 
-logger = logging.get_logger(__name__)
-
 
 def compile_glob_rule(source_glob: str, target_glob: str) -> tuple[re.Pattern, str]:
     """
@@ -82,13 +81,13 @@ def build_glob_alternation(
     i = 0
     for glob in globs:
         if isinstance(glob, (WeightRenaming, WeightConverter)):
-            for src in glob.source_keys:
+            for src in glob.source_patterns:
                 group_name = f"g{i}"
                 src_group_to_glob[group_name] = src
                 i += 1
                 body = src.replace("*", r".*")
                 branches.append(f"(?P<{group_name}>{body})")
-                tgt_group_to_glob[group_name] = glob.target_keys[0]  # we index witht the first target
+                tgt_group_to_glob[group_name] = glob.target_patterns[0]  # we index witht the first target
         else:
             group_name = f"g{i}"
             src_group_to_glob[group_name] = glob
@@ -105,15 +104,12 @@ def build_glob_alternation(
 class ConversionOps:
     """Base class for weight conversion operations."""
 
-    # The inverse operation class, will be used when saving the checkpoint
-    reverse_op: type[ConversionOps]
-
     @abstractmethod
     def convert(
         self,
         value: dict[str, Any],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
@@ -122,64 +118,72 @@ def convert(
     ) -> dict[str, list[torch.Tensor]]:
         raise NotImplementedError
 
+    @property
+    def reverse_op(self) -> ConversionOps:
+        raise NotImplementedError
 
-class Chunk(ConversionOps):
-    """Split a tensor along ``dim`` into equally sized chunks or using explicit ``sizes``."""
 
-    reverse_op: type[ConversionOps]
+class Chunk(ConversionOps):
+    """Split a tensor along ``dim`` into equally sized chunks."""
 
-    def __init__(self, dim: int = 0, chunks: Optional[int] = None, sizes: Optional[Sequence[int]] = None):
+    def __init__(self, dim: int = 0):
         self.dim = dim
-        self.chunks = chunks
-        self.sizes = list(sizes) if sizes is not None else None
-        self.reverse_op = Concatenate
 
     def convert(
         self,
         value: dict[str, list[torch.Tensor]],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
         config,
     ) -> dict[str, list[torch.Tensor]]:
-        tensors = next(iter(value.values()))
-        tensor = tensors[0]
-        sizes = len(target_keys)
+        if len(value) > 1:
+            raise ValueError("Unexpected value in `Chunk`!")
+        tensor = next(iter(value.values()))[0]
+        sizes = len(target_patterns)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
-        return {full_layer_name.replace(target_keys[0], target): [chunk] for target, chunk in zip(target_keys, chunks)}
+        return {
+            full_layer_name.replace(target_patterns[0], target): [chunk]
+            for target, chunk in zip(target_patterns, chunks)
+        }
 
+    @property
+    def reverse_op(self) -> ConversionOps:
+        return Concatenate(self.dim)
 
-class Concatenate(ConversionOps):
-    """Concatenate tensors along `dim` using a reusable buffer."""
 
-    reverse_op: type[ConversionOps]
+class Concatenate(ConversionOps):
+    """Concatenate tensors along `dim`."""
 
     def __init__(self, dim: int = 0):
         self.dim = dim
-        self.reverse_op = Chunk
 
     @torch.no_grad
     def convert(
         self,
         value: dict[str, list[torch.Tensor]],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
-        if len(target_keys) != 1:
+        if len(target_patterns) != 1:
             raise ValueError("Concatenate expects a single target key.")
-        if len(value) != len(source_keys):
+        if len(value) != len(source_patterns):
             raise ValueError("Concatenate received an unexpected number of tensors compared to source keys.")
 
         return {full_layer_name: torch.cat(tuple(value.values()), dim=self.dim)}
 
+    @property
+    def reverse_op(self) -> ConversionOps:
+        return Chunk(self.dim)
 
-class MergeModulelist(Concatenate):
+
+class MergeModulelist(ConversionOps):
     """
     Merge a list of tensors into a single tensor along the first dimension.
     We explicitly define this because for EP or TP you want to make sure you know what you are doing!
@@ -187,62 +191,61 @@ class MergeModulelist(Concatenate):
     """
 
     def __init__(self, dim: int = 0):
-        super().__init__(dim=dim)
-        self.reverse_op = SplitModulelist
+        self.dim = dim
 
     @torch.no_grad
     def convert(
         self,
         value: dict[str, list[torch.Tensor]],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
-        for idx, key in enumerate(value.keys()):
-            tensors = value.get(key, [])
-            if len(source_keys) == 1:
+        for key, tensors in value.items():
+            if len(source_patterns) == 1:
                 key = full_layer_name
-            stacked = torch.stack(tensors, dim=self.dim)
-            merged[key] = stacked
+            merged[key] = torch.stack(tensors, dim=self.dim)
         return merged
 
+    @property
+    def reverse_op(self) -> ConversionOps:
+        return SplitModulelist(self.dim)
+
 
 class SplitModulelist(ConversionOps):
     """Inverse of :class:`MergeModulelist` using explicit split sizes per group."""
 
-    def __init__(self, sizes: Sequence[Sequence[int]], dim: int = 0):
-        if not isinstance(sizes, Sequence) or not all(isinstance(sub, Sequence) and sub for sub in sizes):
-            raise ValueError("`sizes` must be a sequence of non-empty sequences of integers.")
-        self.sizes = [list(sub) for sub in sizes]
+    def __init__(self, dim: int = 0):
         self.dim = dim
-        self.reverse_op = MergeModulelist
 
     @torch.no_grad
     def convert(
         self,
         value: dict[str, list[torch.Tensor]],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
         config,
     ) -> dict[str, list[torch.Tensor]]:
-        if len(value) != len(self.sizes):
-            raise ValueError("SplitModulelist received an unexpected number of tensors.")
-        result: dict[str, list[torch.Tensor]] = {}
-        for (key, tensors), split_sizes in zip(value.items(), self.sizes):
-            if len(tensors) != 1:
-                raise ValueError("SplitModulelist expects exactly one tensor per key.")
-            current_tensor = tensors[0]
-            if not isinstance(current_tensor, torch.Tensor):
-                raise TypeError("SplitModulelist can only split torch.Tensor instances.")
-            result[key] = list(torch.split(current_tensor, split_sizes, dim=self.dim))
-        return result
+        if len(value) > 1:
+            raise ValueError("Unexpected value in `SplitModulelist`!")
+        tensor = next(iter(value.values()))[0]
+        sizes = len(target_patterns)
+        chunks = torch.split(tensor, sizes, dim=self.dim)
+        return {
+            full_layer_name.replace(target_patterns[0], target): [chunk]
+            for target, chunk in zip(target_patterns, chunks)
+        }
+
+    @property
+    def reverse_op(self) -> ConversionOps:
+        return MergeModulelist(self.dim)
 
 
 class PermuteForRope(ConversionOps):
@@ -265,8 +268,8 @@ def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
     def convert(
         self,
         value: dict[str, list[torch.Tensor]],
-        source_keys: list[str],
-        target_keys: list[str],
+        source_patterns: list[str],
+        target_patterns: list[str],
         full_layer_name: str,
         model,
         missing_keys,
@@ -283,8 +286,8 @@ def convert(
 
 @dataclass(slots=True)
 class WeightTransform:
-    source_keys: Union[str, list[str]] = field(init=True)
-    target_keys: Union[str, list[str]] = field(init=True)
+    source_patterns: Union[str, list[str]] = field(init=True)
+    target_patterns: Union[str, list[str]] = field(init=True)
 
     distributed_operation: Optional[TensorParallelLayer] = None
     quantization_operation: Optional[ConversionOps] = None
@@ -293,10 +296,10 @@ class WeightTransform:
     layer_targets: dict[str, set[str]] = field(default_factory=dict, init=False)
 
     def __post_init__(self):
-        if isinstance(self.source_keys, str):
-            self.source_keys = [self.source_keys]
-        if isinstance(self.target_keys, str):
-            self.target_keys = [self.target_keys]
+        if isinstance(self.source_patterns, str):
+            self.source_patterns = [self.source_patterns]
+        if isinstance(self.target_patterns, str):
+            self.target_patterns = [self.target_patterns]
 
     def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):
         bucket = self.collected_tensors.setdefault(source_pattern, [])
@@ -305,6 +308,36 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         bucket = self.layer_targets.setdefault(target_key, set())
         bucket.add(source_key)
 
+    def reset(self) -> None:
+        """Clean-up the collected tensors."""
+        self.collected_tensors = {}
+
+    def reverse_transform(self) -> WeightTransform:
+        if self.distributed_operation is not None or self.quantization_operation is not None:
+            raise ValueError("Cannot reverse the transform with TP or quantization")
+        if len(self.collected_tensors) == 0:
+            raise ValueError(
+                "You can only call `reverse_transform` after the `WeightTransform` instance has been populated"
+                "with the keys and tensors!"
+            )
+        reverse_transform = self.__class__(source_patterns=self.target_patterns, target_patterns=self.source_patterns)
+        # Find the full names of the params we will need to use later for the reverse transform
+        reverse_layer_targets = defaultdict(set)
+        reverse_collected_tensors = defaultdict(list)
+        for target_key, all_sources in self.layer_targets.items():
+            matched_target_pattern = next(pat for pat in self.target_patterns if re.search(pat, target_key))
+            for source in all_sources:
+                reverse_layer_targets[source].add(target_key)
+                reverse_collected_tensors[matched_target_pattern].append(source)
+        reverse_collected_tensors = sorted(set(reverse_collected_tensors))
+        reverse_transform.layer_targets = reverse_layer_targets
+        reverse_transform.collected_tensors = reverse_collected_tensors
+        # Add the reverse ops if applicable
+        if hasattr(reverse_transform, "operations"):
+            # All reverse ops, in reverse order
+            reverse_transform.operations = [op.reverso_op for op in self.operations[::-1]]
+        return reverse_transform
+
 
 @dataclass(slots=True)
 class WeightRenaming(WeightTransform):
@@ -320,15 +353,17 @@ def convert(
         misc: Optional[MutableMapping[str, str]] = None,
     ):
         for pattern, futures in self.collected_tensors.items():
-            self.collected_tensors[pattern] = [future.result() for future in futures]
+            self.collected_tensors[pattern] = (
+                futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]
+            )
 
         collected_tensors = self.collected_tensors
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     self.collected_tensors,
-                    source_keys=self.source_keys,
-                    target_keys=self.target_keys,
+                    source_patterns=self.source_patterns,
+                    target_patterns=self.target_patterns,
                     full_layer_name=layer_name,
                     model=model,
                     config=config,
@@ -344,9 +379,9 @@ class WeightConverter(WeightTransform):
 
     def __post_init__(self):
         WeightTransform.__post_init__(self)
-        if bool(len(self.source_keys) - 1) + bool(len(self.target_keys) - 1) >= 2:
+        if bool(len(self.source_patterns) - 1) + bool(len(self.target_patterns) - 1) >= 2:
             raise ValueError(
-                f"source keys={self.source_keys}, target_keys={self.target_keys} but you can only have one to many, one to one or many to one."
+                f"source keys={self.source_patterns}, target_patterns={self.target_patterns} but you can only have one to many, one to one or many to one."
             )
         if not self.operations:
             raise ValueError("WeightConverter requires at least one operation.")
@@ -361,15 +396,17 @@ def convert(
         misc: Optional[MutableMapping[str, str]] = None,
     ):
         for pattern, futures in self.collected_tensors.items():
-            self.collected_tensors[pattern] = [future.result() for future in futures]
+            self.collected_tensors[pattern] = (
+                futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]
+            )
 
         collected_tensors = self.collected_tensors
         for op in self.operations:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), op):
                 collected_tensors = op.convert(
                     collected_tensors,
-                    source_keys=self.source_keys,
-                    target_keys=self.target_keys,
+                    source_patterns=self.source_patterns,
+                    target_patterns=self.target_patterns,
                     full_layer_name=layer_name,
                     model=model,
                     config=config,
@@ -379,8 +416,8 @@ def convert(
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     collected_tensors,
-                    source_keys=self.source_keys,
-                    target_keys=self.target_keys,
+                    source_patterns=self.source_patterns,
+                    target_patterns=self.target_patterns,
                     full_layer_name=layer_name,
                     config=config,
                     model=model,
@@ -576,8 +613,8 @@ def convert_and_load_state_dict_in_model(
     {
         "model.layers.0.attention.q.weight": # Notice here there is only the first key of the target keys
             WeightConverter(
-                source_keys=["qkv"],
-                target_keys=["q", "k","v"],
+                source_patterns=["qkv"],
+                target_patterns=["q", "k","v"],
                 operations=[Chunk(dim=0, chunks=3)]),
                 collected_tensors={
                     "qkv": [Future, Future, Future]},
@@ -596,8 +633,8 @@ def convert_and_load_state_dict_in_model(
     For example for:
     ```python
     WeightConverter(
-        source_keys=["mlp.experts.*.gate_proj.weight","mlp.experts.*.up_proj.weight"],
-        target_keys="mlp.experts.gate_up_proj",
+        source_patterns=["mlp.experts.*.gate_proj.weight","mlp.experts.*.up_proj.weight"],
+        target_patterns="mlp.experts.gate_up_proj",
         operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
     )
     ```
@@ -633,8 +670,8 @@ def convert_and_load_state_dict_in_model(
     ```python
     # for "medmekk/llama-3.2-1b-float8-torchao"
     WeightConverter(
-        source_keys=[":qdata", ":scale"],
-        target_keys="",
+        source_patterns=[":qdata", ":scale"],
+        target_patterns="",
         operations=[TorchaoDeserialize()],
     )
     ```
@@ -642,8 +679,8 @@ def convert_and_load_state_dict_in_model(
     ```python
     all_weight_mapping = {
         "model.layers.13.self_attn.o_proj.weight": WeightConverter(
-            source_keys=[":qdata", ":scale"],
-            target_keys="",
+            source_patterns=[":qdata", ":scale"],
+            target_patterns="",
             operations=[TorchaoDeserialize()],
             collected_tensors={
                 ":qdata": [Future],
@@ -687,7 +724,7 @@ def convert_and_load_state_dict_in_model(
     if dtype_plan != {}:
         dtype_policy_alt, dtype_policy_by_group_name, _ = build_glob_alternation(list(dtype_plan.keys()))
 
-    pattern_to_converter = {k: converter for converter in converters for k in converter.source_keys}
+    pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns}
 
     state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
     for original_key, tensor in state_dict:
@@ -769,8 +806,8 @@ def convert_and_load_state_dict_in_model(
             mapping.add_tensor(renamed_key, original_key, source_pattern, future)
         elif matched_pattern:  # add all target keys as unexpected
             mapping = pattern_to_converter[src_group_to_glob[matched_pattern.lastgroup]]
-            for k in mapping.target_keys:
-                unexpected_keys.add(renamed_key.replace(mapping.target_keys[0], k))
+            for k in mapping.target_patterns:
+                unexpected_keys.add(renamed_key.replace(mapping.target_patterns[0], k))
         else:
             unexpected_keys.add(renamed_key)
 
@@ -810,27 +847,35 @@ def convert_and_load_state_dict_in_model(
                             mapping.distributed_operation,
                             hf_quantizer,
                         )
+
+                # Cleanup the tensors
+                mapping.reset()
             except SkipLayer:
                 continue
 
+    # Keep computed mapping as attribute for later saving
+    model._weight_loading_mapping = param_name_to_load
     thread_pool.shutdown(wait=False)
     return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, misc
 
 
-# TODO this is not done yet!
 def revert_weight_conversion(model, state_dict):
-    mapping = getattr(model, "_checkpoint_conversion_mapping", {})  # IDK why but setting this will fail all llava.
-    reverse_key_mapping = [(v, k) for k, v in mapping.items()]
-    original_state_dict = {}
-    for key, value in state_dict.items():
-        for pattern, inverse_converter in reverse_key_mapping:
-            # TODO FIXME you name it
-            replacement = inverse_converter.lstrip("^")  # strip off un-needed chars and patterns
-            replacement = re.sub(r"\(.*\)", "", replacement)
-            key, n_replace = re.subn(pattern, replacement, key)
-            # Early exit of the loop
-            if n_replace > 0:
-                break
-        original_state_dict[key] = value
-    state_dict = original_state_dict
-    return state_dict
+    original_mapping = getattr(model, "_weight_loading_mapping", None)
+    if original_mapping is None:
+        return state_dict
+
+    new_state_dict = {}
+    for converter in original_mapping.values():
+        reversed_converter = converter.reverse_transform()
+        reversed_converter.collected_tensors = {
+            k: model.get_parameter_or_buffer(v) for k, v in reversed_converter.collected_tensors.items()
+        }
+        first_param_name = next(iter(reversed_converter.layer_targets.keys()))
+
+        # Apply the reverse converter
+        realized_value, misc = reversed_converter.convert(first_param_name, model=model, config=model.config)
+        for target_name, param in realized_value.items():
+            param = param[0] if isinstance(param, list) else param
+            new_state_dict[target_name] = param
+
+    return new_state_dict

From 69e5c9ea9311ccd6f4622d9b350b3a9fa70442c6 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 16:44:47 +0100
Subject: [PATCH 02/36] default to reversing

---
 src/transformers/modeling_utils.py | 15 +++------------
 1 file changed, 3 insertions(+), 12 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 5172d74cc397..c46461a87146 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -2955,7 +2955,7 @@ def save_pretrained(
         variant: Optional[str] = None,
         token: Optional[Union[str, bool]] = None,
         save_peft_format: bool = True,
-        save_original_format: bool = False,  # TODO next PR will make it go to True
+        save_original_format: bool = True,
         **kwargs,
     ):
         """
@@ -3223,17 +3223,8 @@ def save_pretrained(
                     "This can also just mean that the module's tied weight keys are wrong vs the actual tied weights in the model.",
                 )
 
-        if (
-            any(
-                allowed_name in class_name.__name__.lower()
-                for class_name in self.__class__.__mro__[:-1]
-                for allowed_name in VLMS
-            )
-            or save_original_format
-        ):
-            # MEGA BIG TODO HERE: self._conversion_ops needs to be used to save the final ckpt
-            # using what was loaded. Actually self._conversion_ops wont work because we need it
-            # even if the files are not legacy -> thus no conversion happened
+        # Revert all renaming and/or weight operations
+        if save_original_format:
             state_dict = revert_weight_conversion(self, state_dict)
 
         # Shard the model if it is too big.

From 8fdc41f7a58215bdcf6d1f1343a049dc0ef0934b Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 16:53:19 +0100
Subject: [PATCH 03/36] oupso

---
 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 84eda0c9f8f8..396096eed814 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -315,7 +315,7 @@ def reset(self) -> None:
     def reverse_transform(self) -> WeightTransform:
         if self.distributed_operation is not None or self.quantization_operation is not None:
             raise ValueError("Cannot reverse the transform with TP or quantization")
-        if len(self.collected_tensors) == 0:
+        if len(self.layer_targets) == 0:
             raise ValueError(
                 "You can only call `reverse_transform` after the `WeightTransform` instance has been populated"
                 "with the keys and tensors!"

From 938835423a51b0dbbac87d275a50f72ce57999dc Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 17:02:21 +0100
Subject: [PATCH 04/36] oupsi 2

---
 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 396096eed814..612494b1679a 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -329,7 +329,7 @@ def reverse_transform(self) -> WeightTransform:
             for source in all_sources:
                 reverse_layer_targets[source].add(target_key)
                 reverse_collected_tensors[matched_target_pattern].append(source)
-        reverse_collected_tensors = sorted(set(reverse_collected_tensors))
+        reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
         reverse_transform.layer_targets = reverse_layer_targets
         reverse_transform.collected_tensors = reverse_collected_tensors
         # Add the reverse ops if applicable

From a9eb2111d8e07b461866cefd6909e33abcaa8f35 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 17:07:59 +0100
Subject: [PATCH 05/36] oupsi 3

---
 src/transformers/core_model_loading.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 612494b1679a..21b694c464c1 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -868,7 +868,8 @@ def revert_weight_conversion(model, state_dict):
     for converter in original_mapping.values():
         reversed_converter = converter.reverse_transform()
         reversed_converter.collected_tensors = {
-            k: model.get_parameter_or_buffer(v) for k, v in reversed_converter.collected_tensors.items()
+            k: [model.get_parameter_or_buffer(param) for param in params]
+            for k, params in reversed_converter.collected_tensors.items()
         }
         first_param_name = next(iter(reversed_converter.layer_targets.keys()))
 

From ab055781bc117f8dd610eb8ef1501ae84ade62a7 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 17:19:48 +0100
Subject: [PATCH 06/36] fix renamed kwargs

---
 src/transformers/modeling_utils.py               | 2 +-
 src/transformers/quantizers/quantizer_torchao.py | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index c46461a87146..57a5003faed5 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3882,7 +3882,7 @@ def from_pretrained(
                 weight_conversions = get_checkpoint_conversion_mapping("legacy")
             if key_mapping is not None:
                 weight_conversions.extend(
-                    [WeightRenaming(source_keys=k, target_keys=v) for k, v in key_mapping.items()]
+                    [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()]
                 )
             if hf_quantizer is not None:
                 weight_conversions.extend(hf_quantizer.get_weight_conversions())
diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py
index 777ae193db0d..e7919b7f81b7 100644
--- a/src/transformers/quantizers/quantizer_torchao.py
+++ b/src/transformers/quantizers/quantizer_torchao.py
@@ -548,13 +548,13 @@ def get_weight_conversions(self):
         if self.pre_quantized:
             return [
                 WeightConverter(
-                    source_keys=["weight:qdata", "weight:scale", "weight:zero_point"],
-                    target_keys="weight",
+                    source_patterns=["weight:qdata", "weight:scale", "weight:zero_point"],
+                    target_patterns="weight",
                     operations=[TorchAoDeserialize(self)],
                 ),
                 WeightConverter(
-                    source_keys=["weight:_data"],
-                    target_keys="weight",
+                    source_patterns=["weight:_data"],
+                    target_patterns="weight",
                     operations=[TorchAoDeserialize(self)],
                 ),
                 # used for unsafe serialization

From 96880e0bf5f3dec99ad15a9ef6853209f4fd2f3f Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 18:51:46 +0100
Subject: [PATCH 07/36] fix timm_wrapper

---
 src/transformers/conversion_mapping.py        |  7 +++++++
 src/transformers/core_model_loading.py        |  6 +++++-
 src/transformers/modeling_utils.py            | 18 -----------------
 .../timm_wrapper/modeling_timm_wrapper.py     | 20 +------------------
 4 files changed, 13 insertions(+), 38 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index 02c837425bfe..6d1b45ed2b69 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -67,6 +67,13 @@ def _build_checkpoint_conversion_mapping():
                 operations=[MergeModulelist(dim=0)],
             ),
         ],
+        "timm_wrapper": [
+            # Simply add the prefix `timm_model`
+            WeightRenaming(
+                source_patterns=r"(.+)",
+                target_patterns=r"timm_model.\1",
+            )
+        ],
         "legacy": [
             WeightRenaming(
                 source_patterns="LayerNorm.gamma",
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 21b694c464c1..1757a67b910d 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -589,7 +589,11 @@ def repl(m, repl_map: dict[str, str]) -> str:
 
     # Exactly one match => return replacement
     name = matched_groups[0]
-    return repl_map[name]
+    replacement = repl_map[name]
+    # Allow capturing groups in patterns, i.e. to add a prefix to all keys (timm_wrapper)
+    if len(m.groups()) > 1:
+        return re.sub(rf"({m.group(1)})", replacement, m.group(0))
+    return replacement
 
 
 def convert_and_load_state_dict_in_model(
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 57a5003faed5..ae6fae729a05 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3145,9 +3145,6 @@ def save_pretrained(
                 if ignore_key in state_dict:
                     del state_dict[ignore_key]
 
-        # Rename state_dict keys before saving to file. Do nothing unless overridden in a particular model.
-        # (initially introduced with TimmWrapperModel to remove prefix and make checkpoints compatible with timm)
-        state_dict = self._fix_state_dict_keys_on_save(state_dict)
         # If model was sharded, we cannot properly determine sizes of tensors that `local_*` strategy was used,
         # therefore we replace them with DTensors that are equivalently sharded
         if self._tp_size is not None:
@@ -4026,21 +4023,6 @@ def from_pretrained(
             return model, loading_info
         return model
 
-    @staticmethod
-    def _fix_state_dict_key_on_save(key) -> tuple[str, bool]:
-        """
-        Similar to `_fix_state_dict_key_on_load` allows to define hook for state dict key renaming on model save.
-        Do nothing by default, but can be overridden in particular models.
-        """
-        return key, False
-
-    def _fix_state_dict_keys_on_save(self, state_dict):
-        """
-        Similar to `_fix_state_dict_keys_on_load` allows to define hook for state dict key renaming on model save.
-        Apply `_fix_state_dict_key_on_save` to all keys in `state_dict`.
-        """
-        return {self._fix_state_dict_key_on_save(key)[0]: value for key, value in state_dict.items()}
-
     @classmethod
     def _load_pretrained_model(
         cls,
diff --git a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
index aa541e0b15fb..e7dca5068ebc 100644
--- a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
+++ b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
@@ -98,30 +98,12 @@ def post_init(self):
         self.supports_gradient_checkpointing = self._timm_model_supports_gradient_checkpointing()
         super().post_init()
 
-    @staticmethod
-    def _fix_state_dict_key_on_load(key) -> tuple[str, bool]:
-        """
-        Overrides original method that renames `gamma` and `beta` to `weight` and `bias`.
-        We don't want this behavior for timm wrapped models. Instead, this method adds a
-        "timm_model." prefix to enable loading official timm Hub checkpoints.
-        """
-        if "timm_model." not in key:
-            return f"timm_model.{key}", True
-        return key, False
-
-    def _fix_state_dict_key_on_save(self, key):
-        """
-        Overrides original method to remove "timm_model." prefix from state_dict keys.
-        Makes the saved checkpoint compatible with the `timm` library.
-        """
-        return key.replace("timm_model.", ""), True
-
     def load_state_dict(self, state_dict, *args, **kwargs):
         """
         Override original method to fix state_dict keys on load for cases when weights are loaded
         without using the `from_pretrained` method (e.g., in Trainer to resume from checkpoint).
         """
-        state_dict = {self._fix_state_dict_key_on_load(k)[0]: v for k, v in state_dict.items()}
+        state_dict = {f"timm_model.{k}" if "timm_model." not in k else k: v for k, v in state_dict.items()}
         return super().load_state_dict(state_dict, *args, **kwargs)
 
     @torch.no_grad()

From a73101b04b262baa852920b4c8b6bcdedee3397b Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 18:59:21 +0100
Subject: [PATCH 08/36] remove fix_state_dict methods

---
 .../granite_speech/modeling_granite_speech.py      | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/src/transformers/models/granite_speech/modeling_granite_speech.py b/src/transformers/models/granite_speech/modeling_granite_speech.py
index b1c124931d46..524da4b16f1d 100644
--- a/src/transformers/models/granite_speech/modeling_granite_speech.py
+++ b/src/transformers/models/granite_speech/modeling_granite_speech.py
@@ -540,20 +540,6 @@ def save_pretrained(self, save_directory, *args, **kwargs):
         super().save_pretrained(save_directory, *args, **kwargs)
         self._hf_peft_config_loaded = prev_val
 
-    @staticmethod
-    def _fix_state_dict_key_on_save(key) -> tuple[str, bool]:
-        # save the model with the original weights format
-        return key.replace(".base_layer", ""), False
-
-    def _fix_state_dict_keys_on_save(self, state_dict):
-        if is_peft_available and self._hf_peft_config_loaded:
-            # state dict is only adapter, should keep the same
-            return state_dict
-        # rename back the base model state dict
-        return {
-            self._fix_state_dict_key_on_save(key)[0]: value for key, value in state_dict.items() if ".lora_" not in key
-        }
-
     def _get_adapter_name(self):
         return list(self.peft_config.keys())[0]
 

From eaf9ef374eecc9bf1223ebeeed1b8285f8ffaeab Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 19:54:28 +0100
Subject: [PATCH 09/36] can do it all the time, with __init__ as well

---
 src/transformers/core_model_loading.py | 130 +++++++++++++++++++------
 1 file changed, 102 insertions(+), 28 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 1757a67b910d..74354272ad3a 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -309,33 +309,36 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         bucket.add(source_key)
 
     def reset(self) -> None:
-        """Clean-up the collected tensors."""
+        """Clean-up the collected tensors to make sure we don't keep references to past tensors in memory."""
         self.collected_tensors = {}
 
     def reverse_transform(self) -> WeightTransform:
+        """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations."""
         if self.distributed_operation is not None or self.quantization_operation is not None:
             raise ValueError("Cannot reverse the transform with TP or quantization")
-        if len(self.layer_targets) == 0:
-            raise ValueError(
-                "You can only call `reverse_transform` after the `WeightTransform` instance has been populated"
-                "with the keys and tensors!"
-            )
+
         reverse_transform = self.__class__(source_patterns=self.target_patterns, target_patterns=self.source_patterns)
-        # Find the full names of the params we will need to use later for the reverse transform
-        reverse_layer_targets = defaultdict(set)
-        reverse_collected_tensors = defaultdict(list)
-        for target_key, all_sources in self.layer_targets.items():
-            matched_target_pattern = next(pat for pat in self.target_patterns if re.search(pat, target_key))
-            for source in all_sources:
-                reverse_layer_targets[source].add(target_key)
-                reverse_collected_tensors[matched_target_pattern].append(source)
-        reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
-        reverse_transform.layer_targets = reverse_layer_targets
-        reverse_transform.collected_tensors = reverse_collected_tensors
         # Add the reverse ops if applicable
         if hasattr(reverse_transform, "operations"):
             # All reverse ops, in reverse order
             reverse_transform.operations = [op.reverso_op for op in self.operations[::-1]]
+
+        # In this case, we already called `add_tensor` at least once, and we will use the saved keys to revert
+        # (i.e. we will be able to retrieve directly the correct params thanks to the `collected_tensors` which
+        # will hold the full param names entries)
+        if len(self.layer_targets) > 0:
+            # Find the full names of the params we will need to use later for the reverse transform
+            reverse_layer_targets = defaultdict(set)
+            reverse_collected_tensors = defaultdict(list)
+            for target_key, all_sources in self.layer_targets.items():
+                matched_target_pattern = next(pat for pat in self.target_patterns if re.search(pat, target_key))
+                for source in all_sources:
+                    reverse_layer_targets[source].add(target_key)
+                    reverse_collected_tensors[matched_target_pattern].append(source)
+            reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
+            reverse_transform.layer_targets = reverse_layer_targets
+            reverse_transform.collected_tensors = reverse_collected_tensors
+
         return reverse_transform
 
 
@@ -863,19 +866,90 @@ def convert_and_load_state_dict_in_model(
     return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, misc
 
 
-def revert_weight_conversion(model, state_dict):
-    original_mapping = getattr(model, "_weight_loading_mapping", None)
-    if original_mapping is None:
-        return state_dict
+def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
+    """ "
+    Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one
+    if the model was created in another way and is part of the default mappings.
+    """
+    conversion_mapping = getattr(model, "_weight_loading_mapping", None)
+    need_to_reverse = True
+    # In this case, the model was not created with `from_pretrained` -> let's check if it's in the hardcoded
+    # mappings, and recreate the mapping from there if it is
+    if conversion_mapping is None:
+        from .conversion_mapping import get_checkpoint_conversion_mapping
+        from .modeling_utils import VLMS
+
+        weight_conversions = []
+        # Hardcoded name mappings for some vlms
+        if any(
+            allowed_name in class_name.__name__.lower() for class_name in model.__mro__[:-1] for allowed_name in VLMS
+        ):
+            weight_conversions = [
+                WeightRenaming(source_patterns=k, target_patterns=v)
+                for k, v in model._checkpoint_conversion_mapping.items()
+            ]
+        # Hardcoded mapping for some models
+        # TODO: should be checked recursively on submodels, and similarly in `from_pretrained`
+        model_type = getattr(model.config, "model_type", None)
+        if model_type is not None:
+            model_conversions = get_checkpoint_conversion_mapping(model_type)
+            if model_conversions is not None:
+                weight_conversions.extend(model_conversions)
+
+        # We did not find any operations to perform -> quick escape
+        if len(weight_conversions) == 0:
+            return state_dict
+
+        conversion_mapping = {}
+        need_to_reverse = False
+
+        # Already 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}
+
+        # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
+        # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
+        rename_alt, _, rename_by_group = build_glob_alternation(renamings)
+        if converters != []:
+            weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = build_glob_alternation(converters)
+
+        state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
+        for original_key, tensor in state_dict:
+            # 1. apply all renamings
+            renamed_key = rename_alt.sub(lambda m: repl(m, rename_by_group), original_key).replace("\\", "")
+
+            # 2. apply 1 weight conversion on the key
+            matched_pattern = weight_pattern_alt.search(renamed_key) if converters != [] else None
+            if matched_pattern is not None:
+                renamed_key = weight_pattern_alt.sub(lambda m: repl(m, tgt_group_to_glob), renamed_key).replace(
+                    "\\", ""
+                )
+                new_converter = deepcopy(pattern_to_converter[src_group_to_glob[matched_pattern.lastgroup]])
+                # each target key gets its own converter instance
+                mapping = conversion_mapping.setdefault(renamed_key, new_converter)
+                source_pattern = src_group_to_glob[matched_pattern.lastgroup]
+            else:
+                mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(renamed_key, renamed_key))
+                source_pattern = renamed_key
+
+            mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
 
     new_state_dict = {}
-    for converter in original_mapping.values():
-        reversed_converter = converter.reverse_transform()
-        reversed_converter.collected_tensors = {
-            k: [model.get_parameter_or_buffer(param) for param in params]
-            for k, params in reversed_converter.collected_tensors.items()
-        }
-        first_param_name = next(iter(reversed_converter.layer_targets.keys()))
+    for first_param_name, converter in conversion_mapping.items():
+        # In this case, the mapping was obtained from the loaded model, we need to reverse the ops
+        if need_to_reverse:
+            reversed_converter = converter.reverse_transform()
+            reversed_converter.collected_tensors = {
+                k: [model.get_parameter_or_buffer(param) for param in params]
+                for k, params in reversed_converter.collected_tensors.items()
+            }
+            first_param_name = next(iter(reversed_converter.layer_targets.keys()))
+        # In this case, we just created the mapping above, the ops are already reversed
+        else:
+            reversed_converter = converter
 
         # Apply the reverse converter
         realized_value, misc = reversed_converter.convert(first_param_name, model=model, config=model.config)

From aa71dc45600de0c3ec3d4cf5fca31a5a68c1464d Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 19:56:02 +0100
Subject: [PATCH 10/36] doc

---
 src/transformers/core_model_loading.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 74354272ad3a..62ad1c7da9b9 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -942,6 +942,8 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
         # In this case, the mapping was obtained from the loaded model, we need to reverse the ops
         if need_to_reverse:
             reversed_converter = converter.reverse_transform()
+            # The `collected_tensors` only contain the full param keys after `reverse_transform` is called, we
+            # need to populate it with the actual tensors
             reversed_converter.collected_tensors = {
                 k: [model.get_parameter_or_buffer(param) for param in params]
                 for k, params in reversed_converter.collected_tensors.items()

From 04eeab9f5358ae807d364e7a1202365ae1cf2266 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 20:03:23 +0100
Subject: [PATCH 11/36] oupsi

---
 src/transformers/core_model_loading.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 62ad1c7da9b9..48294c376479 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -882,7 +882,9 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
         weight_conversions = []
         # Hardcoded name mappings for some vlms
         if any(
-            allowed_name in class_name.__name__.lower() for class_name in model.__mro__[:-1] for allowed_name in VLMS
+            allowed_name in class_name.__name__.lower()
+            for class_name in model.__class__.__mro__[:-1]
+            for allowed_name in VLMS
         ):
             weight_conversions = [
                 WeightRenaming(source_patterns=k, target_patterns=v)

From 67d9288946fb28b4ce76d21e6a3e3a62ac9f5e27 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 21:05:21 +0100
Subject: [PATCH 12/36] fix

---
 src/transformers/core_model_loading.py | 8 ++++----
 tests/utils/test_core_model_loading.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 48294c376479..359f956eedb9 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -761,8 +761,8 @@ def convert_and_load_state_dict_in_model(
                 mapping = param_name_to_load.setdefault(renamed_key, new_converter)
                 source_pattern = src_group_to_glob[matched_pattern.lastgroup]
             else:
-                mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(renamed_key, renamed_key))
-                source_pattern = renamed_key
+                mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
+                source_pattern = original_key
 
             # 5. Handle dtype casting
             if (
@@ -934,8 +934,8 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
                 mapping = conversion_mapping.setdefault(renamed_key, new_converter)
                 source_pattern = src_group_to_glob[matched_pattern.lastgroup]
             else:
-                mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(renamed_key, renamed_key))
-                source_pattern = renamed_key
+                mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
+                source_pattern = original_key
 
             mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
 
diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py
index 8973e9900f0f..883be48f6f60 100644
--- a/tests/utils/test_core_model_loading.py
+++ b/tests/utils/test_core_model_loading.py
@@ -1,4 +1,4 @@
-# Copyright 2024 HuggingFace Inc.
+# Copyright 2025 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From ac250748b0ec7083db31c11beaf8005f6ac5eb1c Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Tue, 25 Nov 2025 21:37:35 +0100
Subject: [PATCH 13/36] create helper

---
 src/transformers/conversion_mapping.py | 76 +++++++++++++++++++++++++-
 src/transformers/core_model_loading.py | 29 ++--------
 src/transformers/modeling_utils.py     | 49 ++---------------
 3 files changed, 84 insertions(+), 70 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index 6d1b45ed2b69..90d270ea7d74 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from copy import deepcopy
+from typing import TYPE_CHECKING
 
 from .core_model_loading import Concatenate, MergeModulelist, WeightConverter, WeightRenaming
 from .utils import is_torch_available
@@ -23,6 +24,11 @@
     import torch
 
 
+if TYPE_CHECKING:
+    from .modeling_utils import PreTrainedModel
+    from .quantizers import HfQuantizer
+
+
 def _build_checkpoint_conversion_mapping():
     mapping = {
         "mixtral": [
@@ -134,5 +140,71 @@ def _build_checkpoint_conversion_mapping():
 def get_checkpoint_conversion_mapping(model_type):
     global _checkpoint_conversion_mapping_cache
     _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping()
-    globals()["_checkpoint_conversion_mapping"] = _checkpoint_conversion_mapping_cache
-    return deepcopy(_checkpoint_conversion_mapping_cache.get(model_type, None))
+    return deepcopy(_checkpoint_conversion_mapping_cache.get(model_type))
+
+
+# DO NOT MODIFY, KEPT FOR BC ONLY
+VLMS = [
+    "aria",
+    "ayavision",
+    "colpali",
+    "emu3",
+    "fuyu",
+    "gotocr2",
+    "gemma3",
+    "internvl",
+    "llava",  # all llava prefixed models fall under this check
+    "mistral3",
+    "mllama",
+    "paligemma",
+    "shieldgemma2",
+    "qwen2vl",
+    "qwen2_5_vl",
+    "videollava",
+    "vipllava",
+    "sam3_video",
+    "sam3",
+    "sam3_tracker",
+    "sam3_tracker_video",
+]
+
+
+def get_model_conversion_mapping(
+    model: PreTrainedModel,
+    key_mapping: dict[str, str] | None = None,
+    hf_quantizer: HfQuantizer | None = None,
+    add_legacy: bool = True,
+) -> list[WeightConverter | WeightRenaming]:
+    """
+    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.
+    """
+    weight_conversions = []
+
+    # Load models with key mapping
+    if key_mapping is not None:
+        weight_conversions = [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()]
+    elif any(
+        allowed_name in class_name.__name__.lower()
+        for class_name in model.__class__.__mro__[:-1]
+        for allowed_name in VLMS
+    ):
+        weight_conversions = [
+            WeightRenaming(source_patterns=k, target_patterns=v)
+            for k, v in model._checkpoint_conversion_mapping.items()
+        ]
+
+    # TODO: should be checked recursively on submodels!!
+    model_type = getattr(model.config, "model_type", None)
+    if model_type is not None:
+        model_specific_conversions = get_checkpoint_conversion_mapping(model_type)
+        if model_specific_conversions is not None:
+            weight_conversions.extend(model_specific_conversions)
+        elif add_legacy:
+            weight_conversions.extend(get_checkpoint_conversion_mapping("legacy"))
+
+    # Add the ones from the quantizer as well if provided
+    if hf_quantizer is not None:
+        weight_conversions.extend(hf_quantizer.get_weight_conversions())
+
+    return weight_conversions
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 359f956eedb9..ff7aaee98ded 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -876,27 +876,10 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
     # In this case, the model was not created with `from_pretrained` -> let's check if it's in the hardcoded
     # mappings, and recreate the mapping from there if it is
     if conversion_mapping is None:
-        from .conversion_mapping import get_checkpoint_conversion_mapping
-        from .modeling_utils import VLMS
-
-        weight_conversions = []
-        # Hardcoded name mappings for some vlms
-        if any(
-            allowed_name in class_name.__name__.lower()
-            for class_name in model.__class__.__mro__[:-1]
-            for allowed_name in VLMS
-        ):
-            weight_conversions = [
-                WeightRenaming(source_patterns=k, target_patterns=v)
-                for k, v in model._checkpoint_conversion_mapping.items()
-            ]
-        # Hardcoded mapping for some models
-        # TODO: should be checked recursively on submodels, and similarly in `from_pretrained`
-        model_type = getattr(model.config, "model_type", None)
-        if model_type is not None:
-            model_conversions = get_checkpoint_conversion_mapping(model_type)
-            if model_conversions is not None:
-                weight_conversions.extend(model_conversions)
+        from .conversion_mapping import get_model_conversion_mapping
+
+        # Do not resave with the legacy renaming, if present
+        weight_conversions = get_model_conversion_mapping(model, add_legacy=False)
 
         # We did not find any operations to perform -> quick escape
         if len(weight_conversions) == 0:
@@ -920,10 +903,10 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
 
         state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
         for original_key, tensor in state_dict:
-            # 1. apply all renamings
+            # apply all renamings
             renamed_key = rename_alt.sub(lambda m: repl(m, rename_by_group), original_key).replace("\\", "")
 
-            # 2. apply 1 weight conversion on the key
+            # apply weight conversion on the key
             matched_pattern = weight_pattern_alt.search(renamed_key) if converters != [] else None
             if matched_pattern is not None:
                 renamed_key = weight_pattern_alt.sub(lambda m: repl(m, tgt_group_to_glob), renamed_key).replace(
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index ae6fae729a05..4acc87d3fa3e 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -46,7 +46,7 @@
 
 from . import initialization as init
 from .configuration_utils import PreTrainedConfig
-from .conversion_mapping import get_checkpoint_conversion_mapping
+from .conversion_mapping import get_model_conversion_mapping
 from .core_model_loading import (
     WeightConverter,
     WeightRenaming,
@@ -187,31 +187,6 @@ def is_local_dist_rank_0():
     "orthogonal_": nn.init.orthogonal_,
 }
 
-# DO NOT MODIFY, KEPT FOR BC ONLY
-VLMS = [
-    "aria",
-    "ayavision",
-    "colpali",
-    "emu3",
-    "fuyu",
-    "gotocr2",
-    "gemma3",
-    "internvl",
-    "llava",  # all llava prefixed models fall under this check
-    "mistral3",
-    "mllama",
-    "paligemma",
-    "shieldgemma2",
-    "qwen2vl",
-    "qwen2_5_vl",
-    "videollava",
-    "vipllava",
-    "sam3_video",
-    "sam3",
-    "sam3_tracker",
-    "sam3_tracker_video",
-]
-
 
 @contextmanager
 def no_init_weights():
@@ -3773,13 +3748,7 @@ def from_pretrained(
         trust_remote_code = kwargs.pop("trust_remote_code", None)
         use_kernels = kwargs.pop("use_kernels", False)
         kernel_config = kwargs.pop("kernel_config", None)
-
         key_mapping = kwargs.pop("key_mapping", None)
-        # Load models with key mapping
-        if key_mapping is None and any(
-            allowed_name in class_name.__name__.lower() for class_name in cls.__mro__[:-1] for allowed_name in VLMS
-        ):
-            key_mapping = copy.copy(cls._checkpoint_conversion_mapping)
 
         if distributed_config is not None and tp_plan is None:
             tp_plan = "auto"
@@ -3871,19 +3840,6 @@ def from_pretrained(
             config, quantization_config, dtype, device_map, weights_only, user_agent
         )
 
-        weight_conversions: Optional[list[WeightConverter | WeightRenaming]] = None
-        model_type = getattr(config, "model_type", None)
-        if model_type is not None:
-            weight_conversions = get_checkpoint_conversion_mapping(model_type)
-            if weight_conversions is None:
-                weight_conversions = get_checkpoint_conversion_mapping("legacy")
-            if key_mapping is not None:
-                weight_conversions.extend(
-                    [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()]
-                )
-            if hf_quantizer is not None:
-                weight_conversions.extend(hf_quantizer.get_weight_conversions())
-
         if gguf_file:
             if hf_quantizer is not None:
                 raise ValueError(
@@ -3939,6 +3895,9 @@ def from_pretrained(
             # Let's make sure we don't run the init function of buffer modules
             model = cls(config, *model_args, **model_kwargs)
 
+        # Obtain the weight conversion mapping for this model if any are registered
+        weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
+
         # make sure we use the model's config since the __init__ call might have copied it
         config = model.config
 

From 9df26279727ab51392972187de4fd05810caf407 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 10:26:39 +0100
Subject: [PATCH 14/36] fix annotation annoying isue

---
 src/transformers/conversion_mapping.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index 90d270ea7d74..e1cbeb2e460f 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from copy import deepcopy
 from typing import TYPE_CHECKING
 

From 2530a52b854c3dd599d7436ad65ba5137d279b4f Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 10:30:34 +0100
Subject: [PATCH 15/36] small fix

---
 src/transformers/core_model_loading.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index ff7aaee98ded..58b7ec1a8a0c 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -317,11 +317,15 @@ def reverse_transform(self) -> WeightTransform:
         if self.distributed_operation is not None or self.quantization_operation is not None:
             raise ValueError("Cannot reverse the transform with TP or quantization")
 
-        reverse_transform = self.__class__(source_patterns=self.target_patterns, target_patterns=self.source_patterns)
-        # Add the reverse ops if applicable
-        if hasattr(reverse_transform, "operations"):
+        kwargs = {}
+        # Add the reverse ops if applicable (it needs to be provided at __init__)
+        if hasattr(self, "operations"):
             # All reverse ops, in reverse order
-            reverse_transform.operations = [op.reverso_op for op in self.operations[::-1]]
+            kwargs["operations"] = [op.reverso_op for op in self.operations[::-1]]
+
+        reverse_transform = self.__class__(
+            source_patterns=self.target_patterns, target_patterns=self.source_patterns, **kwargs
+        )
 
         # In this case, we already called `add_tensor` at least once, and we will use the saved keys to revert
         # (i.e. we will be able to retrieve directly the correct params thanks to the `collected_tensors` which

From d996f7ca17b422ca465950eecacdcab47683c8ad Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 10:32:01 +0100
Subject: [PATCH 16/36] small fixes

---
 src/transformers/core_model_loading.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 58b7ec1a8a0c..70e4f236cbad 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -314,14 +314,15 @@ def reset(self) -> None:
 
     def reverse_transform(self) -> WeightTransform:
         """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations."""
-        if self.distributed_operation is not None or self.quantization_operation is not None:
+        # TODO: check this and relax when quantizer have `reverse_op`
+        if self.quantization_operation is not None:
             raise ValueError("Cannot reverse the transform with TP or quantization")
 
         kwargs = {}
         # Add the reverse ops if applicable (it needs to be provided at __init__)
         if hasattr(self, "operations"):
             # All reverse ops, in reverse order
-            kwargs["operations"] = [op.reverso_op for op in self.operations[::-1]]
+            kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]]
 
         reverse_transform = self.__class__(
             source_patterns=self.target_patterns, target_patterns=self.source_patterns, **kwargs

From 9a80e7bb268d124deddc5c8c912ba0605d42411b Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 14:55:43 +0100
Subject: [PATCH 17/36] alright commit all that already

---
 src/transformers/core_model_loading.py | 118 +++++++++++++++----------
 1 file changed, 73 insertions(+), 45 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 70e4f236cbad..14530c57e932 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -104,13 +104,16 @@ def build_glob_alternation(
 class ConversionOps:
     """Base class for weight conversion operations."""
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}(dim={self.dim})"
+
     @abstractmethod
     def convert(
         self,
         value: dict[str, Any],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
@@ -131,23 +134,20 @@ def __init__(self, dim: int = 0):
 
     def convert(
         self,
-        value: dict[str, list[torch.Tensor]],
+        value: dict[str, torch.Tensor],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
-    ) -> dict[str, list[torch.Tensor]]:
+    ) -> dict[str, torch.Tensor]:
         if len(value) > 1:
             raise ValueError("Unexpected value in `Chunk`!")
-        tensor = next(iter(value.values()))[0]
+        tensor = next(iter(value.values()))
         sizes = len(target_patterns)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
-        return {
-            full_layer_name.replace(target_patterns[0], target): [chunk]
-            for target, chunk in zip(target_patterns, chunks)
-        }
+        return {target: chunk.squeeze() for target, chunk in zip(target_patterns, chunks)}
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -166,17 +166,15 @@ def convert(
         value: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
-        if len(target_patterns) != 1:
+        if len(all_target_keys) != 1:
             raise ValueError("Concatenate expects a single target key.")
-        if len(value) != len(source_patterns):
-            raise ValueError("Concatenate received an unexpected number of tensors compared to source keys.")
 
-        return {full_layer_name: torch.cat(tuple(value.values()), dim=self.dim)}
+        return {all_target_keys[0]: torch.cat(tuple(value.values()), dim=self.dim)}
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -199,15 +197,18 @@ def convert(
         value: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
-        for key, tensors in value.items():
-            if len(source_patterns) == 1:
-                key = full_layer_name
+        for source_pattern, tensors in value.items():
+            # If only 1 source pattern, use the full target name in the result, otherwise keep the pattern
+            if len(value) == 1:
+                key = all_target_keys[0]
+            else:
+                key = source_pattern
             merged[key] = torch.stack(tensors, dim=self.dim)
         return merged
 
@@ -225,23 +226,24 @@ def __init__(self, dim: int = 0):
     @torch.no_grad
     def convert(
         self,
-        value: dict[str, list[torch.Tensor]],
+        value: dict[str, torch.Tensor],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
-    ) -> dict[str, list[torch.Tensor]]:
-        if len(value) > 1:
-            raise ValueError("Unexpected value in `SplitModulelist`!")
-        tensor = next(iter(value.values()))[0]
-        sizes = len(target_patterns)
-        chunks = torch.split(tensor, sizes, dim=self.dim)
-        return {
-            full_layer_name.replace(target_patterns[0], target): [chunk]
-            for target, chunk in zip(target_patterns, chunks)
-        }
+    ) -> dict[str, torch.Tensor]:
+        all_tensors = {}
+        for source_pattern, tensor in value.items():
+            sizes = len(all_target_keys) // len(value)
+            chunks = torch.chunk(tensor, sizes, dim=self.dim)
+            if len(value) > 1:
+                targets = [target for target in all_target_keys if re.search(source_pattern, target)]
+            else:
+                targets = all_target_keys
+            all_tensors.update({target: chunk.squeeze() for target, chunk in zip(targets, chunks)})
+        return all_tensors
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -270,7 +272,7 @@ def convert(
         value: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
-        full_layer_name: str,
+        all_target_keys: list[str],
         model,
         missing_keys,
         config,
@@ -292,8 +294,8 @@ class WeightTransform:
     distributed_operation: Optional[TensorParallelLayer] = None
     quantization_operation: Optional[ConversionOps] = None
 
-    collected_tensors: dict[str, list[Future]] = field(default_factory=dict, init=False)
-    layer_targets: dict[str, set[str]] = field(default_factory=dict, init=False)
+    collected_tensors: dict[str, list[Future]] = field(default_factory=lambda: defaultdict(list), init=False)
+    layer_targets: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set), init=False)
 
     def __post_init__(self):
         if isinstance(self.source_patterns, str):
@@ -302,15 +304,38 @@ def __post_init__(self):
             self.target_patterns = [self.target_patterns]
 
     def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):
-        bucket = self.collected_tensors.setdefault(source_pattern, [])
-        bucket += [future]
-
-        bucket = self.layer_targets.setdefault(target_key, set())
-        bucket.add(source_key)
+        self.collected_tensors[source_pattern].append(future)
+        self.layer_targets[target_key].add(source_key)
+
+        # For 1 to many operations (Chunk and SplitModuleList), we need to infer here all the tensors that we need
+        # to create from the source as `add_tensor` will only be called once with this given source for many targets
+        if (ops := getattr(self, "operations", None)) is not None:
+            # TODO: Here we assume this only happens during saving if the model was created from __init__, i.e.
+            # the future are actually Tensors, and we use heuristics to grab the sizes and the names
+            # This is brittle but works for the default mappings we have now
+            all_created_targets = []
+            if len(ops) == 2 and isinstance(ops[0], Chunk) and isinstance(ops[1], SplitModulelist):
+                all_created_targets = [
+                    re.sub(source_pattern, target_pattern, source_key) for target_pattern in self.target_patterns
+                ]
+                tensor = future
+                size = tensor.size(ops[1].dim)
+                all_created_targets = [
+                    target.replace("*", f"{i}") for i in range(size) for target in all_created_targets
+                ]
+            elif len(ops) == 1 and isinstance(ops[0], SplitModulelist):
+                tensor = future
+                size = tensor.size(ops[0].dim)
+                all_created_targets = [target_key.replace("*", f"{i}") for i in range(size)]
+            # Perform replacement if we took any of the above branches
+            if len(all_created_targets) > 0:
+                self.layer_targets.pop(target_key)
+                for target in all_created_targets:
+                    self.layer_targets[target].add(source_key)
 
     def reset(self) -> None:
         """Clean-up the collected tensors to make sure we don't keep references to past tensors in memory."""
-        self.collected_tensors = {}
+        self.collected_tensors = defaultdict(list)
 
     def reverse_transform(self) -> WeightTransform:
         """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations."""
@@ -337,9 +362,10 @@ def reverse_transform(self) -> WeightTransform:
             reverse_collected_tensors = defaultdict(list)
             for target_key, all_sources in self.layer_targets.items():
                 matched_target_pattern = next(pat for pat in self.target_patterns if re.search(pat, target_key))
+                reverse_collected_tensors[matched_target_pattern].append(target_key)
                 for source in all_sources:
                     reverse_layer_targets[source].add(target_key)
-                    reverse_collected_tensors[matched_target_pattern].append(source)
+                    # reverse_collected_tensors[matched_target_pattern].append(source)
             reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
             reverse_transform.layer_targets = reverse_layer_targets
             reverse_transform.collected_tensors = reverse_collected_tensors
@@ -366,13 +392,14 @@ def convert(
             )
 
         collected_tensors = self.collected_tensors
+        all_target_keys = sorted(self.layer_targets.keys(), key=dot_natural_key)
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     self.collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    full_layer_name=layer_name,
+                    all_target_keys=all_target_keys,
                     model=model,
                     config=config,
                     missing_keys=missing_keys,
@@ -409,13 +436,14 @@ def convert(
             )
 
         collected_tensors = self.collected_tensors
+        all_target_keys = sorted(self.layer_targets.keys(), key=dot_natural_key)
         for op in self.operations:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), op):
                 collected_tensors = op.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    full_layer_name=layer_name,
+                    all_target_keys=all_target_keys,
                     model=model,
                     config=config,
                     missing_keys=missing_keys,
@@ -426,7 +454,7 @@ def convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    full_layer_name=layer_name,
+                    all_target_keys=all_target_keys,
                     config=config,
                     model=model,
                     missing_keys=missing_keys,
@@ -599,8 +627,8 @@ def repl(m, repl_map: dict[str, str]) -> str:
     name = matched_groups[0]
     replacement = repl_map[name]
     # Allow capturing groups in patterns, i.e. to add a prefix to all keys (timm_wrapper)
-    if len(m.groups()) > 1:
-        return re.sub(rf"({m.group(1)})", replacement, m.group(0))
+    # if len(m.groups()) > 1:
+    #     return re.sub(rf"({m.group(1)})", replacement, m.group(0))
     return replacement
 
 

From 6b161e04bd1ed5361c42b89a3307d4ca5651385f Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 15:05:19 +0100
Subject: [PATCH 18/36] oupsi

---
 src/transformers/core_model_loading.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 14530c57e932..4c0bc1d4ec01 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -144,7 +144,8 @@ def convert(
     ) -> dict[str, torch.Tensor]:
         if len(value) > 1:
             raise ValueError("Unexpected value in `Chunk`!")
-        tensor = next(iter(value.values()))
+        tensors = next(iter(value.values()))
+        tensor = tensors[0] if isinstance(tensors, list) else tensors
         sizes = len(target_patterns)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
         return {target: chunk.squeeze() for target, chunk in zip(target_patterns, chunks)}
@@ -235,7 +236,8 @@ def convert(
         config,
     ) -> dict[str, torch.Tensor]:
         all_tensors = {}
-        for source_pattern, tensor in value.items():
+        for source_pattern, tensors in value.items():
+            tensor = tensors[0] if isinstance(tensors, list) else tensors
             sizes = len(all_target_keys) // len(value)
             chunks = torch.chunk(tensor, sizes, dim=self.dim)
             if len(value) > 1:

From d5cecba659ccf5275cb5213e425e0686d2b8ec0f Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 16:48:27 +0100
Subject: [PATCH 19/36] the fix

---
 src/transformers/core_model_loading.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 4c0bc1d4ec01..4ac9168337ba 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -388,13 +388,19 @@ def convert(
         missing_keys: Optional[MutableSet[str]] = None,
         misc: Optional[MutableMapping[str, str]] = None,
     ):
+        # Collect the tensor if using threading
         for pattern, futures in self.collected_tensors.items():
             self.collected_tensors[pattern] = (
                 futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]
             )
 
+        # Perform renaming op (for a simple WeightRenaming, `self.source_patterns` and `self.target_patterns` can
+        # only be of length 1, and are actually the full key names - we also have only 1 single related tensor)
+        target_key = self.target_patterns[0]
+        self.collected_tensors[target_key] = self.collected_tensors.pop(self.source_patterns[0])
+
+        all_target_keys = [target_key]
         collected_tensors = self.collected_tensors
-        all_target_keys = sorted(self.layer_targets.keys(), key=dot_natural_key)
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
@@ -432,6 +438,7 @@ def convert(
         missing_keys: Optional[MutableSet[str]] = None,
         misc: Optional[MutableMapping[str, str]] = None,
     ):
+        # Collect all tensors if using threading
         for pattern, futures in self.collected_tensors.items():
             self.collected_tensors[pattern] = (
                 futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]

From 17a86280b8f1b9ff3905ed99bc273e9c2330d369 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 17:03:36 +0100
Subject: [PATCH 20/36] update quantizers

---
 src/transformers/integrations/finegrained_fp8.py | 2 --
 src/transformers/integrations/mxfp4.py           | 9 ++++++---
 src/transformers/integrations/torchao.py         | 8 +++++---
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py
index 42854324bc0d..8c2a8b9249cb 100644
--- a/src/transformers/integrations/finegrained_fp8.py
+++ b/src/transformers/integrations/finegrained_fp8.py
@@ -577,8 +577,6 @@ class Fp8Quantize(ConversionOps):
     A quantization operation that creates two tensors, weight and scale out of a weight.
     """
 
-    reverse_op: type[ConversionOps]
-
     def __init__(self, hf_quantizer):
         self.hf_quantizer = hf_quantizer
         self.reverse_op = Fp8Dequantize
diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py
index 6430c0d9d57d..c6e131a5726a 100644
--- a/src/transformers/integrations/mxfp4.py
+++ b/src/transformers/integrations/mxfp4.py
@@ -85,12 +85,13 @@ def convert(
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
         missing_keys: Optional[list[str]] = None,
-        full_layer_name: str | None = None,
+        all_target_keys: list[str] | None = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         _, value = tuple(input_dict.items())[0]
         value = value[0] if isinstance(value, list) else value
 
+        full_layer_name = all_target_keys[0]
         module, _ = get_module_from_name(model, full_layer_name)
 
         with torch.device(value.device):
@@ -132,7 +133,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        full_layer_name: str | None = None,
+        all_target_keys: list[str] | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -150,6 +151,7 @@ def convert(
 
         # Here we are dequantizing the weights
         dequantized = dequantize_convertops(param_data["_blocks"], param_data["_scales"], param_data["_blocks"].device)
+        full_layer_name = all_target_keys[0]
         return {full_layer_name: dequantized}
 
 
@@ -161,7 +163,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        full_layer_name: str | None = None,
+        all_target_keys: list[str] | None = None,
         missing_keys: Optional[list[str]] = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -178,6 +180,7 @@ def convert(
                 param_data["_scales"] = input_dict["_scales"]
 
         # Eagerly set tensors on the module and perform swizzle
+        full_layer_name = all_target_keys[0]
         module, _ = get_module_from_name(model, full_layer_name)
         proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
         swizzle_mxfp4_convertops(
diff --git a/src/transformers/integrations/torchao.py b/src/transformers/integrations/torchao.py
index 3a1fdb0d407e..155446621f24 100644
--- a/src/transformers/integrations/torchao.py
+++ b/src/transformers/integrations/torchao.py
@@ -83,7 +83,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        full_layer_name: str | None = None,
+        all_target_keys: list[str] | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -92,9 +92,10 @@ def convert(
         _, value = tuple(input_dict.items())[0]
         value = value[0] if isinstance(value, list) else value
 
+        full_layer_name = all_target_keys[0]
         module, tensor_name = get_module_from_name(model, full_layer_name)
 
-        module._parameters[tensor_name] = torch.nn.Parameter(value, requires_grad=value.requires_grad).to(value.device)
+        module._parameters[tensor_name] = torch.nn.Parameter(value, requires_grad=value.requires_grad)
         # if we are quantizing tied parameters, to avoid tying the quantized weights
         # the correct order to do it is
         # 1. load the weight to model
@@ -211,10 +212,11 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        full_layer_name: str | None = None,
+        all_target_keys: list[str] | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
+        full_layer_name = all_target_keys[0]
         if isinstance(self.hf_quantizer.quantization_config.quant_type, str):
             is_int_4 = "int4" in self.hf_quantizer.quantization_config.quant_type
         else:

From 5b4e6951fa8cbfab43ff40b18bea8212f7e32341 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 17:12:09 +0100
Subject: [PATCH 21/36] this works

---
 src/transformers/core_model_loading.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 4ac9168337ba..61a8e52c36a2 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -397,14 +397,13 @@ def convert(
         # Perform renaming op (for a simple WeightRenaming, `self.source_patterns` and `self.target_patterns` can
         # only be of length 1, and are actually the full key names - we also have only 1 single related tensor)
         target_key = self.target_patterns[0]
-        self.collected_tensors[target_key] = self.collected_tensors.pop(self.source_patterns[0])
+        collected_tensors = {target_key: self.collected_tensors[self.source_patterns[0]]}
 
         all_target_keys = [target_key]
-        collected_tensors = self.collected_tensors
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
-                    self.collected_tensors,
+                    collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
                     all_target_keys=all_target_keys,

From a4f2cb56b4669ab999066b54d15393a2335e9208 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 18:30:13 +0100
Subject: [PATCH 22/36] the hardcoded regex got me hard....

---
 src/transformers/core_model_loading.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 61a8e52c36a2..c846d955a826 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -634,9 +634,16 @@ def repl(m, repl_map: dict[str, str]) -> str:
     # Exactly one match => return replacement
     name = matched_groups[0]
     replacement = repl_map[name]
+    # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
+    if replacement.startswith("^"):
+        replacement = replacement[1:]
+    # This is ugly but needed for reverse mapping of Qwen2.5!
+    if "(?!\.(language_model|visual))" in replacement:
+        replacement = replacement.replace("(?!\.(language_model|visual))", "")
+
     # Allow capturing groups in patterns, i.e. to add a prefix to all keys (timm_wrapper)
     # if len(m.groups()) > 1:
-    #     return re.sub(rf"({m.group(1)})", replacement, m.group(0))
+        # return re.sub(rf"({m.group(1)})", replacement, m.group(0))
     return replacement
 
 

From faa2179b4bd14326b84aab2db2261cb22bf25566 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 18:30:27 +0100
Subject: [PATCH 23/36] style

---
 src/transformers/core_model_loading.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index c846d955a826..d707de00427f 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -635,15 +635,14 @@ def repl(m, repl_map: dict[str, str]) -> str:
     name = matched_groups[0]
     replacement = repl_map[name]
     # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
-    if replacement.startswith("^"):
-        replacement = replacement[1:]
+    replacement = replacement.removeprefix("^")
     # This is ugly but needed for reverse mapping of Qwen2.5!
-    if "(?!\.(language_model|visual))" in replacement:
-        replacement = replacement.replace("(?!\.(language_model|visual))", "")
+    if r"(?!\.(language_model|visual))" in replacement:
+        replacement = replacement.replace(r"(?!\.(language_model|visual))", "")
 
     # Allow capturing groups in patterns, i.e. to add a prefix to all keys (timm_wrapper)
     # if len(m.groups()) > 1:
-        # return re.sub(rf"({m.group(1)})", replacement, m.group(0))
+    # return re.sub(rf"({m.group(1)})", replacement, m.group(0))
     return replacement
 
 

From 5c64084865281987dc2ce5afe66061683768f00b Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 18:53:50 +0100
Subject: [PATCH 24/36] the final one

---
 src/transformers/core_model_loading.py | 8 ++++----
 tests/utils/test_core_model_loading.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index d707de00427f..38ca5f7bfe7b 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -148,7 +148,7 @@ def convert(
         tensor = tensors[0] if isinstance(tensors, list) else tensors
         sizes = len(target_patterns)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
-        return {target: chunk.squeeze() for target, chunk in zip(target_patterns, chunks)}
+        return dict(zip(target_patterns, chunks))
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -639,10 +639,10 @@ def repl(m, repl_map: dict[str, str]) -> str:
     # This is ugly but needed for reverse mapping of Qwen2.5!
     if r"(?!\.(language_model|visual))" in replacement:
         replacement = replacement.replace(r"(?!\.(language_model|visual))", "")
+    # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper)
+    if "\1" in replacement and len(m.groups()) > 1:
+        replacement = replacement.replace("\1", m.group(1))
 
-    # Allow capturing groups in patterns, i.e. to add a prefix to all keys (timm_wrapper)
-    # if len(m.groups()) > 1:
-    # return re.sub(rf"({m.group(1)})", replacement, m.group(0))
     return replacement
 
 
diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py
index 883be48f6f60..3943238c0412 100644
--- a/tests/utils/test_core_model_loading.py
+++ b/tests/utils/test_core_model_loading.py
@@ -246,7 +246,7 @@ def test_moe_and_qkv_conversion(self):
                     "model.layers.0.self_attn.k_proj.weight",
                     "model.layers.0.self_attn.v_proj.weight",
                 ],
-                operations=[Chunk(dim=0, chunks=3)],
+                operations=[Chunk(dim=0)],
             ),
             WeightRenaming("mlp.w2.weight", "mlp.down_proj.weight"),
         ]

From 30bae5f3fa62106bee8b288a2885b7132ac7f3b0 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 18:58:41 +0100
Subject: [PATCH 25/36] cleanup a bit

---
 src/transformers/core_model_loading.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 38ca5f7bfe7b..598d42afd81a 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -313,7 +313,7 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         # to create from the source as `add_tensor` will only be called once with this given source for many targets
         if (ops := getattr(self, "operations", None)) is not None:
             # TODO: Here we assume this only happens during saving if the model was created from __init__, i.e.
-            # the future are actually Tensors, and we use heuristics to grab the sizes and the names
+            # the Futures are actually Tensors, and we use heuristics to grab the sizes and the names
             # This is brittle but works for the default mappings we have now
             all_created_targets = []
             if len(ops) == 2 and isinstance(ops[0], Chunk) and isinstance(ops[1], SplitModulelist):
@@ -367,7 +367,6 @@ def reverse_transform(self) -> WeightTransform:
                 reverse_collected_tensors[matched_target_pattern].append(target_key)
                 for source in all_sources:
                     reverse_layer_targets[source].add(target_key)
-                    # reverse_collected_tensors[matched_target_pattern].append(source)
             reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
             reverse_transform.layer_targets = reverse_layer_targets
             reverse_transform.collected_tensors = reverse_collected_tensors

From 58cc5b4d9479d15caef2d41d96e9ff369a88517a Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 19:32:19 +0100
Subject: [PATCH 26/36] better

---
 src/transformers/conversion_mapping.py |  1 +
 src/transformers/core_model_loading.py | 33 +++++++++++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index e1cbeb2e460f..033f22d8867b 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -77,6 +77,7 @@ def _build_checkpoint_conversion_mapping():
         ],
         "timm_wrapper": [
             # Simply add the prefix `timm_model`
+            # TODO: Would be probably much cleaner with a `add_prefix` argument in WeightRenaming
             WeightRenaming(
                 source_patterns=r"(.+)",
                 target_patterns=r"timm_model.\1",
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 598d42afd81a..1373f91519f1 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -305,6 +305,27 @@ def __post_init__(self):
         if isinstance(self.target_patterns, str):
             self.target_patterns = [self.target_patterns]
 
+        # Due to how our `_checkpoint_conversion_mapping` mappings are written, we need 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 targets for the reversed mapping
+        for i, pattern in enumerate(self.target_patterns):
+            # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
+            pattern = pattern.removeprefix("^")
+            # This is ugly but needed for reverse mapping of Qwen2.5!
+            if r"(?!\.(language_model|visual))" in pattern:
+                pattern = pattern.replace(r"(?!\.(language_model|visual))", "")
+            # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper)
+            if r"(.+)" in pattern:
+                pattern = pattern.replace(r"(.+)", "")
+
+            self.target_patterns[i] = pattern
+
+        # We also need to check capturing groups in the sources during reverse mapping (e.g. timm_wrapper)
+        for i, pattern in enumerate(self.source_patterns):
+            if r"\1" in pattern:
+                pattern = pattern.replace(r"\1", "")
+            self.source_patterns[i] = pattern
+
     def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):
         self.collected_tensors[source_pattern].append(future)
         self.layer_targets[target_key].add(source_key)
@@ -313,7 +334,7 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         # to create from the source as `add_tensor` will only be called once with this given source for many targets
         if (ops := getattr(self, "operations", None)) is not None:
             # TODO: Here we assume this only happens during saving if the model was created from __init__, i.e.
-            # the Futures are actually Tensors, and we use heuristics to grab the sizes and the names
+            # the future are actually Tensors, and we use heuristics to grab the sizes and the names
             # This is brittle but works for the default mappings we have now
             all_created_targets = []
             if len(ops) == 2 and isinstance(ops[0], Chunk) and isinstance(ops[1], SplitModulelist):
@@ -367,6 +388,7 @@ def reverse_transform(self) -> WeightTransform:
                 reverse_collected_tensors[matched_target_pattern].append(target_key)
                 for source in all_sources:
                     reverse_layer_targets[source].add(target_key)
+                    # reverse_collected_tensors[matched_target_pattern].append(source)
             reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
             reverse_transform.layer_targets = reverse_layer_targets
             reverse_transform.collected_tensors = reverse_collected_tensors
@@ -633,14 +655,9 @@ def repl(m, repl_map: dict[str, str]) -> str:
     # Exactly one match => return replacement
     name = matched_groups[0]
     replacement = repl_map[name]
-    # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
-    replacement = replacement.removeprefix("^")
-    # This is ugly but needed for reverse mapping of Qwen2.5!
-    if r"(?!\.(language_model|visual))" in replacement:
-        replacement = replacement.replace(r"(?!\.(language_model|visual))", "")
     # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper)
-    if "\1" in replacement and len(m.groups()) > 1:
-        replacement = replacement.replace("\1", m.group(1))
+    if r"\1" in replacement and len(m.groups()) > 1:
+        replacement = replacement.replace(r"\1", m.group(1))
 
     return replacement
 

From 2ef739990ab5076576823b52c8e395f7a7cb2d8e Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 19:33:52 +0100
Subject: [PATCH 27/36] style

---
 src/transformers/core_model_loading.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 1373f91519f1..ef9203e82a91 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -317,7 +317,6 @@ def __post_init__(self):
             # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper)
             if r"(.+)" in pattern:
                 pattern = pattern.replace(r"(.+)", "")
-
             self.target_patterns[i] = pattern
 
         # We also need to check capturing groups in the sources during reverse mapping (e.g. timm_wrapper)

From 20e6927a90444ecc83181e28a964d7004c12427b Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Wed, 26 Nov 2025 19:34:50 +0100
Subject: [PATCH 28/36] oupsi readded it

---
 src/transformers/core_model_loading.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index ef9203e82a91..689a58ebc391 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -333,7 +333,7 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         # to create from the source as `add_tensor` will only be called once with this given source for many targets
         if (ops := getattr(self, "operations", None)) is not None:
             # TODO: Here we assume this only happens during saving if the model was created from __init__, i.e.
-            # the future are actually Tensors, and we use heuristics to grab the sizes and the names
+            # the Futures are actually Tensors, and we use heuristics to grab the sizes and the names
             # This is brittle but works for the default mappings we have now
             all_created_targets = []
             if len(ops) == 2 and isinstance(ops[0], Chunk) and isinstance(ops[1], SplitModulelist):
@@ -387,7 +387,6 @@ def reverse_transform(self) -> WeightTransform:
                 reverse_collected_tensors[matched_target_pattern].append(target_key)
                 for source in all_sources:
                     reverse_layer_targets[source].add(target_key)
-                    # reverse_collected_tensors[matched_target_pattern].append(source)
             reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
             reverse_transform.layer_targets = reverse_layer_targets
             reverse_transform.collected_tensors = reverse_collected_tensors

From b1a6621431bbeb52b11ad95742d668c5bc49fa85 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 13:52:45 +0100
Subject: [PATCH 29/36] do it inside the ops instead - no need for full names
 anymore

---
 src/transformers/core_model_loading.py | 119 ++++++++++++++-----------
 1 file changed, 66 insertions(+), 53 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 689a58ebc391..274133f22ba8 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -110,7 +110,7 @@ def __repr__(self):
     @abstractmethod
     def convert(
         self,
-        value: dict[str, Any],
+        input_dict: dict[str, Any],
         source_patterns: list[str],
         target_patterns: list[str],
         all_target_keys: list[str],
@@ -134,7 +134,7 @@ def __init__(self, dim: int = 0):
 
     def convert(
         self,
-        value: dict[str, torch.Tensor],
+        input_dict: dict[str, torch.Tensor],
         source_patterns: list[str],
         target_patterns: list[str],
         all_target_keys: list[str],
@@ -142,13 +142,18 @@ def convert(
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
-        if len(value) > 1:
-            raise ValueError("Unexpected value in `Chunk`!")
-        tensors = next(iter(value.values()))
+        tensors = next(iter(input_dict.values()))
         tensor = tensors[0] if isinstance(tensors, list) else tensors
-        sizes = len(target_patterns)
+        targets = self.get_target_pattern(input_dict, target_patterns)
+        sizes = len(targets)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
-        return dict(zip(target_patterns, chunks))
+        return dict(zip(targets, chunks))
+
+    def get_target_pattern(self, input_dict: dict, target_patterns: list[str]) -> list[str]:
+        # Here we always return the target patterns
+        if len(input_dict) > 1 or len(target_patterns) == 1:
+            raise ValueError("Undefined Operation encountered!")
+        return target_patterns
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -164,7 +169,7 @@ def __init__(self, dim: int = 0):
     @torch.no_grad
     def convert(
         self,
-        value: dict[str, list[torch.Tensor]],
+        input_dict: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
         all_target_keys: list[str],
@@ -172,10 +177,14 @@ def convert(
         missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
-        if len(all_target_keys) != 1:
-            raise ValueError("Concatenate expects a single target key.")
+        target_pattern = self.get_target_pattern(target_patterns)
+        return {target_pattern: torch.cat(tuple(input_dict.values()), dim=self.dim)}
 
-        return {all_target_keys[0]: torch.cat(tuple(value.values()), dim=self.dim)}
+    def get_target_pattern(self, target_patterns: list[str]) -> str:
+        # Here we always return the target pattern
+        if len(target_patterns) > 1:
+            raise ValueError("Undefined Operation encountered!")
+        return target_patterns[0]
 
     @property
     def reverse_op(self) -> ConversionOps:
@@ -195,7 +204,7 @@ def __init__(self, dim: int = 0):
     @torch.no_grad
     def convert(
         self,
-        value: dict[str, list[torch.Tensor]],
+        input_dict: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
         all_target_keys: list[str],
@@ -204,15 +213,22 @@ def convert(
         config,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
-        for source_pattern, tensors in value.items():
-            # If only 1 source pattern, use the full target name in the result, otherwise keep the pattern
-            if len(value) == 1:
-                key = all_target_keys[0]
-            else:
-                key = source_pattern
-            merged[key] = torch.stack(tensors, dim=self.dim)
+        for source_pattern, tensors in input_dict.items():
+            target_pattern = self.get_target_pattern(input_dict, source_pattern, target_patterns)
+            merged[target_pattern] = torch.stack(tensors, dim=self.dim)
         return merged
 
+    def get_target_pattern(self, input_dict: dict, source_pattern: str, target_patterns: list[str]) -> str:
+        # Here it's a single operation, so we use the target
+        if len(input_dict) == 1:
+            if len(target_patterns) == 1:
+                return target_patterns[0]
+            else:
+                raise ValueError("Undefined Operation encountered!")
+        #  Here it's the first operation in a chain, so we use the source as they were replaced before in the chain
+        else:
+            return source_pattern
+
     @property
     def reverse_op(self) -> ConversionOps:
         return SplitModulelist(self.dim)
@@ -227,7 +243,7 @@ def __init__(self, dim: int = 0):
     @torch.no_grad
     def convert(
         self,
-        value: dict[str, torch.Tensor],
+        input_dict: dict[str, torch.Tensor],
         source_patterns: list[str],
         target_patterns: list[str],
         all_target_keys: list[str],
@@ -236,17 +252,29 @@ def convert(
         config,
     ) -> dict[str, torch.Tensor]:
         all_tensors = {}
-        for source_pattern, tensors in value.items():
+        for source_pattern, tensors in input_dict.items():
             tensor = tensors[0] if isinstance(tensors, list) else tensors
-            sizes = len(all_target_keys) // len(value)
+            # We split in the number of tensors present in the given dim
+            sizes = tensor.size(self.dim)
+            targets = self.get_target_patterns(input_dict, source_pattern, target_patterns, sizes)
             chunks = torch.chunk(tensor, sizes, dim=self.dim)
-            if len(value) > 1:
-                targets = [target for target in all_target_keys if re.search(source_pattern, target)]
-            else:
-                targets = all_target_keys
+            # We squeeze each chunk here as well to make sure to give them their original shape
             all_tensors.update({target: chunk.squeeze() for target, chunk in zip(targets, chunks)})
         return all_tensors
 
+    def get_target_patterns(
+        self, input_dict: dict, source_pattern: str, target_patterns: list[str], sizes: int
+    ) -> list[str]:
+        # Here it's a single operation, so we use the target
+        if len(input_dict) == 1:
+            if len(target_patterns) == 1:
+                return [target_patterns[0].replace("*", f"{i}") for i in range(sizes)]
+            else:
+                raise ValueError("Undefined Operation encountered!")
+        # Here it's the last operation in a chain, so we use the source as they were replaced before in the chain
+        else:
+            return [source_pattern.replace("*", f"{i}") for i in range(sizes)]
+
     @property
     def reverse_op(self) -> ConversionOps:
         return MergeModulelist(self.dim)
@@ -329,32 +357,6 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
         self.collected_tensors[source_pattern].append(future)
         self.layer_targets[target_key].add(source_key)
 
-        # For 1 to many operations (Chunk and SplitModuleList), we need to infer here all the tensors that we need
-        # to create from the source as `add_tensor` will only be called once with this given source for many targets
-        if (ops := getattr(self, "operations", None)) is not None:
-            # TODO: Here we assume this only happens during saving if the model was created from __init__, i.e.
-            # the Futures are actually Tensors, and we use heuristics to grab the sizes and the names
-            # This is brittle but works for the default mappings we have now
-            all_created_targets = []
-            if len(ops) == 2 and isinstance(ops[0], Chunk) and isinstance(ops[1], SplitModulelist):
-                all_created_targets = [
-                    re.sub(source_pattern, target_pattern, source_key) for target_pattern in self.target_patterns
-                ]
-                tensor = future
-                size = tensor.size(ops[1].dim)
-                all_created_targets = [
-                    target.replace("*", f"{i}") for i in range(size) for target in all_created_targets
-                ]
-            elif len(ops) == 1 and isinstance(ops[0], SplitModulelist):
-                tensor = future
-                size = tensor.size(ops[0].dim)
-                all_created_targets = [target_key.replace("*", f"{i}") for i in range(size)]
-            # Perform replacement if we took any of the above branches
-            if len(all_created_targets) > 0:
-                self.layer_targets.pop(target_key)
-                for target in all_created_targets:
-                    self.layer_targets[target].add(source_key)
-
     def reset(self) -> None:
         """Clean-up the collected tensors to make sure we don't keep references to past tensors in memory."""
         self.collected_tensors = defaultdict(list)
@@ -418,6 +420,7 @@ def convert(
         target_key = self.target_patterns[0]
         collected_tensors = {target_key: self.collected_tensors[self.source_patterns[0]]}
 
+        # TODO: `all_target_keys` should not be needed, clean it up later in quantization
         all_target_keys = [target_key]
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
@@ -463,7 +466,8 @@ def convert(
             )
 
         collected_tensors = self.collected_tensors
-        all_target_keys = sorted(self.layer_targets.keys(), key=dot_natural_key)
+        # TODO: `all_target_keys` should not be needed, clean it up later in quantization
+        all_target_keys = [layer_name]
         for op in self.operations:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), op):
                 collected_tensors = op.convert(
@@ -475,6 +479,15 @@ def convert(
                     config=config,
                     missing_keys=missing_keys,
                 )
+
+        # Tensors are returned from ops with the target patterns, we need to expand them to full name.
+        # This means we need to grab the prefix and suffix to add to every target key
+        if ".*." in layer_name:
+            layer_name = layer_name.replace(".*.", ".0.")
+        prefix, _, suffix = next(layer_name.partition(k) for k in collected_tensors.keys() if k in layer_name)
+        # Rename the tensors
+        collected_tensors = {prefix + k + suffix: v for k, v in collected_tensors.items()}
+
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(

From 20bcbdd5724ec29f9f053bae66892d103dc693ba Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 14:08:23 +0100
Subject: [PATCH 30/36] reverse quantizers and simplify signatures

---
 src/transformers/core_model_loading.py   | 60 +++++-------------------
 src/transformers/integrations/mxfp4.py   |  9 ++--
 src/transformers/integrations/torchao.py |  6 +--
 3 files changed, 16 insertions(+), 59 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 274133f22ba8..ac372c5d07ab 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -109,15 +109,7 @@ def __repr__(self):
 
     @abstractmethod
     def convert(
-        self,
-        input_dict: dict[str, Any],
-        source_patterns: list[str],
-        target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
-        config,
-        **kwargs,
+        self, input_dict: dict[str, Any], source_patterns: list[str], target_patterns: list[str], **kwargs
     ) -> dict[str, list[torch.Tensor]]:
         raise NotImplementedError
 
@@ -133,14 +125,7 @@ def __init__(self, dim: int = 0):
         self.dim = dim
 
     def convert(
-        self,
-        input_dict: dict[str, torch.Tensor],
-        source_patterns: list[str],
-        target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
-        config,
+        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
     ) -> dict[str, torch.Tensor]:
         tensors = next(iter(input_dict.values()))
         tensor = tensors[0] if isinstance(tensors, list) else tensors
@@ -172,10 +157,7 @@ def convert(
         input_dict: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
-        config,
+        **kwargs,
     ) -> dict[str, torch.Tensor]:
         target_pattern = self.get_target_pattern(target_patterns)
         return {target_pattern: torch.cat(tuple(input_dict.values()), dim=self.dim)}
@@ -203,14 +185,7 @@ def __init__(self, dim: int = 0):
 
     @torch.no_grad
     def convert(
-        self,
-        input_dict: dict[str, list[torch.Tensor]],
-        source_patterns: list[str],
-        target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
-        config,
+        self, input_dict: dict[str, list[torch.Tensor]], source_patterns: list[str], target_patterns: list[str]
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
         for source_pattern, tensors in input_dict.items():
@@ -242,14 +217,7 @@ def __init__(self, dim: int = 0):
 
     @torch.no_grad
     def convert(
-        self,
-        input_dict: dict[str, torch.Tensor],
-        source_patterns: list[str],
-        target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
-        config,
+        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
     ) -> dict[str, torch.Tensor]:
         all_tensors = {}
         for source_pattern, tensors in input_dict.items():
@@ -299,17 +267,15 @@ def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
     @torch.no_grad
     def convert(
         self,
-        value: dict[str, list[torch.Tensor]],
+        input_dict: dict[str, list[torch.Tensor]],
         source_patterns: list[str],
         target_patterns: list[str],
-        all_target_keys: list[str],
-        model,
-        missing_keys,
         config,
+        **kwargs,
     ) -> dict[str, list[torch.Tensor]]:
         self.config = config
         output: dict[str, list[torch.Tensor]] = {}
-        for key, tensors in value.items():
+        for key, tensors in input_dict.items():
             if len(tensors) != 1:
                 raise ValueError("PermuteForRope expects a single tensor per key.")
             output[key] = [self._apply(tensors[0])]
@@ -420,15 +386,13 @@ def convert(
         target_key = self.target_patterns[0]
         collected_tensors = {target_key: self.collected_tensors[self.source_patterns[0]]}
 
-        # TODO: `all_target_keys` should not be needed, clean it up later in quantization
-        all_target_keys = [target_key]
         if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    all_target_keys=all_target_keys,
+                    full_layer_name=target_key,
                     model=model,
                     config=config,
                     missing_keys=missing_keys,
@@ -466,15 +430,13 @@ def convert(
             )
 
         collected_tensors = self.collected_tensors
-        # TODO: `all_target_keys` should not be needed, clean it up later in quantization
-        all_target_keys = [layer_name]
         for op in self.operations:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), op):
                 collected_tensors = op.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    all_target_keys=all_target_keys,
+                    # Additional kwargs, ususally not used
                     model=model,
                     config=config,
                     missing_keys=missing_keys,
@@ -494,7 +456,7 @@ def convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
                     target_patterns=self.target_patterns,
-                    all_target_keys=all_target_keys,
+                    full_layer_name=layer_name,
                     config=config,
                     model=model,
                     missing_keys=missing_keys,
diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py
index c6e131a5726a..6430c0d9d57d 100644
--- a/src/transformers/integrations/mxfp4.py
+++ b/src/transformers/integrations/mxfp4.py
@@ -85,13 +85,12 @@ def convert(
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
         missing_keys: Optional[list[str]] = None,
-        all_target_keys: list[str] | None = None,
+        full_layer_name: str | None = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         _, value = tuple(input_dict.items())[0]
         value = value[0] if isinstance(value, list) else value
 
-        full_layer_name = all_target_keys[0]
         module, _ = get_module_from_name(model, full_layer_name)
 
         with torch.device(value.device):
@@ -133,7 +132,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        all_target_keys: list[str] | None = None,
+        full_layer_name: str | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -151,7 +150,6 @@ def convert(
 
         # Here we are dequantizing the weights
         dequantized = dequantize_convertops(param_data["_blocks"], param_data["_scales"], param_data["_blocks"].device)
-        full_layer_name = all_target_keys[0]
         return {full_layer_name: dequantized}
 
 
@@ -163,7 +161,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        all_target_keys: list[str] | None = None,
+        full_layer_name: str | None = None,
         missing_keys: Optional[list[str]] = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -180,7 +178,6 @@ def convert(
                 param_data["_scales"] = input_dict["_scales"]
 
         # Eagerly set tensors on the module and perform swizzle
-        full_layer_name = all_target_keys[0]
         module, _ = get_module_from_name(model, full_layer_name)
         proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj"
         swizzle_mxfp4_convertops(
diff --git a/src/transformers/integrations/torchao.py b/src/transformers/integrations/torchao.py
index 155446621f24..22a776a7ec74 100644
--- a/src/transformers/integrations/torchao.py
+++ b/src/transformers/integrations/torchao.py
@@ -83,7 +83,7 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        all_target_keys: list[str] | None = None,
+        full_layer_name: str | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -92,7 +92,6 @@ def convert(
         _, value = tuple(input_dict.items())[0]
         value = value[0] if isinstance(value, list) else value
 
-        full_layer_name = all_target_keys[0]
         module, tensor_name = get_module_from_name(model, full_layer_name)
 
         module._parameters[tensor_name] = torch.nn.Parameter(value, requires_grad=value.requires_grad)
@@ -212,11 +211,10 @@ def convert(
         self,
         input_dict: dict[str, torch.Tensor],
         model: Optional[torch.nn.Module] = None,
-        all_target_keys: list[str] | None = None,
+        full_layer_name: str | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
-        full_layer_name = all_target_keys[0]
         if isinstance(self.hf_quantizer.quantization_config.quant_type, str):
             is_int_4 = "int4" in self.hf_quantizer.quantization_config.quant_type
         else:

From 9e1bb419c95cd9306afa8b842c5f90af07b01a73 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 14:09:33 +0100
Subject: [PATCH 31/36] small thingy

---
 src/transformers/core_model_loading.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index ac372c5d07ab..04b17c119292 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -444,9 +444,10 @@ def convert(
 
         # Tensors are returned from ops with the target patterns, we need to expand them to full name.
         # This means we need to grab the prefix and suffix to add to every target key
+        full_name = layer_name
         if ".*." in layer_name:
-            layer_name = layer_name.replace(".*.", ".0.")
-        prefix, _, suffix = next(layer_name.partition(k) for k in collected_tensors.keys() if k in layer_name)
+            full_name = layer_name.replace(".*.", ".0.")
+        prefix, _, suffix = next(full_name.partition(k) for k in collected_tensors.keys() if k in full_name)
         # Rename the tensors
         collected_tensors = {prefix + k + suffix: v for k, v in collected_tensors.items()}
 

From 1a240e60f254161c3450076eb9f2e2972ff19095 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 14:11:26 +0100
Subject: [PATCH 32/36] add no_grad decorator

---
 src/transformers/core_model_loading.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 04b17c119292..bda7516e6f64 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -124,6 +124,7 @@ class Chunk(ConversionOps):
     def __init__(self, dim: int = 0):
         self.dim = dim
 
+    @torch.no_grad
     def convert(
         self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
     ) -> dict[str, torch.Tensor]:

From 8cda0fd98e88cec03b5f82a32d23b32a9e0379ea Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 14:43:42 +0100
Subject: [PATCH 33/36] utils to rename keys

---
 src/transformers/core_model_loading.py | 86 +++++++++++++++++---------
 1 file changed, 58 insertions(+), 28 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index bda7516e6f64..65a5166515cd 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -105,7 +105,7 @@ class ConversionOps:
     """Base class for weight conversion operations."""
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(dim={self.dim})"
+        return f"{self.__class__.__name__}(dim={self.dim}, {self.reverse_op}"
 
     @abstractmethod
     def convert(
@@ -637,6 +637,44 @@ def repl(m, repl_map: dict[str, str]) -> str:
     return replacement
 
 
+def rename_source_key(
+    source_key: str,
+    rename_alternation: re.Pattern,
+    rename_by_group: dict,
+    weight_pattern_alternation: re.Pattern,
+    weight_pattern_by_group: dict,
+    prefix: str | None = None,
+    meta_state_dict: dict | None = None,
+) -> tuple[str, re.Match | None]:
+    """
+    Rename a source key given all the renaming and weight conversion patterns we have. Also takes care of adding/removing
+    the base model prefix during loading if necesary.
+    """
+    # 1. apply all renamings
+    renamed_key = rename_alternation.sub(lambda m: repl(m, rename_by_group), source_key).replace("\\", "")
+
+    # 2. apply renaming through weight conversions on the key if we have any WeightConverter
+    matched_converter_pattern = (
+        weight_pattern_alternation.search(renamed_key) if weight_pattern_alternation is not None else None
+    )
+    if matched_converter_pattern is not None:
+        renamed_key = weight_pattern_alternation.sub(lambda m: repl(m, weight_pattern_by_group), renamed_key).replace(
+            "\\", ""
+        )
+
+    # 3. check if we need to add or remove prefix if necesary (only during loading, not saving)
+    if prefix is not None and meta_state_dict is not None:
+        if (
+            renamed_key.startswith(prefix)
+            and meta_state_dict.get(re.sub(f"^{prefix}.", "", renamed_key, count=1)) is not None
+        ):
+            renamed_key = re.sub(f"^{prefix}.", "", renamed_key, count=1)
+        elif meta_state_dict.get(f"{prefix}.{renamed_key}") is not None:
+            renamed_key = f"{prefix}.{renamed_key}"
+
+    return renamed_key, matched_converter_pattern
+
+
 def convert_and_load_state_dict_in_model(
     model: PreTrainedModel,
     state_dict: dict[str, Any],
@@ -762,6 +800,7 @@ def convert_and_load_state_dict_in_model(
     # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
     # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
     rename_alt, _, rename_by_group = build_glob_alternation(renamings)
+    weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = None, None, None
     if converters != []:
         weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = build_glob_alternation(converters)
     if tp_plan != {}:
@@ -773,24 +812,18 @@ 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. apply all renamings
-        renamed_key = rename_alt.sub(lambda m: repl(m, rename_by_group), original_key).replace("\\", "")
-
-        # 2. apply 1 weight conversion on the key
-        matched_pattern = weight_pattern_alt.search(renamed_key) if converters != [] else None
-        if matched_pattern is not None:  # we have a converter to apply
-            renamed_key = weight_pattern_alt.sub(lambda m: repl(m, tgt_group_to_glob), renamed_key).replace("\\", "")
-
-        # 3. check if we need to add or remove prefix
-        if (
-            renamed_key.startswith(prefix)
-            and meta_model_state_dict.get(re.sub(f"^{prefix}.", "", renamed_key, count=1)) is not None
-        ):
-            renamed_key = re.sub(f"^{prefix}.", "", renamed_key, count=1)
-        elif meta_model_state_dict.get(f"{prefix}.{renamed_key}") is not None:
-            renamed_key = f"{prefix}.{renamed_key}"
+        # 1. Rename the key according to all renaming pattern and optional weight converter patterns
+        renamed_key, matched_pattern = rename_source_key(
+            original_key,
+            rename_alt,
+            rename_by_group,
+            weight_pattern_alt,
+            tgt_group_to_glob,
+            prefix,
+            meta_model_state_dict,
+        )
 
-        # 4. finally, collect the tensor into the proper converter
+        # 2. finally, collect the tensor into the proper converter
         if renamed_key in missing_keys:
             empty_param = meta_model_state_dict.get(renamed_key)
             if matched_pattern:
@@ -802,7 +835,7 @@ def convert_and_load_state_dict_in_model(
                 mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
                 source_pattern = original_key
 
-            # 5. Handle dtype casting
+            # 3. Handle dtype casting
             if (
                 hf_quantizer
                 and not hf_quantizer.pre_quantized
@@ -822,7 +855,7 @@ def convert_and_load_state_dict_in_model(
             elif empty_param is not None and empty_param.dtype != _dtype:
                 _dtype = empty_param.dtype  # usually correct when initializing
 
-            # 6. Handle TP sharding or device_map placement -> scheduled materialization
+            # 4. Handle TP sharding or device_map placement -> scheduled materialization
             future = None
             if device_mesh:
                 if matched_tp_pattern := tp_plan_alt.search(renamed_key):
@@ -936,20 +969,17 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
         # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
         # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
         rename_alt, _, rename_by_group = build_glob_alternation(renamings)
+        weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = None, None, None
         if converters != []:
             weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = build_glob_alternation(converters)
 
         state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
         for original_key, tensor in state_dict:
-            # apply all renamings
-            renamed_key = rename_alt.sub(lambda m: repl(m, rename_by_group), original_key).replace("\\", "")
-
-            # apply weight conversion on the key
-            matched_pattern = weight_pattern_alt.search(renamed_key) if converters != [] else None
+            # Rename the key according to all renaming pattern and optional weight converter patterns
+            renamed_key, matched_pattern = rename_source_key(
+                original_key, rename_alt, rename_by_group, weight_pattern_alt, tgt_group_to_glob
+            )
             if matched_pattern is not None:
-                renamed_key = weight_pattern_alt.sub(lambda m: repl(m, tgt_group_to_glob), renamed_key).replace(
-                    "\\", ""
-                )
                 new_converter = deepcopy(pattern_to_converter[src_group_to_glob[matched_pattern.lastgroup]])
                 # each target key gets its own converter instance
                 mapping = conversion_mapping.setdefault(renamed_key, new_converter)

From ca6445be2cee03b1541229606b9892a0496b4beb Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 15:40:05 +0100
Subject: [PATCH 34/36] oupssii again

---
 src/transformers/conversion_mapping.py |  5 +++--
 src/transformers/core_model_loading.py | 12 ++++++++----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index 033f22d8867b..83d3790fc2fb 100644
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -203,8 +203,9 @@ def get_model_conversion_mapping(
         model_specific_conversions = get_checkpoint_conversion_mapping(model_type)
         if model_specific_conversions is not None:
             weight_conversions.extend(model_specific_conversions)
-        elif add_legacy:
-            weight_conversions.extend(get_checkpoint_conversion_mapping("legacy"))
+
+    if add_legacy:
+        weight_conversions.extend(get_checkpoint_conversion_mapping("legacy"))
 
     # Add the ones from the quantizer as well if provided
     if hf_quantizer is not None:
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 65a5166515cd..0aabbbfe45ad 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -105,7 +105,7 @@ class ConversionOps:
     """Base class for weight conversion operations."""
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(dim={self.dim}, {self.reverse_op}"
+        return f"{self.__class__.__name__}(dim={self.dim})"
 
     @abstractmethod
     def convert(
@@ -186,7 +186,11 @@ def __init__(self, dim: int = 0):
 
     @torch.no_grad
     def convert(
-        self, input_dict: dict[str, list[torch.Tensor]], source_patterns: list[str], target_patterns: list[str]
+        self,
+        input_dict: dict[str, list[torch.Tensor]],
+        source_patterns: list[str],
+        target_patterns: list[str],
+        **kwargs,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
         for source_pattern, tensors in input_dict.items():
@@ -641,8 +645,8 @@ def rename_source_key(
     source_key: str,
     rename_alternation: re.Pattern,
     rename_by_group: dict,
-    weight_pattern_alternation: re.Pattern,
-    weight_pattern_by_group: dict,
+    weight_pattern_alternation: re.Pattern | None,
+    weight_pattern_by_group: dict | None,
     prefix: str | None = None,
     meta_state_dict: dict | None = None,
 ) -> tuple[str, re.Match | None]:

From 636290166b4b04889be61c72611b81f438221465 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 17:32:11 +0100
Subject: [PATCH 35/36] add test

---
 tests/test_modeling_common.py | 119 ++++++++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index 4359368a4fdb..1f343d1992c0 100755
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -31,6 +31,7 @@
 from packaging import version
 from parameterized import parameterized
 from pytest import mark
+from safetensors.torch import load_file
 
 from transformers import (
     AutoModel,
@@ -42,6 +43,8 @@
     logging,
     set_seed,
 )
+from transformers.conversion_mapping import get_model_conversion_mapping
+from transformers.core_model_loading import WeightRenaming
 from transformers.integrations import HfDeepSpeedConfig
 from transformers.integrations.deepspeed import (
     is_deepspeed_available,
@@ -4055,10 +4058,126 @@ def test_tp_plan_matches_params(self):
                 len(unused_entries) == 0, f"The following entries of the TP-plan are not valid: {unused_entries}"
             )
 
+    @unittest.skip("Some models have wrong mappings....")
+    def test_reverse_loading_mapping(self):
+        """Make sure we can load and save correctly the models having any weight renaming mapping or weight conversion
+        mapping.
+        Note that this test would be better if we could start from the serialized keys, and check that the model
+        keys correspond to the weight converions. However, when instantiating a model, it already has the "target"
+        keys (or modified keys after mapping) of the conversion mapping, so we have to do it the other way, i.e.
+        reverse the conversion and then check that those converted keys match correctly the conversions.
+
+        However, all the checks performed here should ensure everything is going as it should.
+        """
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+
+        #  Some MoE models alternate between a classic MLP and a MoE layer, in which case we want to have at
+        # lest one MoE layer here to check the mapping
+        config_to_set = config.get_text_config()
+        config_to_set.first_k_dense_replace = 1  # means that the first layer (idx 0) will be MLP, then MoE
+        config_to_set.moe_layer_start_index = 1  # same as above but for Ernie 4.5...
+        config_to_set.mlp_only_layers = [0]  # same but for qwens
+        config_to_set.num_dense_layers = 1  # lfm2_moe
+
+        for model_class in self.all_model_classes:
+            # Each individual model is a subtest
+            with self.subTest(model_class.__name__):
+                model = model_class(copy.deepcopy(config))
+                # Skip if no conversions
+                conversions = get_model_conversion_mapping(model, add_legacy=False)
+                if len(conversions) == 0:
+                    self.skipTest("No conversion found for this model")
+
+                # Find the model keys, so the targets according to the conversions
+                model_keys = list(model.state_dict().keys())
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    # Serialize with reverse mapping
+                    model.save_pretrained(tmpdirname)
+                    state_dict = load_file(os.path.join(tmpdirname, "model.safetensors"))
+                    # Get all the serialized keys that we just saved according to the reverse mapping
+                    serialized_keys = list(state_dict.keys())
+
+                # They should be different, otherwise we did not perform any mapping
+                self.assertNotEqual(sorted(serialized_keys), sorted(model_keys), "No key mapping was performed!")
+
+                # Check that for each conversion entry, we at least map to one key
+                for conversion in conversions:
+                    for source_pattern in conversion.source_patterns:
+                        # Sometimes the mappings specify keys that are tied, so absent from the saved state dict
+                        if isinstance(conversion, WeightRenaming):
+                            if any(
+                                re.search(conversion.target_patterns[0], 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"`{source_pattern}` in `{conversion}` did not match any of the source keys. "
+                            "This indicates whether that the pattern is not properly written, ot that it could not be reversed correctly",
+                        )
+
+                # If everything is still good at this point, let's test that we perform the same operations both when
+                # reverting ops from `from_pretrained` and from `__init__`
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    # The model was instantiated from __init__ before being saved
+                    model.save_pretrained(tmpdirname)
+                    state_dict_saved_from_init = load_file(os.path.join(tmpdirname, "model.safetensors"))
+
+                    # Now reload it
+                    model_reloaded = model_class.from_pretrained(tmpdirname)
+
+                    # Make sure both loaded state_dict are identical
+                    self.assertTrue(compare_state_dicts(model_reloaded.state_dict(), model.state_dict()))
+
+                    # The model was instantiated from `from_pretrained` before being saved
+                    model_reloaded.save_pretrained(tmpdirname)
+                    state_dict_saved_from_pretrained = load_file(os.path.join(tmpdirname, "model.safetensors"))
+
+                    # Make sure both saved state_dict are identical
+                    self.assertTrue(compare_state_dicts(state_dict_saved_from_init, state_dict_saved_from_pretrained))
+
+    @unittest.skip("Some models have wrong mappings....")
+    def test_can_load_from_already_mapped_keys(self):
+        """Test that we can correctly reload a model if we chose `save_original_format=False` in `save_pretrained`,
+        i.e. we do not reapply weight conversions when reloading if it was saved correctly already.
+        """
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            # Each individual model is a subtest
+            with self.subTest(model_class.__name__):
+                model = model_class(copy.deepcopy(config))
+
+                # Skip if no conversions
+                conversions = get_model_conversion_mapping(model, add_legacy=False)
+                if len(conversions) == 0:
+                    self.skipTest("No conversion found for this model")
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    # Serialize without reverting the mapping
+                    model.save_pretrained(tmpdirname, save_original_format=False)
+                    model_reloaded = model_class.from_pretrained(tmpdirname)
+                    # Make sure both saved state_dict are identical
+                    self.assertTrue(compare_state_dicts(model.state_dict(), model_reloaded.state_dict()))
+
 
 global_rng = random.Random()
 
 
+def compare_state_dicts(state_dict1, state_dict2) -> bool:
+    """Make sure 2 state dicts are the exact same"""
+    # Make sure the keys are the exact same
+    if sorted(state_dict1.keys()) != sorted(state_dict2.keys()):
+        raise ValueError("The keys of both state dict are not the same")
+
+    for k, v1 in state_dict1.items():
+        v2 = state_dict2[k]
+        torch.testing.assert_close(v1, v2)
+
+    return True
+
+
 def ids_tensor(shape, vocab_size, rng=None, name=None):
     #  Creates a random int32 tensor of the shape within the vocab size
     if rng is None:

From 59a97ad77a1b3dc2d1feac39c77b38828e2cb080 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Thu, 27 Nov 2025 19:15:36 +0100
Subject: [PATCH 36/36] simplify nicely

---
 src/transformers/core_model_loading.py | 123 ++++++++++---------------
 tests/test_modeling_common.py          |   5 +-
 tests/utils/test_core_model_loading.py |  66 ++++++++++++-
 3 files changed, 119 insertions(+), 75 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 0aabbbfe45ad..f4ec3a7f38ca 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -161,7 +161,16 @@ def convert(
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         target_pattern = self.get_target_pattern(target_patterns)
-        return {target_pattern: torch.cat(tuple(input_dict.values()), dim=self.dim)}
+        all_tensors = []
+        # Very important to keep the relative order of the source patterms here, so we iterate over them not the
+        # input directly as it's unordered!
+        for source_pattern in source_patterns:
+            tensors = input_dict[source_pattern]
+            if isinstance(tensors, list):
+                all_tensors.extend(tensors)
+            else:
+                all_tensors.append(tensors)
+        return {target_pattern: torch.cat(all_tensors, dim=self.dim)}
 
     def get_target_pattern(self, target_patterns: list[str]) -> str:
         # Here we always return the target pattern
@@ -348,22 +357,6 @@ def reverse_transform(self) -> WeightTransform:
             source_patterns=self.target_patterns, target_patterns=self.source_patterns, **kwargs
         )
 
-        # In this case, we already called `add_tensor` at least once, and we will use the saved keys to revert
-        # (i.e. we will be able to retrieve directly the correct params thanks to the `collected_tensors` which
-        # will hold the full param names entries)
-        if len(self.layer_targets) > 0:
-            # Find the full names of the params we will need to use later for the reverse transform
-            reverse_layer_targets = defaultdict(set)
-            reverse_collected_tensors = defaultdict(list)
-            for target_key, all_sources in self.layer_targets.items():
-                matched_target_pattern = next(pat for pat in self.target_patterns if re.search(pat, target_key))
-                reverse_collected_tensors[matched_target_pattern].append(target_key)
-                for source in all_sources:
-                    reverse_layer_targets[source].add(target_key)
-            reverse_collected_tensors = {k: sorted(set(v)) for k, v in reverse_collected_tensors.items()}
-            reverse_transform.layer_targets = reverse_layer_targets
-            reverse_transform.collected_tensors = reverse_collected_tensors
-
         return reverse_transform
 
 
@@ -935,81 +928,65 @@ def convert_and_load_state_dict_in_model(
             except SkipLayer:
                 continue
 
-    # Keep computed mapping as attribute for later saving
-    model._weight_loading_mapping = param_name_to_load
+    # Keep the current weight conversion mapping for later saving (in case it was coming directly from the user)
+    model._weight_conversions = weight_mapping
     thread_pool.shutdown(wait=False)
     return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, misc
 
 
 def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
-    """ "
+    """
     Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one
     if the model was created in another way and is part of the default mappings.
     """
-    conversion_mapping = getattr(model, "_weight_loading_mapping", None)
-    need_to_reverse = True
+    weight_conversions = getattr(model, "_weight_conversions", None)
     # In this case, the model was not created with `from_pretrained` -> let's check if it's in the hardcoded
     # mappings, and recreate the mapping from there if it is
-    if conversion_mapping is None:
+    if weight_conversions is None:
         from .conversion_mapping import get_model_conversion_mapping
 
         # Do not resave with the legacy renaming, if present
         weight_conversions = get_model_conversion_mapping(model, add_legacy=False)
+        weight_conversions = weight_conversions if len(weight_conversions) > 0 else None
 
-        # We did not find any operations to perform -> quick escape
-        if len(weight_conversions) == 0:
-            return state_dict
-
-        conversion_mapping = {}
-        need_to_reverse = False
-
-        # Already 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}
-
-        # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
-        # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
-        rename_alt, _, rename_by_group = build_glob_alternation(renamings)
-        weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = None, None, None
-        if converters != []:
-            weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = build_glob_alternation(converters)
-
-        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, matched_pattern = rename_source_key(
-                original_key, rename_alt, rename_by_group, weight_pattern_alt, tgt_group_to_glob
-            )
-            if matched_pattern is not None:
-                new_converter = deepcopy(pattern_to_converter[src_group_to_glob[matched_pattern.lastgroup]])
-                # each target key gets its own converter instance
-                mapping = conversion_mapping.setdefault(renamed_key, new_converter)
-                source_pattern = src_group_to_glob[matched_pattern.lastgroup]
-            else:
-                mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
-                source_pattern = original_key
+    # We did not find any operations to perform -> quick escape
+    if weight_conversions is None:
+        return state_dict
 
-            mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
+    # 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 = {}
 
-    new_state_dict = {}
-    for first_param_name, converter in conversion_mapping.items():
-        # In this case, the mapping was obtained from the loaded model, we need to reverse the ops
-        if need_to_reverse:
-            reversed_converter = converter.reverse_transform()
-            # The `collected_tensors` only contain the full param keys after `reverse_transform` is called, we
-            # need to populate it with the actual tensors
-            reversed_converter.collected_tensors = {
-                k: [model.get_parameter_or_buffer(param) for param in params]
-                for k, params in reversed_converter.collected_tensors.items()
-            }
-            first_param_name = next(iter(reversed_converter.layer_targets.keys()))
-        # In this case, we just created the mapping above, the ops are already reversed
+    # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
+    # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
+    rename_alt, _, rename_by_group = build_glob_alternation(renamings)
+    weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = None, None, None
+    if converters != []:
+        weight_pattern_alt, src_group_to_glob, tgt_group_to_glob = build_glob_alternation(converters)
+
+    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, matched_pattern = rename_source_key(
+            original_key, rename_alt, rename_by_group, weight_pattern_alt, tgt_group_to_glob
+        )
+        if matched_pattern is not None:
+            new_converter = deepcopy(pattern_to_converter[src_group_to_glob[matched_pattern.lastgroup]])
+            # each target key gets its own converter instance
+            mapping = conversion_mapping.setdefault(renamed_key, new_converter)
+            source_pattern = src_group_to_glob[matched_pattern.lastgroup]
         else:
-            reversed_converter = converter
+            mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
+            source_pattern = original_key
+
+        mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)
 
+    new_state_dict = {}
+    for first_param_name, reversed_converter in conversion_mapping.items():
         # Apply the reverse converter
         realized_value, misc = reversed_converter.convert(first_param_name, model=model, config=model.config)
         for target_name, param in realized_value.items():
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index 1f343d1992c0..c1f7c1b83ed0 100755
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -4173,7 +4173,10 @@ def compare_state_dicts(state_dict1, state_dict2) -> bool:
 
     for k, v1 in state_dict1.items():
         v2 = state_dict2[k]
-        torch.testing.assert_close(v1, v2)
+        try:
+            torch.testing.assert_close(v1, v2)
+        except Exception as e:
+            raise AssertionError(f"For key {k}: {e}")
 
     return True
 
diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py
index 3943238c0412..88bdb27256ba 100644
--- a/tests/utils/test_core_model_loading.py
+++ b/tests/utils/test_core_model_loading.py
@@ -28,9 +28,12 @@
     build_glob_alternation,
     convert_and_load_state_dict_in_model,
     repl,
+    revert_weight_conversion,
 )
 from transformers.utils.import_utils import is_triton_available
 
+from ..test_modeling_common import compare_state_dicts
+
 
 class TestWeightGlobMatching(unittest.TestCase):
     def setUp(self):
@@ -315,6 +318,67 @@ def stack_down(layer_prefix: str) -> torch.Tensor:
 
         torch.testing.assert_close(model_state["mlp.down_proj.weight"], raw_tensors["mlp.w2.weight"])
 
+    def test_moe_and_qkv_conversion_reversed(self):
+        model = DummyRoot()
+        model.config = PretrainedConfig()
+
+        raw_tensors = {
+            "model.layers.0.experts.0.w1.weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]]),
+            "model.layers.0.experts.1.w1.weight": torch.tensor([[10.0, 11.0], [12.0, 13.0]]),
+            "model.layers.0.experts.0.w3.weight": torch.tensor([[4.0, 5.0], [6.0, 7.0]]),
+            "model.layers.0.experts.1.w3.weight": torch.tensor([[14.0, 15.0], [16.0, 17.0]]),
+            "model.layers.0.experts.0.w2.weight": torch.tensor([[20.0, 21.0], [22.0, 23.0]]),
+            "model.layers.0.experts.1.w2.weight": torch.tensor([[24.0, 25.0], [26.0, 27.0]]),
+            "model.layers.1.experts.0.w1.weight": torch.tensor([[30.0, 31.0], [32.0, 33.0]]),
+            "model.layers.1.experts.1.w1.weight": torch.tensor([[34.0, 35.0], [36.0, 37.0]]),
+            "model.layers.1.experts.0.w3.weight": torch.tensor([[38.0, 39.0], [40.0, 41.0]]),
+            "model.layers.1.experts.1.w3.weight": torch.tensor([[42.0, 43.0], [44.0, 45.0]]),
+            "model.layers.1.experts.0.w2.weight": torch.tensor([[46.0, 47.0], [48.0, 49.0]]),
+            "model.layers.1.experts.1.w2.weight": torch.tensor([[50.0, 51.0], [52.0, 53.0]]),
+            "model.layers.0.self_attn.qkv_proj.weight": torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
+            "model.layers.1.self_attn.qkv_proj.weight": torch.tensor([[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]),
+            "mlp.w2.weight": torch.tensor([[60.0, 61.0], [62.0, 63.0]]),
+        }
+        state_dict = {k: v.clone() for k, v in raw_tensors.items()}
+
+        weight_mapping = [
+            WeightConverter(
+                ["experts.*.w1.weight", "experts.*.w3.weight"],
+                "experts.gate_up_proj.weight",
+                operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
+            ),
+            WeightConverter(
+                "experts.*.w2.weight",
+                "experts.down_proj.weight",
+                operations=[MergeModulelist(dim=0)],
+            ),
+            WeightConverter(
+                "self_attn.qkv_proj.weight",
+                [
+                    "self_attn.q_proj.weight",
+                    "self_attn.k_proj.weight",
+                    "self_attn.v_proj.weight",
+                ],
+                operations=[Chunk(dim=0)],
+            ),
+            WeightRenaming("mlp.w2.weight", "mlp.down_proj.weight"),
+        ]
+
+        # Use the mapping to load
+        missing, unexpected, mismatch, _, misc = convert_and_load_state_dict_in_model(
+            model, state_dict, weight_mapping, tp_plan=None, hf_quantizer=None
+        )
+        self.assertTrue(len(missing) == 0)
+        self.assertTrue(len(unexpected) == 0)
+        self.assertTrue(len(mismatch) == 0)
+        self.assertTrue(len(misc) == 0)
+
+        # Try to revert the mapping
+        reversed_state_dict = revert_weight_conversion(model, model.state_dict())
+
+        # Make sure both saved state_dict are identical
+        self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict))
+
     def test_qkv_chunk_rope_permute_with_fp8_quantization(self):
         if is_triton_available():
             from transformers.integrations.finegrained_fp8 import Fp8Dequantize
@@ -396,7 +460,7 @@ def __init__(self):
                     "model.layers.*.self_attn.k_proj.weight",
                     "model.layers.*.self_attn.v_proj.weight",
                 ],
-                operations=[Chunk(dim=0, chunks=3), PermuteForRope()],
+                operations=[Chunk(dim=0), PermuteForRope()],
             )
         ]