From 281d01dfe21432682e9c9a3517750611e221306c Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 15 Sep 2026 17:11:34 +0900 Subject: [PATCH 01/12] initial draft --- .../distributed/tensor_parallel.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 3151f6264a12..b5d0bb0dce79 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -241,6 +241,25 @@ def transform_output_post_forward(self, module, output, mesh): return output.to_local() if self.use_local_output else output +class UnitColwiseParallel(TensorParallelLayer): + """This forces all ranks to own at least 1 unit of the tensor. + It is designed to be used on `query/key/value` layers when running TP. You cannot + For `query = nn.Linear(hidden_dim, head_dim * num_heads)` you can't divide the second dimension blindly by the world size as the attention REQUIRES to split per-head. You cannot split 2 heads of dim 64 into 4 GPUs by just separating into 4 x 32. This would mix the heads. + + The unit colwise handles 2 cases: + - Less heads than world size: 2 heads, 4 process -> we replicate the heads. + Rank: | 0 | 1 | 2 | 3 | + This gives: |h0 | h0 | h1 | h1| + And we normalize the contributions from rank 0,1 and from 2,3. + - More heads than world size: 6 heads, 4 process -> we also replicate the heads. + Rank: | 0 | 1 | 2 | 3 | + This gives: |h0h1|h2h3|h4h5|h4h5| + And we "ignore" the results from rank 2,3. + """ + + pass + + class RowwiseParallel(TensorParallelLayer): """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). From 40569a1339a19c4ef890bea1a26a9410eb263c18 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 10:43:23 +0900 Subject: [PATCH 02/12] draft 1. --- src/transformers/core_model_loading.py | 1 + src/transformers/distributed/sharding_utils.py | 5 ++++- .../distributed/tensor_parallel.py | 18 +++++++++++++++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 580f46849c3c..ab130c7beae3 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1724,6 +1724,7 @@ def convert_and_load_state_dict_in_model( # accelerate will take care of putting back on correct device after loading # Note that we only do it with `device_map` but not with `tp_plan`, as tp will perform local sharding before, so memory # spike during conversion ops should be fine + # TODO: move commment to the doc of force cpu, the flow is broken. if sharding_op is None and isinstance(mapping, WeightConverter) and mapping.force_cpu: param_device = "cpu" diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 083fbe5f915f..304ca135baa7 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -65,6 +65,10 @@ class DtensorShardOperation: | TP + FSDP, same tensor dim (fused/interleaved TP) | [_StridedShard(d, sf=tp_size), Shard(d)] | | TP + FSDP, different dims | [Shard(d1), Shard(d2)] | + There is an additional case when we have to force non-empty chunks. This will transform any Shard(n) placements + into a [Shard(n), Replicate()]. + This is intended to solved cases where you have 6 key value heads and 8 process -> 2 process shard(n) replicated 4 times. + Loading (this class) -------------------- During from_pretrained, each rank looks up every tensor key, but does not load the weight bytes yet (safetensors get_slice). @@ -138,7 +142,6 @@ def shard_tensor( has_strided_shard = any(not placement.is_shard() for _, placement in dim_placements) # finally fetch from the disk only the slices - # finally fetch from the disk only the slices if has_strided_shard: # Multi-interval dim: read each piece separately, then concatenate. return self._slice_and_cat(source, intervals_by_dim, device, dtype) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index b5d0bb0dce79..59640aa01c9b 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -255,9 +255,25 @@ class UnitColwiseParallel(TensorParallelLayer): Rank: | 0 | 1 | 2 | 3 | This gives: |h0h1|h2h3|h4h5|h4h5| And we "ignore" the results from rank 2,3. + + 1. Update the rank's expected shape. + 2. Prepare the slices appropriately when sharding the param. + + We need to have [Shard(0), Replicate()] . """ - pass + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + placements = [Shard(meta.ndim - 2), Replicate()] + remainder = meta.shape[meta.ndim - 2] % mesh.world_size + split = meta.shape[meta.ndim - 2] // remainder + mesh = mesh.reshape(remainder, split) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, placements, src_data_rank=None), + requires_grad=meta.requires_grad, + ) class RowwiseParallel(TensorParallelLayer): From 6b865e983cfa880ce01ca7adaa4a510ea0ac16c9 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 10:52:32 +0900 Subject: [PATCH 03/12] update --- src/transformers/distributed/tensor_parallel.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 59640aa01c9b..f39a6da4d615 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -263,15 +263,24 @@ class UnitColwiseParallel(TensorParallelLayer): """ def shard_param(self, module, param, mesh): + from torch.distributed.device_mesh import DeviceMesh + meta = module._parameters.get(param) if meta is None: return placements = [Shard(meta.ndim - 2), Replicate()] remainder = meta.shape[meta.ndim - 2] % mesh.world_size split = meta.shape[meta.ndim - 2] // remainder - mesh = mesh.reshape(remainder, split) + # We need to create a device mesh anew + new_mesh = mesh.mesh.reshape(remainder, split) + new_device_mesh = DeviceMesh.from_group( + mesh.get_group(), + device_type=mesh.device_type, + mesh=new_mesh, + mesh_dim_names=[mesh.mesh_dim_names[0], "unit_replicate"], + ) module._parameters[param] = torch.nn.Parameter( - distribute_tensor(meta, mesh, placements, src_data_rank=None), + distribute_tensor(meta, new_device_mesh, placements, src_data_rank=None), requires_grad=meta.requires_grad, ) From 9a15232aba88e2af4eb37d31c4200c00d02c1448 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 11:19:28 +0900 Subject: [PATCH 04/12] up --- src/transformers/distributed/tensor_parallel.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index f39a6da4d615..bf2ca66ce907 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -250,11 +250,10 @@ class UnitColwiseParallel(TensorParallelLayer): - Less heads than world size: 2 heads, 4 process -> we replicate the heads. Rank: | 0 | 1 | 2 | 3 | This gives: |h0 | h0 | h1 | h1| - And we normalize the contributions from rank 0,1 and from 2,3. - - More heads than world size: 6 heads, 4 process -> we also replicate the heads. - Rank: | 0 | 1 | 2 | 3 | - This gives: |h0h1|h2h3|h4h5|h4h5| - And we "ignore" the results from rank 2,3. + + - More heads than world size: 6 heads, 8 process -> we also replicate the heads. + Rank: | 0 | 1 | 2 | 3 | + This gives: |h0h1h2|h0h1h2|h3h4h5|h3h4h5| 1. Update the rank's expected shape. 2. Prepare the slices appropriately when sharding the param. From 332beee134778ec5398c18fe20ba26693eef4d6d Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 14:05:13 +0900 Subject: [PATCH 05/12] this should be better? --- src/transformers/distributed/tensor_parallel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index bf2ca66ce907..47d8d1a48183 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -268,15 +268,15 @@ def shard_param(self, module, param, mesh): if meta is None: return placements = [Shard(meta.ndim - 2), Replicate()] - remainder = meta.shape[meta.ndim - 2] % mesh.world_size - split = meta.shape[meta.ndim - 2] // remainder + remainder = mesh._layout.numel() % meta.shape[meta.ndim - 2] + split = mesh._layout.numel() // remainder # We need to create a device mesh anew new_mesh = mesh.mesh.reshape(remainder, split) - new_device_mesh = DeviceMesh.from_group( - mesh.get_group(), + new_device_mesh = DeviceMesh( device_type=mesh.device_type, mesh=new_mesh, mesh_dim_names=[mesh.mesh_dim_names[0], "unit_replicate"], + _init_backend=False, ) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, new_device_mesh, placements, src_data_rank=None), From d7543fe22e745659e232a6045ac1ea5581396417 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 14:13:19 +0900 Subject: [PATCH 06/12] let's introduce the unit dimension? --- src/transformers/distributed/tensor_parallel.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 47d8d1a48183..96fd22b8a53a 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -268,7 +268,10 @@ def shard_param(self, module, param, mesh): if meta is None: return placements = [Shard(meta.ndim - 2), Replicate()] - remainder = mesh._layout.numel() % meta.shape[meta.ndim - 2] + unit_number = meta.shape[meta.ndim - 2] // module.unit_dim + + remainder = mesh._layout.numel() % unit_number + # if the remainder is 0, there is nothing to do. split = mesh._layout.numel() // remainder # We need to create a device mesh anew new_mesh = mesh.mesh.reshape(remainder, split) From 89d72b7b5a867b3857235c18d460c540297d1565 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 14:23:42 +0900 Subject: [PATCH 07/12] inherit --- src/transformers/distributed/tensor_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 96fd22b8a53a..c798ed2976e6 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -241,7 +241,7 @@ def transform_output_post_forward(self, module, output, mesh): return output.to_local() if self.use_local_output else output -class UnitColwiseParallel(TensorParallelLayer): +class UnitColwiseParallel(ColwiseParallel): """This forces all ranks to own at least 1 unit of the tensor. It is designed to be used on `query/key/value` layers when running TP. You cannot For `query = nn.Linear(hidden_dim, head_dim * num_heads)` you can't divide the second dimension blindly by the world size as the attention REQUIRES to split per-head. You cannot split 2 heads of dim 64 into 4 GPUs by just separating into 4 x 32. This would mix the heads. From 9992df7bee16e34e389065be0b15fe943e8a6fb8 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 16:02:26 +0900 Subject: [PATCH 08/12] add rowise --- .../distributed/tensor_parallel.py | 102 ++++++++++++++---- 1 file changed, 82 insertions(+), 20 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index c798ed2976e6..4a2e273f8671 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -14,6 +14,7 @@ from __future__ import annotations import contextlib +import math import re from ..utils import logging @@ -28,6 +29,7 @@ if is_torch_distributed_available(): import torch.distributed as dist + from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor from torch.distributed.tensor.placement_types import _StridedShard @@ -241,6 +243,21 @@ def transform_output_post_forward(self, module, output, mesh): return output.to_local() if self.use_local_output else output +def _unit_mesh(mesh, unit_number: int): + """ + A 2-D unit-view of `mesh`. We reshape the mesh as (units, replicas), so every rank owns units. + """ + world = mesh.size() + groups = math.gcd(unit_number, world) + replicas = world // groups + outer = mesh.mesh_dim_names[0] if mesh.mesh_dim_names else "unit" + return DeviceMesh( + device_type=mesh.device_type, + mesh=mesh.mesh.reshape(groups, replicas), + mesh_dim_names=[outer, "unit_replicate"], + ) + + class UnitColwiseParallel(ColwiseParallel): """This forces all ranks to own at least 1 unit of the tensor. It is designed to be used on `query/key/value` layers when running TP. You cannot @@ -262,30 +279,37 @@ class UnitColwiseParallel(ColwiseParallel): """ def shard_param(self, module, param, mesh): - from torch.distributed.device_mesh import DeviceMesh - meta = module._parameters.get(param) if meta is None: return - placements = [Shard(meta.ndim - 2), Replicate()] - unit_number = meta.shape[meta.ndim - 2] // module.unit_dim - - remainder = mesh._layout.numel() % unit_number - # if the remainder is 0, there is nothing to do. - split = mesh._layout.numel() // remainder - # We need to create a device mesh anew - new_mesh = mesh.mesh.reshape(remainder, split) - new_device_mesh = DeviceMesh( - device_type=mesh.device_type, - mesh=new_mesh, - mesh_dim_names=[mesh.mesh_dim_names[0], "unit_replicate"], - _init_backend=False, - ) + weight = module._parameters["weight"] + unit = weight.shape[-1] // module.unit_dim + unit_mesh = _unit_mesh(mesh, unit) + module._unit_mesh = unit_mesh module._parameters[param] = torch.nn.Parameter( - distribute_tensor(meta, new_device_mesh, placements, src_data_rank=None), + distribute_tensor(meta, unit_mesh, [Shard(meta.ndim - 2), Replicate()], src_data_rank=None), requires_grad=meta.requires_grad, ) + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + # The inputs are on a different mesh. We HAVE to swap the meshs as cross mesh coms + # do not work. + x = args[0] + if isinstance(x, DTensor): + replicated = [Replicate()] * x.device_mesh.ndim + if x.placements != tuple(replicated): + x = x.redistribute(placements=replicated) + x = x.to_local() + if not self.should_use_local_tensors(module): + x = DTensor.from_local(x, module._unit_mesh, [Replicate(), Replicate()], run_check=True) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + # Already on the unit mesh, correctly laid out; the parent handles the local case. + if isinstance(output, DTensor): + return output + return super().transform_output_post_forward(module, output, mesh) + class RowwiseParallel(TensorParallelLayer): """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). @@ -341,9 +365,6 @@ def context_around_forward(self, module, mesh): if not self.should_use_local_tensors(module): yield else: - # A rowwise local forward must produce only its partial matmul. If we don't hide - # the bias, we will be adding the bias x world_size times which is not correct. - # We should add it once after the all_reduce (redistribute). bias = module._parameters.get("bias") if bias is not None: module._parameters["bias"] = None @@ -380,6 +401,45 @@ def transform_output_post_forward(self, module, output, mesh): return output +class UnitRowwiseParallel(RowwiseParallel): + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + weight = module._parameters["weight"] + unit = weight.shape[-1] // module.unit_dim + module._unit_mesh = _unit_mesh(mesh, unit) + # bias is added once after the reduce, so it stays replicated + placements = [Replicate(), Replicate()] if param == "bias" else [Shard(meta.ndim - 1), Replicate()] + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, module._unit_mesh, placements, src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if self.should_use_local_tensors(module): + return (x.to_local() if isinstance(x, DTensor) else x,) + args[1:], kwargs + # `UnitColwiseParallel` already leaves its output on an equivalent unit mesh. + if not isinstance(x, DTensor): + x = DTensor.from_local(x, module._unit_mesh, [Shard(-1), Replicate()], run_check=True) + elif x.device_mesh != module._unit_mesh: + x = DTensor.from_local(x.to_local(), module._unit_mesh, [Shard(-1), Replicate()], run_check=True) + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + # The parent reduces over the whole world, which would also sum the replicas. + if not isinstance(output, DTensor): + dist.all_reduce(output, group=module._unit_mesh.get_group(0)) + if (bias := module._parameters.get("bias")) is not None: + output = output + (bias.to_local() if isinstance(bias, DTensor) else bias) + return output if self.use_local_output else DTensor.from_local(output, mesh, [Replicate()], run_check=True) + # Resolve Partial while the unit axis still exists, then swap back (metadata only). + local = output.redistribute(placements=[Replicate(), Replicate()]).to_local() + output = DTensor.from_local(local, mesh, [Replicate()], run_check=True) + return output.to_local() if self.use_local_output else output + + class ReplicatedWithGradAllReduce(TensorParallelLayer): """Replicated parameter whose gradient is partial. @@ -814,7 +874,9 @@ class ParallelInterface(GeneralInterface): "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), + "unit_colwise": UnitColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), "rowwise": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + "unit_rowwise": UnitRowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), "rowwise_split_input": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "rowwise_rep": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "packed_colwise": PackedColwiseParallel(), From 8367d34720933451aceac2d96264fe3d2b922be7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 16:16:23 +0900 Subject: [PATCH 09/12] update --- src/transformers/distributed/tensor_parallel.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 4a2e273f8671..f5330ac3c710 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -283,7 +283,7 @@ def shard_param(self, module, param, mesh): if meta is None: return weight = module._parameters["weight"] - unit = weight.shape[-1] // module.unit_dim + unit = weight.shape[0] // module.unit_dim unit_mesh = _unit_mesh(mesh, unit) module._unit_mesh = unit_mesh module._parameters[param] = torch.nn.Parameter( @@ -294,6 +294,7 @@ def shard_param(self, module, param, mesh): def transform_inputs_pre_forward(self, module, args, kwargs, mesh): # The inputs are on a different mesh. We HAVE to swap the meshs as cross mesh coms # do not work. + # Swapping meshes is metadata only while the tensor is replicated, so gather first. x = args[0] if isinstance(x, DTensor): replicated = [Replicate()] * x.device_mesh.ndim From b72a595737e82003d886fd2984d9ed0062e43da9 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 16:39:50 +0900 Subject: [PATCH 10/12] small fix --- src/transformers/distributed/tensor_parallel.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index f5330ac3c710..e1c39a037109 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -278,12 +278,16 @@ class UnitColwiseParallel(ColwiseParallel): We need to have [Shard(0), Replicate()] . """ + needs_config = True + def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return weight = module._parameters["weight"] - unit = weight.shape[0] // module.unit_dim + config = module.config + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + unit = weight.shape[0] // head_dim unit_mesh = _unit_mesh(mesh, unit) module._unit_mesh = unit_mesh module._parameters[param] = torch.nn.Parameter( @@ -403,12 +407,16 @@ def transform_output_post_forward(self, module, output, mesh): class UnitRowwiseParallel(RowwiseParallel): + needs_config = True + def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return weight = module._parameters["weight"] - unit = weight.shape[-1] // module.unit_dim + config = module.config + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + unit = weight.shape[-1] // head_dim module._unit_mesh = _unit_mesh(mesh, unit) # bias is added once after the reduce, so it stays replicated placements = [Replicate(), Replicate()] if param == "bias" else [Shard(meta.ndim - 1), Replicate()] @@ -922,6 +930,9 @@ def apply_tensor_parallelism(model, tp_mesh): style_name = _get_parameter_tp_plan(parameter_name=full, tp_plan=model.tp_plan, is_weight=True) if style_name is not None and style_name in ALL_PARALLEL_STYLES: style = ALL_PARALLEL_STYLES[style_name] + if getattr(style, "needs_config", False): + # Head-aware styles need the head size before they can size the shard. + module.config = model.config.get_text_config() style.validate_param(module, p_name, tp_mesh, parameter_name=full) style.shard_param(module, p_name, tp_mesh) From b3176ffdd382885f73faee5e378e5cadd1948781 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 16:43:10 +0900 Subject: [PATCH 11/12] Test unit sharding layouts Covers replicated, non-dividing, coprime, exact, GQA and MQA head counts. --- tests/tensor_parallel/test_tensor_parallel.py | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index c83741361fd5..3b03f753e339 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -11,10 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import math import warnings +from types import SimpleNamespace from unittest.mock import patch import torch +from parameterized import parameterized +from torch import nn from transformers import AutoModelForCausalLM from transformers.distributed import tensor_parallel @@ -25,6 +29,7 @@ PackedColwiseParallel, PackedRowwiseParallel, RowwiseParallel, + _unit_mesh, ) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test @@ -365,3 +370,114 @@ def test_sharding_does_not_create_unrelated_module_attributes(self): self.assertEqual(module.random_attr, 123) self.assertFalse(hasattr(module, "num_experts")) + + +class MockAttention(nn.Module): + """The four projections `unit_colwise`/`unit_rowwise` are meant to be used on.""" + + def __init__(self, num_heads, num_key_value_heads, head_dim, hidden_size): + super().__init__() + self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False) + # `apply_tensor_parallelism` attaches this for styles that set `needs_config`. + config = SimpleNamespace( + head_dim=head_dim, + hidden_size=hidden_size, + num_attention_heads=num_heads, + num_key_value_heads=num_key_value_heads, + ) + for proj in (self.q_proj, self.k_proj, self.v_proj, self.o_proj): + proj.config = config + + +@is_tensor_parallel_test +class TestUnitParallel(TestCasePlus): + """`unit_colwise`/`unit_rowwise` must give every rank whole heads, and never split one.""" + + HEAD_DIM = 8 + HIDDEN = 64 + + class FakeMesh: + """Enough of a DeviceMesh for `_unit_mesh` to reshape, with no process groups.""" + + def __init__(self, ranks, rank=0): + self.mesh, self.rank = ranks, rank + self.device_type, self.mesh_dim_names = "cpu", None + self.ndim = ranks.ndim + + def size(self): + return int(self.mesh.numel()) + + def get_local_rank(self): + return self.rank + + def _shard(self, module, name, style, mesh): + """Shard one projection, returning the placements it asked for.""" + with ( + patch.object( + tensor_parallel, "DeviceMesh", lambda device_type, mesh, mesh_dim_names=None: self.FakeMesh(mesh) + ), + patch.object( + tensor_parallel, "distribute_tensor", side_effect=lambda tensor, *a, **kw: tensor + ) as distribute, + ): + style.shard_param(module, name, mesh) + return distribute.call_args.args[2] + + @parameterized.expand( + [ + (2, 2, 4), # fewer heads than ranks -> replicated + (6, 6, 4), # heads do not divide the world + (7, 7, 4), # coprime -> fully replicated + (8, 8, 4), # exact fit, no replication + (12, 12, 8), + (2, 2, 8), + (32, 8, 8), # GQA + (32, 8, 16), # GQA, more ranks than kv heads + (8, 1, 4), # MQA + ] + ) + def test_every_rank_owns_whole_heads(self, num_heads, num_key_value_heads, world_size): + mesh = self.FakeMesh(torch.arange(world_size)) + with patch.object( + tensor_parallel, "DeviceMesh", lambda device_type, mesh, mesh_dim_names=None: self.FakeMesh(mesh) + ): + for units in (num_heads, num_key_value_heads): + groups, replicas = _unit_mesh(mesh, units).mesh.shape + self.assertEqual(groups * replicas, world_size, "the mesh must cover every rank") + self.assertEqual(groups, math.gcd(units, world_size), "replication must be the least possible") + self.assertEqual(units % groups, 0, "a head may not be split across ranks") + self.assertGreaterEqual(units // groups, 1, "no rank may be left without a head") + + @parameterized.expand([(2, 2, 4), (6, 6, 4), (8, 8, 4), (32, 8, 8), (32, 8, 16), (8, 1, 4)]) + def test_projections_are_sharded_on_the_head_axis(self, num_heads, num_key_value_heads, world_size): + attention = MockAttention(num_heads, num_key_value_heads, self.HEAD_DIM, self.HIDDEN) + mesh = self.FakeMesh(torch.arange(world_size)) + colwise, rowwise = ALL_PARALLEL_STYLES["unit_colwise"], ALL_PARALLEL_STYLES["unit_rowwise"] + + for name, units in (("q_proj", num_heads), ("k_proj", num_key_value_heads), ("v_proj", num_key_value_heads)): + proj = getattr(attention, name) + placements = self._shard(proj, "weight", colwise, mesh) + groups = math.gcd(units, world_size) + self.assertEqual(placements[0].dim, 0, f"{name} shards its output features") + self.assertTrue(placements[1].is_replicate(), f"{name} replicates along the second axis") + self.assertEqual(proj.weight.shape[0] // groups % self.HEAD_DIM, 0, f"{name} split a head") + + placements = self._shard(attention.o_proj, "weight", rowwise, mesh) + self.assertEqual(placements[0].dim, 1, "o_proj shards its input features") + self.assertTrue(placements[1].is_replicate()) + self.assertEqual(attention.o_proj.weight.shape[1] // math.gcd(num_heads, world_size) % self.HEAD_DIM, 0) + + @parameterized.expand([(2, 4), (6, 4), (8, 4), (12, 8), (32, 16)]) + def test_colwise_and_rowwise_agree_on_the_mesh(self, num_heads, world_size): + """o_proj consumes q_proj's output, so the two must land on the same unit mesh.""" + attention = MockAttention(num_heads, num_heads, self.HEAD_DIM, self.HIDDEN) + mesh = self.FakeMesh(torch.arange(world_size)) + self._shard(attention.q_proj, "weight", ALL_PARALLEL_STYLES["unit_colwise"], mesh) + self._shard(attention.o_proj, "weight", ALL_PARALLEL_STYLES["unit_rowwise"], mesh) + self.assertEqual( + tuple(attention.q_proj._unit_mesh.mesh.shape), + tuple(attention.o_proj._unit_mesh.mesh.shape), + ) From 27cfbf8f058c8677f46d1ed3895457dbc8922e11 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 16 Sep 2026 17:04:12 +0900 Subject: [PATCH 12/12] Build the unit sharding tests on a real single-layer Llama Replaces the hand-rolled attention mock with LlamaModel, so the styles read the config the same way they do in a real plan. --- tests/tensor_parallel/test_tensor_parallel.py | 44 ++++++++----------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 3b03f753e339..e302000fb87b 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -13,14 +13,12 @@ # limitations under the License. import math import warnings -from types import SimpleNamespace from unittest.mock import patch import torch from parameterized import parameterized -from torch import nn -from transformers import AutoModelForCausalLM +from transformers import AutoModelForCausalLM, LlamaConfig, LlamaModel from transformers.distributed import tensor_parallel from transformers.distributed.sharding_utils import DtensorShardOperation from transformers.distributed.tensor_parallel import ( @@ -372,26 +370,6 @@ def test_sharding_does_not_create_unrelated_module_attributes(self): self.assertFalse(hasattr(module, "num_experts")) -class MockAttention(nn.Module): - """The four projections `unit_colwise`/`unit_rowwise` are meant to be used on.""" - - def __init__(self, num_heads, num_key_value_heads, head_dim, hidden_size): - super().__init__() - self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False) - self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) - self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) - self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False) - # `apply_tensor_parallelism` attaches this for styles that set `needs_config`. - config = SimpleNamespace( - head_dim=head_dim, - hidden_size=hidden_size, - num_attention_heads=num_heads, - num_key_value_heads=num_key_value_heads, - ) - for proj in (self.q_proj, self.k_proj, self.v_proj, self.o_proj): - proj.config = config - - @is_tensor_parallel_test class TestUnitParallel(TestCasePlus): """`unit_colwise`/`unit_rowwise` must give every rank whole heads, and never split one.""" @@ -413,6 +391,22 @@ def size(self): def get_local_rank(self): return self.rank + def _attention(self, num_heads, num_key_value_heads): + """A real attention layer from a one-layer Llama, with the config the styles read.""" + config = LlamaConfig( + hidden_size=num_heads * self.HEAD_DIM, + intermediate_size=2 * num_heads * self.HEAD_DIM, + num_hidden_layers=1, + num_attention_heads=num_heads, + num_key_value_heads=num_key_value_heads, + head_dim=self.HEAD_DIM, + vocab_size=32, + ) + attention = LlamaModel(config).layers[0].self_attn + for proj in (attention.q_proj, attention.k_proj, attention.v_proj, attention.o_proj): + proj.config = config # `apply_tensor_parallelism` attaches this for `needs_config` styles + return attention + def _shard(self, module, name, style, mesh): """Shard one projection, returning the placements it asked for.""" with ( @@ -453,7 +447,7 @@ def test_every_rank_owns_whole_heads(self, num_heads, num_key_value_heads, world @parameterized.expand([(2, 2, 4), (6, 6, 4), (8, 8, 4), (32, 8, 8), (32, 8, 16), (8, 1, 4)]) def test_projections_are_sharded_on_the_head_axis(self, num_heads, num_key_value_heads, world_size): - attention = MockAttention(num_heads, num_key_value_heads, self.HEAD_DIM, self.HIDDEN) + attention = self._attention(num_heads, num_key_value_heads) mesh = self.FakeMesh(torch.arange(world_size)) colwise, rowwise = ALL_PARALLEL_STYLES["unit_colwise"], ALL_PARALLEL_STYLES["unit_rowwise"] @@ -473,7 +467,7 @@ def test_projections_are_sharded_on_the_head_axis(self, num_heads, num_key_value @parameterized.expand([(2, 4), (6, 4), (8, 4), (12, 8), (32, 16)]) def test_colwise_and_rowwise_agree_on_the_mesh(self, num_heads, world_size): """o_proj consumes q_proj's output, so the two must land on the same unit mesh.""" - attention = MockAttention(num_heads, num_heads, self.HEAD_DIM, self.HIDDEN) + attention = self._attention(num_heads, num_heads) mesh = self.FakeMesh(torch.arange(world_size)) self._shard(attention.q_proj, "weight", ALL_PARALLEL_STYLES["unit_colwise"], mesh) self._shard(attention.o_proj, "weight", ALL_PARALLEL_STYLES["unit_rowwise"], mesh)