LUNAROPS · OPERATIONAL UPLINK 100% UPTIME 1,247d POSTS 893 JEFF.MOON@LUNAROPS.DEV UTC --:--:--

Distributed Training Parallelism

17 min read AI & Machine Learning
distributed-traininggpuparallelismfsdpdeepspeedmegatronmoe
Contents

Every distributed training strategy is an answer to one of two questions: the model doesn’t fit, or the model fits but training is too slow. Those are different problems with different solutions, and the most common mistake in scaling a training job is reaching for the wrong axis — sharding a model that fits fine when you just needed more data throughput, or throwing more data-parallel replicas at a model that was never going to fit on one GPU in the first place. There are exactly five axes you can split along — data, tensor, pipeline, expert, and sequence — and at thousands of GPUs you are using four or five of them at once, composed in a specific nested order that is dictated entirely by which communication operations are cheap on which network links. This post is the map: what each axis splits, what it costs in communication, and why the order you nest them in is not a matter of taste.

The thing that makes this hard is that GPUs are fast and networks are slow, and the gap is enormous. An H100 does on the order of 1,000 TFLOP/s of useful BF16 work; the NVLink fabric inside a node moves ~900 GB/s per GPU; the InfiniBand fabric between nodes moves ~50 GB/s per GPU at best. Every parallelism axis introduces a collective communication operation, and the entire art is arranging things so that the heavy, frequent collectives ride the fast intra-node links and only the light, infrequent ones cross the slow inter-node fabric. Get that mapping wrong and your $40,000-per-hour cluster sits at 20% utilization waiting on all-reduces. Get it right and you hold 50-55% of peak FLOPs across thousands of GPUs, which is roughly the ceiling anyone achieves in practice.

The two questions and the five axes

Start with what each axis actually partitions. A transformer training step has three things consuming GPU memory: the model parameters, the optimizer states (with Adam, two extra tensors per parameter — momentum and variance — usually in FP32, so optimizer state is the single largest memory consumer), and the activations (the intermediate tensors saved during the forward pass for use in the backward pass, which scale with batch size and sequence length). Each parallelism axis attacks a different combination of these.

WHAT EACH AXIS SPLITS, AND WHAT IT COSTS

axis        splits...                 communication per step      rides on
--------    ----------------------    ------------------------    ---------
data (DP)   the batch (replicas)      all-reduce gradients        inter-node OK
  + ZeRO/   ...also opt state/params  all-gather + reduce-scatter inter-node OK
    FSDP
tensor(TP)  individual matmuls        all-reduce per layer (x2)   NVLink ONLY
pipeline    layers into stages        point-to-point activations  inter-node OK
  (PP)
expert(EP)  MoE experts across GPUs   all-to-all token routing    NVLink pref.
sequence    the sequence dimension    all-gather / ring P2P        NVLink pref.
  (SP/CP)

The “rides on” column is the whole game. Tensor parallelism communicates twice per transformer layer — that is dozens of all-reduces of full-width activation tensors every single forward and backward pass, the heaviest communication pattern of any axis. It is only viable inside a single node where NVLink gives you ~900 GB/s. Pipeline parallelism, by contrast, only passes the activation tensor at the boundary between stages — a couple of point-to-point sends per micro-batch — so it tolerates the slow inter-node fabric fine. That single fact determines the standard nesting order, which we will build up axis by axis.

Data parallelism: the default, and where ZeRO changes everything

Data parallelism is the one everyone starts with. Replicate the entire model on every GPU, give each replica a different slice of the batch, run forward and backward independently, then all-reduce the gradients so every replica applies the same averaged update. This is DistributedDataParallel (DDP) in PyTorch, and it is the right default whenever the model, its gradients, and its optimizer states fit on a single GPU. Communication is one gradient all-reduce per step, it overlaps cleanly with the backward pass (PyTorch fires the all-reduce for each bucket of gradients as soon as that bucket is ready, while the rest of the backward is still computing), and it scales near-linearly to large GPU counts. If DDP works for your model, use it and stop reading about the other four axes.

The problem is that pure DDP replicates everything, and the redundancy is brutal. A 7B-parameter model in mixed precision needs roughly 14 GB for BF16 parameters, 14 GB for BF16 gradients, and — with Adam — about 56 GB for FP32 optimizer states plus an FP32 master copy of the weights. That is ~84 GB of model state replicated identically on every GPU, before a single activation. On 64 GPUs you are storing 64 identical copies of that 84 GB. ZeRO (Zero Redundancy Optimizer, the idea behind DeepSpeed-ZeRO and PyTorch’s FSDP) observes that this is insane and shards the redundant state across the data-parallel group instead of replicating it.

