Distributed training

Problem

Training Llama 2 7B (6,738,415,616 parameters) with Adam in mixed precision keeps a bf16 copy of the weights, bf16 gradients, and an fp32 master copy plus two fp32 Adam moments: 2 + 2 + 12 = 16 bytes per parameter, 107.8 GB before a single activation. An 80 GB GPU can't hold it. Even if it could, one GPU doing the trillions of token-steps a pretraining run needs would take years.

So the run is spread over N GPUs, and two questions decide how. First, which of the training state gets replicated on every GPU and which gets split: the batch, the optimizer state, the gradients, the weights, or the layers themselves. Second, what that split costs in communication, because every GPU still has to end each step with the same update the single-GPU program would have made (see the training loop).

Intuition

The gradient of a mean loss is the mean of per-example gradients, so data parallelism gives each of N GPUs a full model copy and 1/N of the batch, then averages the gradients. That's the same gradient as the big batch (the argument gradient accumulation makes for micro-batches in time, made here across space). Every GPU then applies the same update to the same weights, so they stay in sync without ever sending weights.

The waste is that every GPU stores the same optimizer state, gradients and weights. Sharded data parallelism gives each GPU 1/N of them. A GPU only needs a layer's full weights while it computes that layer, so it all-gathers them just before, then frees them. Tensor parallelism splits each weight matrix so a single layer's matmul runs on several GPUs, and pipeline parallelism gives each GPU a contiguous slice of the layers and streams micro-batches through them like an assembly line. Large runs combine all of them: tensor parallel inside a node where links are fastest, pipeline across nodes, and data parallel over the rest.

Mechanics

All-reduce, and the ring that makes it cheap

An all-reduce takes one tensor per rank and leaves every rank holding their elementwise sum (dist.all_reduce, default op=ReduceOp.SUM). It splits into two halves that are collectives in their own right: a reduce-scatter leaves rank r holding the sum of only its 1/p chunk (reduce_scatter_single), and an all-gather copies every rank's chunk to everyone (all_gather_single). In a ring, each of p ranks sends one chunk of N/p elements to its neighbour per step: p − 1 steps of reduce-scatter, then p − 1 of all-gather.

bytes sent per rank = 2(p − 1)/p · N · bytes per element     (→ 2N as p grows, independent of p)

The point is the last part: each rank's traffic barely grows with the number of ranks, where a naive "send everything to rank 0 and broadcast back" puts 2(p − 1)N through rank 0's link. NCCL picks rings, trees or other algorithms per message size and topology, but the bandwidth term of any all-reduce can't beat this 2(p − 1)/p factor.

DDP: buckets that overlap with backward

DistributedDataParallel broadcasts rank 0's weights and buffers at construction (_sync_module_states) so every replica starts identical, then registers a hook on each parameter's gradient. Calling all-reduce per parameter would launch thousands of tiny collectives, and waiting until backward finishes would leave the network idle while it runs. So DDP groups gradients into buckets of bucket_cap_mb (default 25 MiB). Parameters are assigned in roughly reverse order of definition, because backward produces the last layer's gradients first (_ddp_init_helper builds the Reducer from that order). When every gradient in a bucket is ready, its all-reduce starts asynchronously while backward keeps computing earlier layers. The bucket for the parameters defined first, whose gradients arrive last, can be kept small (_BucketCapacityConfig), so little communication is left once backward ends.

The reducer averages: it sums and divides by the world size. With gradient_as_bucket_view=True, param.grad is a view into the bucket, saving one gradient-sized copy. The no_sync context skips the all-reduce for gradient accumulation; the flag is read in _pre_forward.

ZeRO: shard the optimizer state, then gradients, then weights

ZeRO (Rajbhandari et al.) counts the model states for Adam in mixed precision as 2Ψ bytes of fp16 weights, 2Ψ of fp16 gradients and KΨ of optimizer state, with K = 12 (fp32 weights, momentum and variance, 4 bytes each), for Ψ parameters. Data parallelism replicates all (4 + K)Ψ on every GPU. ZeRO removes the replication in three stages:

