Tensor Parallelism

Data parallelism requires the model to fit on one device. When a single layer does not, the matrices themselves have to be split.

There are two ways to cut a matmul, and the trick is that they compose.

A column split leaves each rank with a slice of the output and needs no communication. A row split leaves each rank with a partial sum of the whole output and needs an all-reduce.

Why the feed-forward block costs one collective

Column-split the up-projection, row-split the down-projection.

$X W_{\text{up}}$ leaves rank $i$ holding columns $i$ of the hidden activation — which is exactly the input the row-split down-projection wants. The nonlinearity is elementwise, so it applies to the slice without any exchange. One all-reduce at the end of the block, and none in the middle.

Attention has the same structure for free: heads are independent, so column-splitting $W_Q, W_K, W_V$ by head and row-splitting the output projection gives one all-reduce per attention block. Two per layer forward, two more in the backward pass.

Sequence parallelism, the companion

The all-reduce boundaries leave the layer norms and dropout replicated — every rank computes the same thing on the same full activation, wasting both compute and, more importantly, memory.

Korthikanti et al.’s addition splits those regions along the sequence dimension instead, and replaces one of the all-reduces with a reduce-scatter and the other with an all-gather. Same total bytes moved, and the replicated activations disappear.

Since the all-reduce was already going to be paid, splitting it in two and doing the sequence-parallel region in between costs nothing. It is on by default in Megatron and there is no reason to disable it.

import torch
import torch.distributed as dist
from torch import Tensor, nn


class ColumnParallel(nn.Module):
    """Split the output dimension. No collective; each rank keeps its slice."""

    def __init__(self, in_f: int, out_f: int, world: int, rank: int):
        super().__init__()
        assert out_f % world == 0
        self.w = nn.Parameter(torch.empty(out_f // world, in_f))

    def forward(self, x: Tensor) -> Tensor:
        return torch.nn.functional.linear(x, self.w)


class RowParallel(nn.Module):
    """Split the input dimension. Each rank produces a partial sum."""

    def __init__(self, in_f: int, out_f: int, world: int, rank: int):
        super().__init__()
        assert in_f % world == 0
        self.w = nn.Parameter(torch.empty(out_f, in_f // world))

    def forward(self, x: Tensor) -> Tensor:
        y = torch.nn.functional.linear(x, self.w)
        dist.all_reduce(y)                    # the one collective in the block
        return y


class ParallelFFN(nn.Module):
    def __init__(self, dim: int, hidden: int, world: int, rank: int):
        super().__init__()
        self.up = ColumnParallel(dim, hidden, world, rank)
        self.down = RowParallel(hidden, dim, world, rank)

    def forward(self, x: Tensor) -> Tensor:
        # No exchange between them: the activation is elementwise, so each
        # rank's slice of the hidden state is exactly what its rows need.
        return self.down(torch.nn.functional.silu(self.up(x)))

The real implementation needs an autograd function on each boundary: the column-parallel forward is an identity that becomes an all-reduce in the backward, and the row-parallel forward is an all-reduce that becomes an identity. Getting those the wrong way round gives correct forward passes and wrong gradients, which is the worst available failure mode.

Tensor parallelism is the only strategy that shrinks a single layer’s memory during its computation, which is why it is unavoidable for very wide models and for low-latency serving. It is also the most communication-intensive per unit of compute, which is why it is confined to a node and composed with the others rather than scaled on its own.