LLM Training Parallelism: Data, Tensor, Pipeline, Expert

Before choosing a parallelism strategy, account for where the memory goes, because each strategy attacks a different line item. Training state divides into four categories. Parameters: 2 bytes each in bf16. Gradients: another 2 bytes each. Optimizer state: with Adam in mixed precision you keep a 4-byte fp32 master weight plus two 4-byte moments, so 12 bytes per parameter. Activations: everything cached during the forward pass for the backward pass to consume, which scales with batch size times sequence length times hidden dimension times layer count — and is the one term you control at runtime.

Add the first three and you get the well-known figure of about 16 bytes per parameter of persistent state. For our 70B model that is roughly 1.1TB, spread across 64 GPUs with 80GB each — 5.1TB of aggregate HBM. It fits in aggregate, which is precisely the point: the job is not to shrink the model but to arrange the shards so that no device ever needs the whole thing at once, and so that the communication required to reassemble pieces on demand does not exceed the compute it enables.

That framing tells you what each strategy is for. Data parallelism replicates the model and splits the batch — it addresses throughput, and in its sharded form also the 16-bytes-per-parameter problem. Tensor parallelism splits individual weight matrices — it addresses a single layer being too large. Pipeline parallelism splits the layer stack — it addresses the total depth exceeding one device. Expert parallelism splits the experts of a Mixture-of-Experts layer — it addresses a model whose parameter count is deliberately much larger than its per-token compute. Activation memory is attacked separately, by recomputation and by context parallelism.

Plain distributed data parallel replicates the entire model on every GPU, splits the global batch across them, and all-reduces gradients before the optimizer step. It is simple and it scales throughput linearly until the interconnect saturates. It also does nothing about memory: every rank still holds the full 16 bytes per parameter, so DDP alone cannot train a model that does not fit on one device.

ZeRO-style sharding fixes that by observing that each replica's copy of the optimizer state is redundant. Stage 1 shards optimizer states across ranks, stage 2 adds gradients, stage 3 adds the parameters themselves. PyTorch's FSDP implements the stage-3 idea. In FSDP2's fully_shard, parameters become DTensors chunked along dimension 0, and the runtime installs hooks: a pre-forward hook all-gathers a module's parameters into their unsharded form, the forward runs on complete weights, then the unsharded copies are freed. Backward re-all-gathers what it needs, computes gradients, and reduce-scatters them so each rank keeps only its own shard.

The trade is explicit in the docs — sharding saves memory "at the cost of communication." You have replaced one gradient all-reduce per step with an all-gather before forward, an all-gather before backward, and a reduce-scatter after. That is roughly 1.5× the communication volume of DDP, hidden behind compute only if you overlap it well and if your interconnect is fast enough. On 400Gb/s InfiniBand with a well-tuned prefetch it overlaps; on 25GbE it does not, and you will watch GPUs sit at 30% utilization waiting for weights.

Tensor parallelism cuts individual weight matrices across devices so a single layer's math is executed cooperatively. Megatron-LM established the pattern that everyone still uses, and its elegance is in the pairing. In the feed-forward block, the first linear is split column-wise — each rank owns a slice of the output dimension and produces a slice of the intermediate activation. The second linear is split row-wise — each rank owns the matching slice of the input dimension, so it can consume its local intermediate slice directly.

That pairing is what makes the cost bearable. Because the column-parallel output feeds a row-parallel input, no synchronization is needed between the two matmuls. Each rank computes a partial sum of the block's output, and one all-reduce at the end of the block produces the correct result on every rank. The attention block follows the same shape, splitting the query, key, and value projections by head and row-splitting the output projection. Net cost: one all-reduce per block in the forward pass and one in the backward pass, or two all-reduces per transformer layer per direction.

That is a lot of collectives — one every few milliseconds of compute, on the critical path, with every rank blocked until it completes. This is why the operational rule is nearly universal: keep the tensor-parallel group inside a single node, where NVLink provides hundreds of gigabytes per second between GPUs. Stretch a TP group across nodes over Ethernet and the all-reduces stop overlapping with anything; the model will train, slowly and expensively. TP degree is therefore usually 2, 4, or 8 — bounded by GPUs per node, not by what the math allows.

Pipeline parallelism assigns contiguous groups of layers to different devices. With 80 layers and a pipeline degree of 8, each stage owns 10 layers. Communication is minimal and cheap: each stage sends its output activations to the next stage — a point-to-point transfer, not a collective, and only at stage boundaries. That is why pipeline parallelism is the strategy that tolerates slow interconnects and is normally used to span nodes.

The cost is idle time. A naive pipeline has every stage waiting for the one before it, so only one device works at a time. The fix is micro-batching: split the global batch into m micro-batches and push them through so stages overlap. But the pipeline still has to fill at the start and drain at the end, and those two periods are pure idleness — the bubble. For a GPipe-style schedule the bubble fraction is (p − 1) / m, where p is the number of stages. With 8 stages and 8 micro-batches, 47% of your GPU time is spent doing nothing. With 8 stages and 64 micro-batches, it drops to about 11%.

Two refinements matter in practice. The 1F1B schedule interleaves forward and backward passes so a stage runs a backward as soon as one is available; this does not change the bubble fraction but sharply reduces peak activation memory, because fewer forward results are held waiting. Interleaved (virtual) pipeline parallelism gives each device several non-contiguous chunks of the model instead of one contiguous block, cutting the bubble to roughly (p − 1) / (m · v) for v chunks per device — at the price of more point-to-point messages. The knob you reach for first is always m: more micro-batches, smaller bubble, until micro-batches get too small to keep the GPU's matmul units busy.