DDP: (2 + 2 + K)Ψ    stage 1 (P_os): 4Ψ + KΨ/N    stage 2 (P_os+g): 2Ψ + (2 + K)Ψ/N    stage 3 (P_os+g+p): (2 + 2 + K)Ψ/N

Stage 1 gives each rank the Adam state for 1/N of the parameters; it updates only that slice and then all-gathers the new weights. Stage 2 also keeps only that slice's gradient: instead of an all-reduce, a reduce-scatter hands each rank the summed gradient for exactly the slice it updates. The traffic is the same as DDP, since an all-reduce is a reduce-scatter plus an all-gather. Stage 3 also shards the weights. A layer's weights are all-gathered before its forward, freed, all-gathered again before its backward, and its gradients are reduce-scattered. That's one extra all-gather, 1.5× DDP's traffic, for model states that shrink linearly with N.

FSDP in PyTorch

torch.distributed.fsdp exports two APIs. FSDP1's FullyShardedDataParallel names the ZeRO stages in ShardingStrategy: FULL_SHARD (stage 3), SHARD_GRAD_OP (stage 2: parameters are sharded outside computation but stay unsharded from forward through backward), NO_SHARD (DDP) and HYBRID_SHARD (full sharding inside a node, replication across nodes). FSDP2's fully_shard shards each parameter on dim 0 as a DTensor. Its docstring (lines 116–123) is the ZeRO-3 schedule: all-gather before forward, free after forward if reshard_after_forward, all-gather again in backward, then free and reduce-scatter the gradients.

Each fully_shard call makes one communication group, so it is applied bottom-up, once per transformer block and then once on the root. Each block's weights then come in one foreach_all_gather and leave in one foreach_reduce. Peak memory is the shards plus one or two blocks' full weights. reshard_after_forward is the stage-2/stage-3 dial: False keeps the gathered weights until backward (no second all-gather, more memory), and an integer reshards to a smaller group such as the GPUs of one node. For gradient accumulation, set_requires_gradient_sync(False) is FSDP2's equivalent of no_sync. It skips the reduce-scatter, so every rank keeps unsharded gradients until the last micro-batch. That trades the memory ZeRO saved back for communication.

Tensor parallelism: split the matmuls

Megatron-LM (Shoeybi et al.) splits a transformer's MLP Z = ReLU(XA)B so that it needs one all-reduce. Split A by columns into [A₁, A₂]: each GPU computes Y_i = ReLU(XA_i) with no communication, because the activation works elementwise on columns it owns. Split B by rows to match: each GPU computes Y_iB_i, and Z = Σ_i Y_iB_i is a single all-reduce. Splitting A by rows instead would leave partial sums before the nonlinearity, and ReLU(a + b) ≠ ReLU(a) + ReLU(b). Attention is split the same way: heads go column-wise across GPUs, and the output projection is row-wise.

PyTorch spells the two halves ColwiseParallel (input replicated, output sharded on the last dim) and RowwiseParallel (input sharded on the last dim, output replicated), applied with parallelize_module. The cost is two all-reduces of the b × s × h activations per layer in forward (attention and MLP) and two in backward, on the critical path with nothing to overlap. That is why tensor parallelism stays inside a node, on NVLink.

Pipeline parallelism: split the layers

With p stages, stage s can't start a micro-batch until stage s − 1 has sent its activations, and backward runs in the reverse direction. GPipe (ScheduleGPipe) runs every micro-batch's forward, then every backward. 1F1B (Schedule1F1B) runs min(n_microbatches, num_stages − stage_index) warm-up forwards, then alternates one backward and one forward. With equal forward and backward times, both take 2(m + p − 1) slots for 2m slots of work, so both leave the same idle bubble:

bubble fraction = (p − 1) / (m + p − 1)     (Narayanan et al. write (p − 1)/m as bubble time over ideal time)