ZeRO comes in three escalating stages. Stage 1 shards only the optimizer states — the biggest consumer — cutting per-GPU model-state memory dramatically at almost no communication cost. Stage 2 also shards the gradients. Stage 3 (equivalent to FSDP’s full sharding) also shards the parameters themselves: each GPU permanently holds only 1/N of the weights, and just before a layer runs, the group all-gathers that layer’s full parameters, computes, then immediately frees the non-local shards. The trade is explicit — you replaced replication with communication. Stage 3 / FSDP adds an all-gather in the forward pass, an all-gather and a reduce-scatter in the backward, roughly 1.5x the communication volume of plain DDP, in exchange for cutting model-state memory by a factor of N. For the deep mechanics of FSDP sharding strategies, mixed precision, and NCCL tuning, see the dedicated post on FSDP and DDP patterns; here the point is just that ZeRO/FSDP is still data parallelism — it has not split the model’s computation, only its storage, and it still cannot help when a single layer’s activations or a single matmul exceeds one GPU.

ZeRO stage shards per-GPU model state (7B, N large) extra comms vs DDP
DDP (none) nothing ~84 GB baseline
Stage 1 optimizer states ~32 GB negligible
Stage 2 + gradients ~18 GB negligible
Stage 3 / FSDP + parameters ~84/N GB ~1.5x

Tensor parallelism: splitting the matmul itself

When a model is large enough that even sharded parameters or a single layer’s activations blow past one GPU, you have to split the computation of a layer across GPUs. Tensor parallelism (introduced at scale by Megatron-LM) does this by partitioning the weight matrices of the attention and MLP blocks. The clever part is choosing the partition so that consecutive matmuls need synchronization only once. In the MLP block, the first linear layer is split column-wise (each GPU computes a slice of the hidden dimension), the GELU is applied locally with no communication, and the second linear layer is split row-wise so that summing the partial outputs reconstructs the result — one all-reduce. Attention splits the same way across heads: each GPU owns a subset of attention heads end to end, then one all-reduce combines them.

The cost is that all-reduce, and it happens twice per transformer layer (once in attention, once in the MLP) in the forward pass and twice more in the backward. For a model with dozens of layers that is dozens of all-reduces of full activation tensors per step, each one a hard synchronization barrier where every GPU in the tensor-parallel group must wait for the slowest. This is why tensor parallelism is confined to within a single node: across NVLink at ~900 GB/s it is tolerable; across InfiniBand at ~50 GB/s it would dominate the step time and crater utilization. The practical rule is that tensor-parallel degree should not exceed the number of GPUs sharing an NVLink domain — typically 8 on an HGX node. TP buys you the ability to train layers too big for one GPU and reduces activation memory per GPU, at the price of the heaviest, most latency-sensitive communication in the stack.

Pipeline parallelism: stages, bubbles, and micro-batches

Tensor parallelism splits each layer across GPUs; pipeline parallelism splits the layers themselves into sequential stages, with each stage placed on a different GPU (or different node). GPU 0 holds layers 1-8, GPU 1 holds 9-16, and so on. A batch flows through like an assembly line: the activations at each stage boundary are sent point-to-point to the next stage. Because that is just one send of one activation tensor per boundary, pipeline parallelism is communication-light and the only model-splitting axis that crosses the slow inter-node fabric gracefully — which is exactly why it is the axis used to span nodes.

The catch is the pipeline bubble. If you push one batch through naively, stage 1 works while stages 2-4 sit idle, then stage 2 works while the rest sit idle — utilization is terrible. The fix is to split the batch into many micro-batches and keep the pipeline full, the way a CPU pipeline overlaps instructions. With the GPipe schedule, bubble fraction is (p-1)/(m + p-1) where p is the number of stages and m the number of micro-batches, so you want m several times larger than p. Interleaved/1F1B schedules (one-forward-one-backward, as in Megatron and DeepSpeed) shrink the bubble further and cut the activation-memory spike by starting backward passes earlier.

PIPELINE BUBBLE  (4 stages, time flows right; '.' = idle bubble)

naive (1 batch):
  S0: F . . . B . . .
  S1: . F . . . B . .
  S2: . . F . . . B .
  S3: . . . F B . . .      <- huge idle triangles

