From bb3a775dc83cc594cc509171cf3a0431401b9227 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 03:33:55 +0000 Subject: [PATCH 1/6] Load: convert transposed and RoPE-permuted weights in full before taking the DTensor shard --- src/transformers/core_model_loading.py | 14 +++++++++++++ tests/utils/test_core_model_loading.py | 27 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 580f46849c3c..c95658930f84 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1269,6 +1269,14 @@ def _job(): return _job +def shards_after_conversion(mapping: WeightRenaming | WeightConverter) -> bool: + """Whether a DTensor parameter has to be converted in full before this rank's shard is taken: a transpose or a RoPE + permutation moves elements across the sharded dim, so sharding the source first would give the wrong slice.""" + return isinstance(mapping, WeightConverter) and any( + isinstance(op, (Transpose, PermuteForRope)) for op in mapping.operations + ) + + def dot_natural_key(s: str): """ Sort key for state-dict names: split on `"."` and sort digits numerically and strings alphabetically. It emits a @@ -1718,6 +1726,10 @@ def convert_and_load_state_dict_in_model( sharding_op = None if is_dtensor(empty_param): sharding_op = DtensorShardOperation(empty_param) + # A transpose or 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. + if sharding_op is not None and shards_after_conversion(mapping): + sharding_op = None # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations # directly on the device, as it will completely blow up the memory during the ops memory spike. So defer to "cpu", then @@ -1760,6 +1772,8 @@ 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 + if shards_after_conversion(mapping) and is_dtensor(meta_model_state_dict.get(target_name)): + param = DtensorShardOperation(meta_model_state_dict[target_name]).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 diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 309879ab5091..51a1984ef88d 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -35,6 +35,7 @@ MergeModulelist, PermuteForRope, PrefixChange, + Transpose, VisionFuseAndPermuteForRope, VisionUnfuseAndPermuteForRope, WeightConverter, @@ -43,6 +44,7 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, + shards_after_conversion, spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig @@ -334,6 +336,31 @@ 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_transpose_and_rope_permute_are_sharded_after_conversion(self): + """A transpose or a RoPE permutation moves elements across the sharded dim, so the loader converts the full + tensor and takes this rank's shard afterwards (`shards_after_conversion`). Sharding the source first gives the + wrong slice: for a (4, 2) checkpoint tensor transposed into a (2, 4) parameter sharded on dim 0 over 2 ranks, + rank 0 needs row 0 of the transposed tensor, i.e. column 0 of the source, not rows 0-1 of the source.""" + source = torch.arange(8.0).view(4, 2) + converter = WeightConverter("router.weight", "gate.weight", operations=[Transpose(0, 1)]) + self.assertTrue(shards_after_conversion(converter)) + self.assertTrue( + shards_after_conversion(WeightConverter("q.weight", "q.weight", operations=[PermuteForRope()])) + ) + self.assertFalse( + shards_after_conversion( + WeightConverter(["w1", "w3"], "gate_up", operations=[MergeModulelist(dim=0), Concatenate(dim=1)]) + ) + ) + + converter.add_tensor("gate.weight", "router.weight", "router.weight", spawn_materialize(None, source, "cpu")) + converted = converter.convert("gate.weight")["gate.weight"] + for rank in range(2): + shard_op = _make_dtensor_shard_op( + FakeMesh(shape=(2,), rank=rank), [Shard(0)], param_shape=(2, 4), local_shape=(1, 4) + ) + torch.testing.assert_close(shard_op.shard_tensor(converted), source.T[rank : rank + 1]) + def test_moe_and_qkv_conversion(self): model = DummyRoot(PreTrainedConfig()) From 9e2bca9509e50bd53b959c590548e32992e9afa5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 03:58:42 +0000 Subject: [PATCH 2/6] Let conversion ops declare the dims they move; defer the shard only when one of them is sharded --- src/transformers/core_model_loading.py | 34 ++++++++++++++++------ tests/utils/test_core_model_loading.py | 40 ++++++++++++++------------ 2 files changed, 47 insertions(+), 27 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index c95658930f84..444b79248c51 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,6 +81,11 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" + def moved_dims(self, ndim: int) -> set[int]: + """Dims across which this op moves elements. A DTensor parameter sharded on one of them is converted in full + before this rank's shard is taken, since the shard of the converted tensor is not the converted shard.""" + return set() + def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -319,6 +324,9 @@ def __init__(self, dim0: int = 0, dim1: int = 1, check_dims: bool = False): self.dim1 = dim1 self.check_dims = check_dims + def moved_dims(self, ndim: int) -> set[int]: + return {self.dim0 % ndim, self.dim1 % ndim} + @torch.no_grad def convert( self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs @@ -437,6 +445,9 @@ def __init__( self.inverse = inverse self.permute_layer_names = permute_layer_names + def moved_dims(self, ndim: int) -> set[int]: + return {0} + def _apply(self, tensor: torch.Tensor) -> torch.Tensor: dim0 = tensor.shape[0] config = self.config @@ -1269,12 +1280,14 @@ def _job(): return _job -def shards_after_conversion(mapping: WeightRenaming | WeightConverter) -> bool: - """Whether a DTensor parameter has to be converted in full before this rank's shard is taken: a transpose or a RoPE - permutation moves elements across the sharded dim, so sharding the source first would give the wrong slice.""" - return isinstance(mapping, WeightConverter) and any( - isinstance(op, (Transpose, PermuteForRope)) for op in mapping.operations - ) +def shards_after_conversion(mapping: WeightRenaming | WeightConverter, placements, ndim: int) -> bool: + """Whether a DTensor parameter with these `placements` has to be converted in full before this rank's shard is + taken: an op that moves elements across a sharded dim (a transpose, a RoPE permutation) makes the shard of the + source the wrong slice.""" + if not isinstance(mapping, WeightConverter): + return False + shard_dims = {placement.dim % ndim for placement in placements if placement.is_shard()} + return any(op.moved_dims(ndim) & shard_dims for op in mapping.operations) def dot_natural_key(s: str): @@ -1728,7 +1741,7 @@ def convert_and_load_state_dict_in_model( sharding_op = DtensorShardOperation(empty_param) # A transpose or 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. - if sharding_op is not None and shards_after_conversion(mapping): + if sharding_op is not None and shards_after_conversion(mapping, empty_param.placements, empty_param.ndim): sharding_op = None # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations @@ -1772,8 +1785,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 - if shards_after_conversion(mapping) and is_dtensor(meta_model_state_dict.get(target_name)): - param = DtensorShardOperation(meta_model_state_dict[target_name]).shard_tensor(param) + empty_param = meta_model_state_dict.get(target_name) + if is_dtensor(empty_param) and shards_after_conversion( + mapping, 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 diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 51a1984ef88d..8415c2d8f74f 100644 --- a/tests/utils/test_core_model_loading.py +++ b/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 ( @@ -337,24 +337,28 @@ def test_dtensor_shard_aware_mixtral_conversion_uses_only_local_experts(self): ) def test_transpose_and_rope_permute_are_sharded_after_conversion(self): - """A transpose or a RoPE permutation moves elements across the sharded dim, so the loader converts the full - tensor and takes this rank's shard afterwards (`shards_after_conversion`). Sharding the source first gives the - wrong slice: for a (4, 2) checkpoint tensor transposed into a (2, 4) parameter sharded on dim 0 over 2 ranks, - rank 0 needs row 0 of the transposed tensor, i.e. column 0 of the source, not rows 0-1 of the source.""" - source = torch.arange(8.0).view(4, 2) - converter = WeightConverter("router.weight", "gate.weight", operations=[Transpose(0, 1)]) - self.assertTrue(shards_after_conversion(converter)) - self.assertTrue( - shards_after_conversion(WeightConverter("q.weight", "q.weight", operations=[PermuteForRope()])) - ) - self.assertFalse( - shards_after_conversion( - WeightConverter(["w1", "w3"], "gate_up", operations=[MergeModulelist(dim=0), Concatenate(dim=1)]) - ) - ) + """A transpose or a RoPE permutation moves elements across dims, so when a shard placement sits on one of + those dims the loader converts the full tensor and takes this rank's shard afterwards + (`shards_after_conversion`). Sharding the source first gives the wrong slice: for a (4, 2) checkpoint tensor + transposed into a (2, 4) parameter sharded on dim 0 over 2 ranks, rank 0 needs row 0 of the transposed tensor, + i.e. column 0 of the source, not rows 0-1 of the source.""" + 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.assertTrue(shards_after_conversion(transpose, [Shard(0)], ndim=2)) + self.assertTrue(shards_after_conversion(transpose, [Replicate(), Shard(1)], ndim=2)) + self.assertFalse(shards_after_conversion(transpose, [Replicate()], ndim=2)) + self.assertTrue(shards_after_conversion(permute, [Shard(0)], ndim=2)) + self.assertFalse(shards_after_conversion(permute, [Shard(1)], ndim=2)) + # Qwen3-VL-MoE transposes the two inner dims of the experts, which expert parallelism shards on dim 0 + self.assertFalse(shards_after_conversion(experts, [Shard(0)], ndim=3)) + self.assertTrue(shards_after_conversion(experts, [Shard(0), Shard(2)], ndim=3)) + self.assertFalse(shards_after_conversion(merge, [Shard(0)], ndim=3)) - converter.add_tensor("gate.weight", "router.weight", "router.weight", spawn_materialize(None, source, "cpu")) - converted = converter.convert("gate.weight")["gate.weight"] + source = torch.arange(8.0).view(4, 2) + transpose.add_tensor("gate.weight", "router.weight", "router.weight", spawn_materialize(None, source, "cpu")) + converted = transpose.convert("gate.weight")["gate.weight"] for rank in range(2): shard_op = _make_dtensor_shard_op( FakeMesh(shape=(2,), rank=rank), [Shard(0)], param_shape=(2, 4), local_shape=(1, 4) From d83a1f72b35d4fef34cf19269f14e18466fdec49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 04:00:57 +0000 Subject: [PATCH 3/6] Keep the decision in one helper --- src/transformers/core_model_loading.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 444b79248c51..636027777ae4 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,11 +81,6 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" - def moved_dims(self, ndim: int) -> set[int]: - """Dims across which this op moves elements. A DTensor parameter sharded on one of them is converted in full - before this rank's shard is taken, since the shard of the converted tensor is not the converted shard.""" - return set() - def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -324,9 +319,6 @@ def __init__(self, dim0: int = 0, dim1: int = 1, check_dims: bool = False): self.dim1 = dim1 self.check_dims = check_dims - def moved_dims(self, ndim: int) -> set[int]: - return {self.dim0 % ndim, self.dim1 % ndim} - @torch.no_grad def convert( self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs @@ -445,9 +437,6 @@ def __init__( self.inverse = inverse self.permute_layer_names = permute_layer_names - def moved_dims(self, ndim: int) -> set[int]: - return {0} - def _apply(self, tensor: torch.Tensor) -> torch.Tensor: dim0 = tensor.shape[0] config = self.config @@ -1282,12 +1271,16 @@ def _job(): def shards_after_conversion(mapping: WeightRenaming | WeightConverter, placements, ndim: int) -> bool: """Whether a DTensor parameter with these `placements` has to be converted in full before this rank's shard is - taken: an op that moves elements across a sharded dim (a transpose, a RoPE permutation) makes the shard of the - source the wrong slice.""" + taken: a transpose or a RoPE permutation on a sharded dim makes the shard of the source the wrong slice.""" if not isinstance(mapping, WeightConverter): return False shard_dims = {placement.dim % ndim for placement in placements if placement.is_shard()} - return any(op.moved_dims(ndim) & shard_dims for op in mapping.operations) + for op in mapping.operations: + if isinstance(op, Transpose) and {op.dim0 % ndim, op.dim1 % ndim} & shard_dims: + return True + if isinstance(op, PermuteForRope) and 0 in shard_dims: + return True + return False def dot_natural_key(s: str): From f4562cf15f944662d067fcd9a1316bdf9b366387 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Tue, 8 Sep 2026 17:46:23 +0000 Subject: [PATCH 4/6] Narrow to the RoPE permutation, leave transposes to shard-on-read A transpose only moves axes, so a target shard maps to a slice of the source and shard-on-read still applies. #48373 does that translation, which is better than converting in full, so drop the transpose case here. PermuteForRope is not that: it interleaves rows and takes the head size from tensor.shape[0], which on a shard is the local row count, so it permutes with the wrong head size and raises nothing. That one still has to be converted before the shard is taken. --- src/transformers/core_model_loading.py | 14 +++++++------- tests/utils/test_core_model_loading.py | 26 ++++++++------------------ 2 files changed, 15 insertions(+), 25 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 636027777ae4..2331567fb9a2 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1271,16 +1271,16 @@ def _job(): def shards_after_conversion(mapping: WeightRenaming | WeightConverter, 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 transpose or a RoPE permutation on a sharded dim makes the shard of the source the wrong slice.""" + taken. + + `PermuteForRope` 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. A transpose is not handled here: it only moves axes + around, so the target shard maps to a source slice and shard-on-read still applies (#48373). + """ if not isinstance(mapping, WeightConverter): return False shard_dims = {placement.dim % ndim for placement in placements if placement.is_shard()} - for op in mapping.operations: - if isinstance(op, Transpose) and {op.dim0 % ndim, op.dim1 % ndim} & shard_dims: - return True - if isinstance(op, PermuteForRope) and 0 in shard_dims: - return True - return False + return any(isinstance(op, PermuteForRope) for op in mapping.operations) and 0 in shard_dims def dot_natural_key(s: str): diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 8415c2d8f74f..2f1433593021 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -336,35 +336,25 @@ 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_transpose_and_rope_permute_are_sharded_after_conversion(self): - """A transpose or a RoPE permutation moves elements across dims, so when a shard placement sits on one of - those dims the loader converts the full tensor and takes this rank's shard afterwards - (`shards_after_conversion`). Sharding the source first gives the wrong slice: for a (4, 2) checkpoint tensor - transposed into a (2, 4) parameter sharded on dim 0 over 2 ranks, rank 0 needs row 0 of the transposed tensor, - i.e. column 0 of the source, not rows 0-1 of the source.""" + 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.assertTrue(shards_after_conversion(transpose, [Shard(0)], ndim=2)) - self.assertTrue(shards_after_conversion(transpose, [Replicate(), Shard(1)], ndim=2)) + self.assertFalse(shards_after_conversion(transpose, [Shard(0)], ndim=2)) + self.assertFalse(shards_after_conversion(transpose, [Replicate(), Shard(1)], ndim=2)) self.assertFalse(shards_after_conversion(transpose, [Replicate()], ndim=2)) self.assertTrue(shards_after_conversion(permute, [Shard(0)], ndim=2)) self.assertFalse(shards_after_conversion(permute, [Shard(1)], ndim=2)) # Qwen3-VL-MoE transposes the two inner dims of the experts, which expert parallelism shards on dim 0 self.assertFalse(shards_after_conversion(experts, [Shard(0)], ndim=3)) - self.assertTrue(shards_after_conversion(experts, [Shard(0), Shard(2)], ndim=3)) + self.assertFalse(shards_after_conversion(experts, [Shard(0), Shard(2)], ndim=3)) self.assertFalse(shards_after_conversion(merge, [Shard(0)], ndim=3)) - source = torch.arange(8.0).view(4, 2) - transpose.add_tensor("gate.weight", "router.weight", "router.weight", spawn_materialize(None, source, "cpu")) - converted = transpose.convert("gate.weight")["gate.weight"] - for rank in range(2): - shard_op = _make_dtensor_shard_op( - FakeMesh(shape=(2,), rank=rank), [Shard(0)], param_shape=(2, 4), local_shape=(1, 4) - ) - torch.testing.assert_close(shard_op.shard_tensor(converted), source.T[rank : rank + 1]) - def test_moe_and_qkv_conversion(self): model = DummyRoot(PreTrainedConfig()) From cb61b1aab880f1b5cab690bcf255eb2d4d8635c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 10 Sep 2026 04:26:48 +0000 Subject: [PATCH 5/6] Only build the sharding op when it is used A parameter that has to be converted in full discarded the `DtensorShardOperation` right after building it. --- src/transformers/core_model_loading.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 2331567fb9a2..6d86c37bb3b9 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1730,12 +1730,12 @@ 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 shards_after_conversion( + mapping, empty_param.placements, empty_param.ndim + ): sharding_op = DtensorShardOperation(empty_param) - # A transpose or 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. - if sharding_op is not None and shards_after_conversion(mapping, empty_param.placements, empty_param.ndim): - sharding_op = None # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations # directly on the device, as it will completely blow up the memory during the ops memory spike. So defer to "cpu", then From 071db77579907fb571aa1d7dd2973b38c0639a55 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 10 Sep 2026 14:28:21 +0000 Subject: [PATCH 6/6] Let each op say whether it breaks shard-on-read `shards_after_conversion` was a single function growing an isinstance chain. It is now a method: the default on `ConversionOps` returns False, `PermuteForRope` overrides it, and a `WeightConverter` answers for the ops it holds. A new op that cannot be sharded on read declares it itself. --- src/transformers/core_model_loading.py | 40 ++++++++++++++------------ tests/utils/test_core_model_loading.py | 18 ++++++------ 2 files changed, 31 insertions(+), 27 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 6d86c37bb3b9..3b95c9d6d944 100644 --- a/src/transformers/core_model_loading.py +++ b/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, @@ -1269,20 +1287,6 @@ def _job(): return _job -def shards_after_conversion(mapping: WeightRenaming | WeightConverter, placements, ndim: int) -> bool: - """Whether a DTensor parameter with these `placements` has to be converted in full before this rank's shard is - taken. - - `PermuteForRope` 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. A transpose is not handled here: it only moves axes - around, so the target shard maps to a source slice and shard-on-read still applies (#48373). - """ - if not isinstance(mapping, WeightConverter): - return False - shard_dims = {placement.dim % ndim for placement in placements if placement.is_shard()} - return any(isinstance(op, PermuteForRope) for op in mapping.operations) and 0 in shard_dims - - def dot_natural_key(s: str): """ Sort key for state-dict names: split on `"."` and sort digits numerically and strings alphabetically. It emits a @@ -1732,8 +1736,8 @@ def convert_and_load_state_dict_in_model( sharding_op = None # 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 shards_after_conversion( - mapping, empty_param.placements, empty_param.ndim + if is_dtensor(empty_param) and not mapping.shards_after_conversion( + empty_param.placements, empty_param.ndim ): sharding_op = DtensorShardOperation(empty_param) @@ -1779,8 +1783,8 @@ 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 shards_after_conversion( - mapping, empty_param.placements, empty_param.ndim + 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) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 2f1433593021..02efa38a045a 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -44,7 +44,6 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, - shards_after_conversion, spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig @@ -345,15 +344,16 @@ def test_rope_permute_is_sharded_after_conversion(self): 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(shards_after_conversion(transpose, [Shard(0)], ndim=2)) - self.assertFalse(shards_after_conversion(transpose, [Replicate(), Shard(1)], ndim=2)) - self.assertFalse(shards_after_conversion(transpose, [Replicate()], ndim=2)) - self.assertTrue(shards_after_conversion(permute, [Shard(0)], ndim=2)) - self.assertFalse(shards_after_conversion(permute, [Shard(1)], ndim=2)) + 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(shards_after_conversion(experts, [Shard(0)], ndim=3)) - self.assertFalse(shards_after_conversion(experts, [Shard(0), Shard(2)], ndim=3)) - self.assertFalse(shards_after_conversion(merge, [Shard(0)], ndim=3)) + 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())