What 1F1B changes is memory. A stage must keep each micro-batch's activations from its forward until its backward. GPipe holds all m, while 1F1B frees one per forward after warm-up, so stage 0 never holds more than p. More micro-batches shrink the bubble without growing 1F1B's activation memory, and interleaved schedules (ScheduleInterleaved1F1B) give each GPU several non-adjacent chunks of layers to shrink it further, at the cost of more sends.

Worked example

Four parts: a ring all-reduce by hand, the memory per GPU for Llama 2 7B on 8 GPUs under each scheme, a tensor-parallel MLP split two ways, and the bubble in a 4-stage pipeline. Every number below comes from tested helpers in ml_models_docs.py.

Implementation

FSDP2 on one node, sharding each transformer block, with bf16 compute and fp32 gradient reduction. Launch with torchrun --nproc-per-node 8 train.py:

import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

dist.init_process_group("nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

model = build_model()                       # on the meta device for large models
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
for block in model.layers:                  # bottom-up: one all-gather per block
    fully_shard(block, mp_policy=mp)
fully_shard(model, mp_policy=mp)            # root: embeddings, final norm, head

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)   # sees only shards

for step, micro_batches in enumerate(loader):
    for i, (x, y) in enumerate(micro_batches):
        model.set_requires_gradient_sync(i == len(micro_batches) - 1)
        loss = model(x, labels=y) / len(micro_batches)
        loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # DTensor-aware
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

Swap the two fully_shard lines for model = DistributedDataParallel(model, device_ids=[local_rank]), and set_requires_gradient_sync for no_sync, to get plain data parallelism. Add tensor parallelism by building a 2-D DeviceMesh and calling parallelize_module with ColwiseParallel on the up projections and attention input projections and RowwiseParallel on the down and output projections before fully_shard on the data-parallel dimension. Each rank must also see a different slice of data (DistributedSampler, see data loading), or all of them compute the same gradient.

Tradeoffs

• Data parallel is simplest and fastest when it fits. One overlapped all-reduce per step and no change to the model, but full model states on every GPU.
• Sharding buys memory with latency. ZeRO-3/FSDP sends 1.5× DDP's bytes and puts an all-gather in front of every block, which prefetching must hide. Small blocks, slow links or tiny per-GPU batches expose it. Stage 2 or reshard_after_forward=False sits in between.
• Tensor parallelism needs fast links. Its all-reduces are on the critical path and scale with batch × sequence × hidden, so it stays within NVLink domains (typically ≤ 8 ways) and needs the hidden size and head count to divide evenly.
• Pipelines need many micro-batches. The bubble is (p − 1)/(m + p − 1), so deep pipelines need large global batches, and stages must be balanced or the slowest one sets the pace. The traffic, activations between neighbours only, is the lightest of all, which suits slower inter-node links.
• The global batch grows with data parallelism. 1,024 GPUs at 4 sequences each is a 4,096-sequence batch, which may need learning-rate retuning or exceed the critical batch size (see gradient accumulation).
• Failures scale too. More GPUs mean more frequent hardware faults; frequent sharded checkpoints and elastic restarts become part of the design.

Connections

References

• Rajbhandari, Rasley, Ruwase and He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC, 2020. arXiv:1910.02054.
• Shoeybi, Patwary, Puri, LeGresley, Casper and Catanzaro. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arXiv:1909.08053.
• Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC, 2021. arXiv:2104.04473 (1F1B, interleaving, the bubble).
• Huang et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS, 2019. arXiv:1811.06965.
• Li et al. PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB, 2020. arXiv:2006.15704 (DDP's buckets).
• Zhao et al. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB, 2023. arXiv:2304.11277.
• PyTorch v2.15.0: torch/nn/parallel/distributed.py, torch/distributed/fsdp/__init__.py, _fully_shard.py, tensor/parallel/style.py, pipelining/schedules.py.