micro-batched (8 micro-batches, 1F1B):
  S0: F F F F B F B F B F B B
  S1: . F F F F B F B F B B B
  S2: . . F F F F B F B B B B
  S3: . . . F F B F B B B B B   <- bubble shrinks toward (p-1)/(m+p-1)

Pipeline parallelism’s trade-offs are about scheduling and memory, not raw bandwidth: more micro-batches mean a smaller bubble but more in-flight activations to store, and stages must be balanced in compute or the slowest stage sets the pace. It is indispensable for spanning nodes, but it is the fiddliest axis to tune well.

Expert parallelism: the MoE-specific axis

Mixture-of-Experts models change the arithmetic. Instead of every token passing through one big MLP, an MoE layer has many expert MLPs and a router that sends each token to only one or two of them. Total parameters explode (you might have 64 or 256 experts) but the FLOPs per token stay roughly constant because each token only activates a couple of experts. The natural way to distribute this is expert parallelism: place different experts on different GPUs. Now the router has to physically ship each token’s activation to whichever GPU holds its chosen expert, then ship the result back — and since any token can route to any expert, the communication pattern is an all-to-all, the most demanding collective there is, performed twice per MoE layer (dispatch and combine).

The difficulties of expert parallelism are about load balance and that all-to-all. If the router sends a disproportionate share of tokens to a few popular experts, those GPUs become stragglers and the rest idle — hence auxiliary load-balancing losses and per-expert capacity limits (tokens over capacity are dropped or rerouted) that every production MoE trainer needs. The all-to-all is bandwidth-hungry enough that expert parallelism strongly prefers to stay within NVLink domains where possible, and frameworks overlap the dispatch all-to-all with expert computation to hide it. Expert parallelism pays off precisely when you want a model with a very large parameter count at a fixed compute budget — it is the axis that makes trillion-parameter sparse models trainable — but it adds a routing-and-balancing problem that dense models never have.

Sequence parallelism: when the context is the problem

The fifth axis exists because of long context. Activation memory scales with sequence length, and self-attention’s intermediate memory scales worse; at 128K or 1M tokens, a single sequence’s activations exceed any GPU regardless of how you shard the weights. Sequence parallelism (and its close cousin context parallelism) splits along the sequence dimension — each GPU owns a contiguous chunk of the tokens. The complication is attention: every token attends to every other token, so each GPU needs keys and values from all the other chunks. Ring attention solves this by passing K/V blocks around the GPUs in a ring, overlapping each communication step with computing attention against the block currently held, so the all-gather of K/V is hidden behind compute. A lighter form of sequence parallelism (in Megatron) also shards the layer-norm and dropout activations that tensor parallelism leaves replicated, shaving activation memory at the cost of a couple of extra all-gathers. Sequence parallelism is the axis you add specifically to push context length, and it composes on top of tensor parallelism rather than replacing it.

Composing them: 3D and 4D parallelism

At scale you do not pick one axis — you nest several, and the nesting order is dictated by the communication-cost table from the top of this post. The canonical arrangement for training a large dense model on a multi-node cluster is 3D parallelism: tensor parallelism inside a node (heavy all-reduces on NVLink), pipeline parallelism across a small group of nodes (light point-to-point on InfiniBand), and data parallelism / ZeRO across the remaining replicas of that whole TP×PP block (gradient sync, which overlaps with compute and tolerates the slow fabric). Add expert parallelism for MoE or sequence parallelism for long context and it becomes 4D.

3D PARALLELISM ON A CLUSTER  (TP=8, PP=4, DP=N)

  +----------------- one DATA-PARALLEL replica -----------------+
  | node A (NVLink)   node B          node C          node D    |
  | [TP group of 8]   [TP group 8]    [TP group 8]    [TP grp 8]|
  |  pipeline stage0 -> stage1 ----> stage2 ------> stage3      |
  |  (TP all-reduce    (point-to-point activations between      |
  |   stays on NVLink)  stages, over InfiniBand)                |
  +-------------------------------------------------------------+
        x N replicas, gradients all-reduced across replicas
        (ZeRO/FSDP shards optimizer state within the DP group)

  comms placement:  TP all-reduce  -> NVLink   (heaviest, intra-node)
                    PP activations -> IB        (lightest cross-node)
                    DP gradients   -> IB        (overlaps w/ backward)