A Mixture-of-Experts layer replaces one feed-forward block with many smaller ones and routes each token to a small number of them. Total parameters climb while per-token compute stays flat — the mechanism Mixture of Experts Explained covers in detail. That design creates a parallelism opportunity that dense models do not have: since only a few experts fire per token, the experts can live on different devices.

Expert parallelism places disjoint subsets of experts on each rank. The communication pattern is unique to this strategy: after the router picks destinations, an all-to-all dispatches each token's hidden state to the rank that owns its chosen expert; after the experts compute, a second all-to-all combines results back to the token's original rank. Two all-to-alls per MoE layer, and all-to-all is the least forgiving collective — its cost is governed by the slowest pairwise link in the group.

The failure mode follows directly from the routing. If the router sends 30% of tokens to one popular expert, the rank hosting it does far more work than its peers, and because both all-to-alls are synchronization points, every other rank waits for it. Load-balancing auxiliary losses exist precisely to spread routing probability, and capacity factors cap how many tokens an expert will accept — with the honest consequence that tokens over capacity get dropped or passed through unmodified. Watch tokens-per-expert as a first-class training metric; a skew that looks minor on a chart is a straggler that shows up as a throughput cliff.

Trigger: a team moved a 70B pre-training run from 8 GPUs on one node to 64 GPUs across 8 nodes. The config kept tensor parallel degree 8, which had worked well within the single node, and added pipeline degree 4 and data parallel degree 2 to fill the new hardware. The launcher assigned ranks in the default sequential order.

Symptom: aggregate throughput rose by about 1.6× despite an 8× increase in hardware. nvidia-smi showed GPUs oscillating between 95% and near zero. Step time was three times the projection, and the variance between steps was large.

Root mechanism: rank placement put each tensor-parallel group across nodes rather than within one. With sequential rank assignment and TP=8 as the innermost dimension, ranks 0–7 of a TP group landed on eight different hosts. Every one of the two all-reduces per transformer layer — 160 collectives per forward pass across 80 layers — now traversed the inter-node fabric instead of NVLink. Tensor parallelism issues its collectives on the critical path with no opportunity to overlap, so the GPUs stalled on every single layer. Pipeline parallelism, which would have tolerated the slow links happily, was placed on the fast ones.

Evidence and recovery: the diagnostic was an NCCL topology dump plus a per-collective timing profile, which showed all-reduce latency roughly two orders of magnitude worse than the intra-node baseline. Recovery was purely a placement change — pin TP groups to the 8 GPUs within a host, put pipeline stages across hosts, and let data parallelism be the outermost dimension. Same code, same hyperparameters, throughput rose to the projected range. Prevention: a startup assertion that fails the job if any tensor-parallel group spans more than one node, and a pre-flight bandwidth probe between every rank pair in each TP group.

Reach for them in a fixed order, because each one has a cheaper predecessor. Start with the settings that cost nothing structurally: gradient checkpointing and a smaller micro-batch, which attack activation memory directly. Then add FSDP, which handles the 16-bytes-per-parameter problem and scales cleanly to a few hundred GPUs if the fabric is fast. Only when a single layer exceeds a device do you need tensor parallelism, and only when the layer stack exceeds a node do you need pipeline parallelism. Expert parallelism is not optional or ordered — it is required if and only if you are training an MoE model.

Combined configurations are named by how many of these run at once. Three dimensions — data, tensor, pipeline — is the standard dense recipe. Add expert parallelism for MoE and you have four. Add context parallelism, which splits the sequence axis so very long contexts fit, and you have five. Public descriptions of frontier training runs generally report configurations in this shape, though the exact degrees are workload-specific and rarely transferable.

The evaluation rule is simple and worth writing into your runbook before the first large job: measure model FLOPs utilization, not GPU utilization. A GPU can sit at 100% busy while spending most of its cycles in NCCL kernels. Compute the achieved fraction of peak matmul throughput, compare it to the same model at a smaller scale, and treat any large gap as a communication or placement bug until proven otherwise. For the serving-side counterpart of these trade-offs, GPU vs CPU for AI Inference covers where the constraints differ.

You can predict most of a configuration's behavior with arithmetic before touching a cluster. Write a short script that takes parameter count, GPU memory, tensor-parallel degree, pipeline degree, and micro-batch count. Have it output four numbers: persistent state per GPU at 16 bytes per parameter divided by the sharding degrees, the resulting free memory for activations, the pipeline bubble fraction as (p − 1) / m, and the number of layer-level collectives per forward pass as two times layer count divided by pipeline degree.

Run it for our reference case: 70B parameters, 80GB GPUs, TP 8, PP 4, DP 2, 32 micro-batches, 80 layers. Expected success: persistent state per GPU comes out well under 80GB, leaving headroom for activations; the bubble fraction is about 9%; and each pipeline stage issues roughly 40 all-reduces per forward pass. Those are numbers you can sanity-check against a real profile.

Now break it on purpose: set micro-batches to 4 while keeping pipeline degree 4. The bubble fraction jumps to 75% — three quarters of your cluster idle — and no learning-rate schedule or fused kernel will recover it. Then set TP to 16 on nodes with 8 GPUs and add the assertion that a TP group must fit inside one node; the script should refuse the configuration outright. The pass criterion: your script rejects both bad configurations with a specific reason, and its predicted memory for the good configuration lands within about 20% of what a real run reports. When your arithmetic and your profiler disagree by more than that, you have found either a bug in the script or an unaccounted buffer in the framework — and both are worth knowing before a week-long job starts.