Loading: convert transposed and RoPE-permuted weights before taking the DTensor shard

src/transformers/core_model_loading.py

@@ -97,6 +97,11 @@ def convert(
     def reverse_op(self) -> ConversionOps:
         raise NotImplementedError
 
+    def shards_after_conversion(self, placements, ndim: int) -> bool:
+        """Whether this op forces a DTensor parameter with these `placements` to be converted in full, before this
+        rank's shard is taken. An op that maps every target element to one source element keeps shard-on-read."""
+        return False
+
 
 class _IdentityOp(ConversionOps):
     """Pass-through reverse op for dequantize operations.
@@ -437,6 +442,11 @@ def __init__(
         self.inverse = inverse
         self.permute_layer_names = permute_layer_names
 
+    def shards_after_conversion(self, placements, ndim: int) -> bool:
+        # `_apply` interleaves rows and reads the head size off `tensor.shape[0]`, so on a shard it permutes with the
+        # local row count and silently produces the wrong weights.
+        return 0 in {placement.dim % ndim for placement in placements if placement.is_shard()}
+
     def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
         dim0 = tensor.shape[0]
         config = self.config
@@ -988,6 +998,11 @@ def was_used(self) -> bool:
         """
         return self._was_used
 
+    def shards_after_conversion(self, placements, ndim: int) -> bool:
+        """Whether a DTensor parameter with these `placements` has to be converted in full before this rank's shard is
+        taken. A rename moves no elements, so its shard is a slice of the source and shard-on-read applies."""
+        return False
+
 
 class WeightRenaming(WeightTransform):
     # Special case of WeightTransform that only renames keys without any conversion.
@@ -1172,6 +1187,9 @@ def __init__(
         if not self.operations:
             raise ValueError("WeightConverter requires at least one operation.")
 
+    def shards_after_conversion(self, placements, ndim: int) -> bool:
+        return any(op.shards_after_conversion(placements, ndim) for op in self.operations)
+
     def convert(
         self,
         layer_name: str,
@@ -1716,7 +1734,11 @@ def convert_and_load_state_dict_in_model(
             # 4. Handle DTensor sharding or device_map placement
             param_device = get_device(device_map, renamed_key, valid_torch_device=True)
             sharding_op = None
-            if is_dtensor(empty_param):
+            # A RoPE permutation moves elements across the sharded dim, so this rank's slice of the source is not the
+            # slice of the converted tensor: those are converted in full and sharded afterwards, and never get an op.
+            if is_dtensor(empty_param) and not mapping.shards_after_conversion(
+                empty_param.placements, empty_param.ndim
+            ):
                 sharding_op = DtensorShardOperation(empty_param)
 
             # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations
@@ -1760,6 +1782,11 @@ def convert_and_load_state_dict_in_model(
                 )
                 for target_name, param in realized_value.items():
                     param = param[0] if isinstance(param, list) else param
+                    empty_param = meta_model_state_dict.get(target_name)
+                    if is_dtensor(empty_param) and mapping.shards_after_conversion(
+                        empty_param.placements, empty_param.ndim
+                    ):
+                        param = DtensorShardOperation(empty_param).shard_tensor(param)
                     param_device = get_device(device_map, target_name)
                     # Exception for params that are so huge that they need conversions on cpu no matter what - we put them back on
                     # the correct device now (if the device_map managed to make them fit on any device, as usually they will stay on

tests/utils/test_core_model_loading.py

@@ -17,7 +17,7 @@
 
 import torch
 import torch.nn as nn
-from torch.distributed.tensor.placement_types import Shard
+from torch.distributed.tensor.placement_types import Replicate, Shard
 
 from transformers import PreTrainedConfig, PreTrainedModel
 from transformers.conversion_mapping import (
@@ -35,6 +35,7 @@
     MergeModulelist,
     PermuteForRope,
     PrefixChange,
+    Transpose,
     VisionFuseAndPermuteForRope,
     VisionUnfuseAndPermuteForRope,
     WeightConverter,
@@ -334,6 +335,26 @@ def test_dtensor_shard_aware_mixtral_conversion_uses_only_local_experts(self):
             torch.tensor([[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]),
         )
 
+    def test_rope_permute_is_sharded_after_conversion(self):
+        """`PermuteForRope` interleaves rows and takes the head size from `tensor.shape[0]`, so on a shard it
+        permutes with the local row count and silently produces the wrong weights. Those parameters are converted in
+        full and sharded afterwards (`shards_after_conversion`). A transpose only moves axes, so the target shard maps
+        to a source slice and shard-on-read still applies to it (#48373); the transpose cases below pin that boundary."""
+        transpose = WeightConverter("router.weight", "gate.weight", operations=[Transpose(0, 1)])
+        permute = WeightConverter("q.weight", "q.weight", operations=[PermuteForRope()])
+        experts = WeightConverter("w13", "gate_up", operations=[Transpose(1, 2, check_dims=True)])
+        merge = WeightConverter(["w1", "w3"], "gate_up", operations=[MergeModulelist(dim=0), Concatenate(dim=1)])
+        self.assertFalse(transpose.shards_after_conversion([Shard(0)], ndim=2))
+        self.assertFalse(transpose.shards_after_conversion([Replicate(), Shard(1)], ndim=2))
+        self.assertFalse(transpose.shards_after_conversion([Replicate()], ndim=2))
+        self.assertTrue(permute.shards_after_conversion([Shard(0)], ndim=2))
+        self.assertFalse(permute.shards_after_conversion([Shard(1)], ndim=2))
+        # Qwen3-VL-MoE transposes the two inner dims of the experts, which expert parallelism shards on dim 0
+        self.assertFalse(experts.shards_after_conversion([Shard(0)], ndim=3))
+        self.assertFalse(experts.shards_after_conversion([Shard(0), Shard(2)], ndim=3))
+        self.assertFalse(merge.shards_after_conversion([Shard(0)], ndim=3))
+        self.assertFalse(WeightRenaming("q.weight", "q.weight").shards_after_conversion([Shard(0)], ndim=2))
+
     def test_moe_and_qkv_conversion(self):
         model = DummyRoot(PreTrainedConfig())