Repository: huggingface/transformers
Pull request: Add a dim check mechanism in Transpose and fix qwen3_vl_moe weight mapping

# What does this PR do?

As per the title. Supersedes https://github.com/huggingface/transformers/pull/43913

Changed files:
filename: src/transformers/conversion_mapping.py
patch: @@ -141,17 +141,14 @@ def _build_checkpoint_conversion_mapping():
         ],
         "qwen3_vl_moe": [
             WeightConverter(
-                source_patterns=[
-                    "mlp.experts.*.gate_proj.weight",
-                    "mlp.experts.*.up_proj.weight",
-                ],
+                source_patterns="mlp.experts.gate_up_proj",
                 target_patterns="mlp.experts.gate_up_proj",
-                operations=[MergeModulelist(dim=0), Concatenate(dim=1), Transpose(1, 2)],
+                operations=[Transpose(1, 2, check_dims=True), Force16BytesAlignment()],
             ),
             WeightConverter(
-                source_patterns="mlp.experts.*.down_proj.weight",
+                source_patterns="mlp.experts.down_proj",
                 target_patterns="mlp.experts.down_proj",
-                operations=[MergeModulelist(dim=0), Transpose(1, 2)],
+                operations=[Transpose(1, 2, check_dims=True), Force16BytesAlignment()],
             ),
         ],
         "phimoe": [

filename: src/transformers/core_model_loading.py
patch: @@ -249,9 +249,10 @@ class Transpose(ConversionOps):
     Transposes the given tensor along dim0 and dim1.
     """
 
-    def __init__(self, dim0: int = 0, dim1: int = 1):
+    def __init__(self, dim0: int = 0, dim1: int = 1, check_dims: bool = False):
         self.dim0 = dim0
         self.dim1 = dim1
+        self.check_dims = check_dims
 
     @torch.no_grad
     def convert(
@@ -260,7 +261,18 @@ def convert(
         target_pattern = self.get_target_pattern(input_dict, source_patterns, target_patterns)
         tensors = next(iter(input_dict.values()))
         tensor = tensors[0] if isinstance(tensors, list) else tensors
-        return {target_pattern: torch.transpose(tensor, dim0=self.dim0, dim1=self.dim1).contiguous()}
+        # In this case, always transpose
+        if not self.check_dims:
+            return {target_pattern: torch.transpose(tensor, dim0=self.dim0, dim1=self.dim1).contiguous()}
+        # In this case, check the shapes before transposing
+        else:
+            # NOTE: this rely on the first param name, so cannot be used for many-to-one operation
+            expected_shape = kwargs["model"].get_parameter(kwargs["full_layer_name"]).shape
+            # The shapes are the same: do NOT transpose
+            if tensor.shape == expected_shape:
+                return {target_pattern: tensor}
+            else:
+                return {target_pattern: torch.transpose(tensor, dim0=self.dim0, dim1=self.dim1).contiguous()}
 
     def get_target_pattern(
         self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str]
@@ -279,7 +291,7 @@ def get_target_pattern(
 
     @property
     def reverse_op(self) -> ConversionOps:
-        return Transpose(dim0=self.dim1, dim1=self.dim0)
+        return Transpose(dim0=self.dim1, dim1=self.dim0, check_dims=self.check_dims)
 
 
 class PermuteForRope(ConversionOps):

filename: tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py
patch: @@ -304,6 +304,10 @@ def test_video_forward(self):
             )
             self.assertIsNotNone(outputs)
 
+    # Need to be False as we only use a Transpose without modifying the keys
+    def test_reverse_loading_mapping(self, check_keys_were_modified=False):
+        super().test_reverse_loading_mapping(check_keys_were_modified)
+
 
 @require_torch
 class Qwen3VLMoeIntegrationTest(unittest.TestCase):