The global batch size is the product: effective_batch = micro_batch x grad_accum x DP_degree. The total GPU count is TP x PP x DP (times EP for MoE). Choosing the factorization is the core sizing decision, and the heuristics are fairly stable: set TP to fill the NVLink domain only as far as you must to fit a layer (TP has real efficiency cost from all those all-reduces, so smaller is better — often 2, 4, or 8); set PP to the smallest number of stages that makes the model fit across nodes while keeping enough micro-batches to suppress the bubble; then let DP soak up all remaining GPUs for throughput. Megatron-LM and DeepSpeed both implement this composition; Megatron tends to lead on raw dense-model efficiency, DeepSpeed on ZeRO ergonomics and offload, and they are increasingly interoperable.

Axis Solves Communication Where it rides Add it when
Data + ZeRO/FSDP throughput, redundant memory gradient all-reduce (+all-gather) inter-node OK always start here
Tensor layer/matmul too big all-reduce x2 per layer NVLink only a layer won’t fit / TP cheaper than more PP
Pipeline model too big for a node point-to-point activations inter-node spanning nodes
Expert huge param count, fixed FLOPs all-to-all token routing NVLink preferred training MoE
Sequence context too long ring K/V passing NVLink preferred long-context activations blow up

The engineering reality at thousands of GPUs

The clean diagram above is the easy 10%. The hard 90% at thousands of GPUs is that something is always broken and the run must survive it. Failures are not rare events at this scale — with tens of thousands of GPUs and a mean-time-between-failures per GPU measured in thousands of hours, the aggregate MTBF of the job is measured in hours, so a multi-week run will hit dozens of hardware faults. The practical consequences dominate the actual work:

  • Checkpointing must be fast and frequent, because every failure costs you the work since the last checkpoint, and naive synchronous checkpointing of a multi-terabyte sharded optimizer state can stall the whole cluster for minutes. Asynchronous and sharded checkpointing (each rank writes only its shard, in the background) is mandatory, not a nice-to-have.
  • One straggler stalls everyone. Every collective is a barrier; a single GPU that is thermally throttling, on a degraded NIC, or running ECC-error retries drags the entire job to its speed. Detecting and evicting stragglers is a first-class operational concern, not an afterthought.
  • NCCL and topology tuning is the difference between 30% and 55% MFU. The collectives must be told the real network topology (NVLink domains, rail-optimized InfiniBand, NVSwitch) or they will route all-reduces over slow paths. This is the single highest-leverage tuning surface and it is covered for the data-parallel case in the FSDP/DDP post’s NCCL section.
  • Numerical stability gets harder with scale. Larger global batches, BF16 accumulation, and loss spikes that only appear after thousands of steps mean you keep an FP32 master copy of weights, watch the loss scale, and are ready to roll back to a checkpoint and skip a bad data shard.

Model FLOPs Utilization (MFU) — the fraction of theoretical peak FLOPs actually spent on useful model math — is the number that matters, and 50-55% is a genuinely good result for a large 3D-parallel run. Anything below ~35% means a communication axis is misplaced (almost always tensor parallelism leaking across nodes, or too few micro-batches leaving a fat pipeline bubble). The reason training infrastructure is its own specialty is that getting from a working run to an efficient one is entirely about this placement-and-failure-handling layer, not about the parallelism math, which is the comparatively easy part.

This is also why training and inference are different worlds: inference has no optimizer state, no backward pass, and no gradient all-reduce, so its parallelism story is dominated by latency and KV-cache memory instead — see the KV cache optimization and inference engine posts for that side of the house.

Verdict

There are five axes and two questions. If the model and its Adam states fit on one GPU, use data parallelism — DDP, or ZeRO/FSDP if the redundant optimizer state is the only thing pushing you over — and do not touch the model-splitting axes at all. If the model itself does not fit, you split its computation, and which axis you reach for is decided by the network: tensor parallelism inside a node because its per-layer all-reduces only survive on NVLink, pipeline parallelism across nodes because its point-to-point boundary sends are the only model-split that tolerates InfiniBand, expert parallelism layered in for MoE with its all-to-all routing, sequence parallelism layered in for long context. At scale you compose four or five of them into 3D or 4D parallelism, and the entire skill is placing each axis’s collective on a link fast enough to hide it, while keeping the job alive through the constant hardware failures that are guaranteed at thousands of GPUs. Pick the smallest tensor-parallel degree that makes a layer fit, the fewest pipeline stages that span your nodes, and let data parallelism take the rest — then spend your real effort on NCCL topology, checkpointing, and straggler detection, because that is where the utilization actually lives.

Sources

Comments