[distributed] Add a new parallel plan for unit sharding

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"
 

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)

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,79 @@ 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
+    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|
+
+    - 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.
+
+    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"]
+        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(
+            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.
+        # 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
+            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)).
 
@@ -295,9 +370,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
@@ -334,6 +406,49 @@ def transform_output_post_forward(self, module, output, mesh):
         return output
 
 
+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"]
+        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()]
+        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.
 
@@ -768,7 +883,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(),
@@ -813,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)
 

tests/tensor_parallel/test_tensor_parallel.py

@@ -11,12 +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 unittest.mock import patch
 
 import torch
+from parameterized import parameterized
 
-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 (
@@ -25,6 +27,7 @@
     PackedColwiseParallel,
     PackedRowwiseParallel,
     RowwiseParallel,
+    _unit_mesh,
 )
 from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test
 
@@ -365,3 +368,110 @@ def test_sharding_does_not_create_unrelated_module_attributes(self):
 
                 self.assertEqual(module.random_attr, 123)
                 self.assertFalse(hasattr(module, "num_experts"))
+
+
+@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 _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 (
+            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 = 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"]
+
+        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 = 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)
+        self.assertEqual(
+            tuple(attention.q_proj._unit_mesh.mesh.shape),
+            tuple(attention.o_proj._unit_mesh.mesh.shape),
+        )