Distributed Training Strategies

When a model or dataset is too large for a single GPU, training must be distributed across multiple devices.

Data Parallelism (DP)

Each GPU holds a full copy of the model
Split the batch across GPUs; each computes gradients on its shard
All-reduce to average gradients, then each GPU updates identically
Scales well, but memory per GPU = full model size
DDP (DistributedDataParallel) — PyTorch's implementation, overlaps gradient computation with communication

ZeRO (Zero Redundancy Optimizer) Eliminates redundant memory across GPUs in three stages:

Stage 1 — partition optimizer states (e.g., Adam momentum, variance) → ~4x memory reduction
Stage 2 — also partition gradients → ~8x reduction
Stage 3 — also partition parameters → memory scales as $O(1/N)$ per GPU
FSDP (Fully Sharded Data Parallel) — PyTorch's native ZeRO Stage 3

Tensor Parallelism (TP)

Split individual weight matrices across GPUs (e.g., columns of $W$ on different GPUs)
Requires communication within each forward/backward pass (all-reduce per layer)
Best within a single node (needs fast interconnect like NVLink)
Used in Megatron-LM

For an MLP, the efficient pattern is column-parallel → activation → row-parallel, which needs only one all-reduce. Column-sharding $W_\text{up}$ produces a naturally sharded output; row-sharding $W_\text{down}$ produces partial sums ( $y = h_0 W_{2,0} + h_1 W_{2,1}$ ) that are summed with an all-reduce. Attention shards by head (heads are independent → no communication during attention itself), then row-parallel output projection.

class ColumnParallelLinear(nn.Module):
    """Shards out_features across ranks; output is naturally sharded."""
    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features // world_size,
                                bias=False, device=f"cuda:{rank}")
    def forward(self, x):
        return self.linear(x)
 
class RowParallelLinear(nn.Module):
    """Shards in_features; each rank computes a partial sum, then all-reduce."""
    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        self.linear = nn.Linear(in_features // world_size, out_features,
                                bias=False, device=f"cuda:{rank}")
    def forward(self, x):
        partial = self.linear(x)                          # partial result per GPU
        dist.all_reduce(partial, op=dist.ReduceOp.SUM)    # sum across GPUs
        return partial
 
class TensorParallelMLP(nn.Module):
    """column parallel -> activation -> row parallel = one all-reduce."""
    def __init__(self, d_model, d_ff, world_size, rank):
        super().__init__()
        self.fc1 = ColumnParallelLinear(d_model, d_ff, world_size, rank)
        self.fc2 = RowParallelLinear(d_ff, d_model, world_size, rank)
    def forward(self, x):
        x = nn.functional.silu(self.fc1(x))  # (batch, seq, d_ff // world_size)
        return self.fc2(x)                   # (batch, seq, d_model), all-reduced

Pipeline Parallelism (PP)

Assign different layers to different GPUs
Forward pass flows through GPUs sequentially → bubble problem (GPUs idle while waiting)
Micro-batching — split the batch into micro-batches to fill pipeline bubbles
GPipe, PipeDream

3D Parallelism — combine DP + TP + PP for the largest models. Example: GPT-3 training used all three.

Communication primitives:

Broadcast — one GPU sends an identical copy to all others
All-reduce — every GPU gets the sum/average of all GPUs' data
All-gather — every GPU gets the full concatenation (removes sharding along an axis)
Reduce-scatter — combine via reduction, each GPU keeps a shard (adds sharding)
Bandwidth is the bottleneck: NVLink (~900 GB/s) >> InfiniBand (~400 GB/s) >> PCIe (~64 GB/s)

Ring all-reduce = reduce-scatter + all-gather (each GPU only talks to two neighbors). Cost depends on the array size and bandwidth, not the number of devices. These primitives are each other's backward pass: all-gather in the forward → reduce-scatter in the backward (and vice versa), so the backward of an all-reduce is another all-reduce.

Sharding a matmul (mesh axes $X, Y$ ; matrix axes $I, J$ ): write the operands as blocks so the product becomes block matmuls.

Neither operand's contracting dim sharded → no communication; output is naturally sharded.

One operand's contracting dim sharded → all-gather it first, then multiply.
Both contracting dims sharded → each device holds a partial sum; finish with an all-reduce.

Distributed Training Strategies

Linked from