Distributed Training Parallelism: Data, Tensor, Pipeline, FSDP/ZeRO and NCCL Collectives (2026)
A 70-billion-parameter model does not fit on one GPU, and not for the reason most people assume. The weights in bf16 are 140 GB, which is already too big for an 80 GB card. But the real wall is the training state: gradients, fp32 master weights and Adam moments push the footprint to roughly 16 bytes per parameter, or about 1.1 TB before a single activation is stored. Every technique in distributed training parallelism exists to split that terabyte, and the compute that touches it, across many devices without letting communication eat the gains.
The hard part is that each split trades memory for network traffic in a different place. Data parallelism copies everything and syncs gradients. Sharding (ZeRO and FSDP) removes the copies at the price of extra all-gathers. Tensor parallelism cuts matrices and needs fast links. Pipeline parallelism cuts layers and leaves bubbles. Real jobs combine all of them.
What this covers: the memory anatomy of training, data parallelism and gradient all-reduce, ZeRO stages and FSDP2, Megatron-style tensor parallelism, pipeline schedules and the bubble formula, sequence and expert parallelism, NCCL ring and tree math, a worked MFU calculation, a decision guide by model size, runnable PyTorch code, and the failure modes that stall large runs.
Context and Background
Until about 2019 most large-model training was plain data parallelism: every GPU held a full replica, processed a different slice of the batch, and averaged gradients. That works until the replica no longer fits. The first escape was model parallelism, splitting the network itself. Megatron-LM (Shoeybi et al., 2019) showed that intra-layer splitting of transformer blocks could train an 8.3-billion-parameter model on 512 GPUs, sustaining 15.1 petaFLOPs with 76% scaling efficiency against a single-GPU baseline, according to the paper’s abstract.
The second escape came from noticing that data parallelism wastes memory. In the ZeRO paper, Rajbhandari et al. observed that every replica stores identical optimizer state, gradients and parameters, and proposed partitioning them across the data-parallel group. The authors report training models above 100B parameters on 400 GPUs; treat that as the authors’ claim for their 2019-2020 hardware, not a current benchmark. PyTorch later absorbed the idea as Fully Sharded Data Parallel (FSDP), and its second-generation fully_shard API is now the default path in the PyTorch ecosystem.
The third piece was scheduling. Narayanan et al. combined data, tensor and pipeline parallelism and reported a 1-trillion-parameter configuration on 3,072 GPUs reaching 502 petaFLOP/s, or 52% of theoretical peak per GPU, per the PTD-P paper. That paper is the origin of the “3D parallelism” vocabulary used today, and its interleaved pipeline schedule is still the reference point for reducing bubbles.
Two later developments shape any 2026 design. Mixture-of-experts models add a fourth axis, expert parallelism, which we cover in our guide to MoE architecture. Long-context training adds a fifth, splitting the sequence itself, as in ring attention and context parallelism. And because large runs last weeks on thousands of GPUs, failure handling matters as much as throughput; our piece on spot GPU checkpointing and preemption covers that side.
A word on framing. People often present these strategies as a menu. They are better understood as answers to three separate questions: how do I fit the state, how do I keep every GPU busy, and how do I keep the network off the critical path. Each strategy answers one question well and the other two badly, which is why the final layout is almost always a composition.
The Memory Anatomy of Training
Direct answer: with mixed-precision Adam, each parameter costs about 16 bytes of persistent training state: 2 for bf16 weights, 2 for bf16 gradients, and 12 for fp32 master weights, momentum and variance. A 70B model therefore needs roughly 1.12 TB of model state before activations, which is why it must be sharded across at least 14 80 GB GPUs even in the best case.

Figure 1: Per-parameter memory in mixed-precision Adam (2 + 2 + 12 bytes) and which parallelism strategy removes which slice.
Figure 1 breaks the 16 bytes into its three pieces and shows which strategy eliminates which. The optimizer state dominates: 12 of 16 bytes, or 75%. This is the single most useful fact in the whole topic, because it explains why ZeRO stage 1, which shards only the optimizer, already recovers most of the available memory.
Model state: the 16-byte rule
The accounting follows the ZeRO paper’s notation. Let Psi be the parameter count. In mixed-precision training you keep fp16 or bf16 copies of parameters (2 Psi bytes) and gradients (2 Psi bytes) for the forward and backward passes. The optimizer needs an fp32 master copy of the weights (4 Psi), plus Adam’s first moment (4 Psi) and second moment (4 Psi). The optimizer total is K Psi with K = 12, so the grand total is (2 + 2 + K) Psi = 16 Psi bytes.
Worked example (exact arithmetic, illustrative model): for Psi = 7B the state is 112 GB, so a single 80 GB GPU cannot even hold the state of a 7B model trained with plain Adam. For Psi = 70B it is 1,120 GB. For Psi = 405B it is 6,480 GB, or 6.48 TB. Variants change K: 8-bit optimizers, bf16 moments, or Adafactor reduce it, and pure-bf16 training without fp32 master weights drops it further at the cost of numerical risk. The 16-byte figure is the safe default, not a law.
Activations: the other half of the footprint
Activations are the intermediate tensors kept for the backward pass. Unlike model state, they scale with batch size, sequence length and depth, not with parameter count. Korthikanti et al. (2022) gave an approximate per-layer activation formula for transformers, roughly s b h (34 + 5 a s / h) bytes in 16-bit precision, where s is sequence length, b micro-batch size, h hidden size and a the number of attention heads. The second term is the attention score matrix, which grows quadratically in s unless you use a fused kernel such as FlashAttention that never materializes it.
Two standard remedies exist. Activation checkpointing (recomputation) discards most activations and recomputes them during backward, trading roughly one extra forward pass, about 33% more compute, for a large memory cut. Sequence parallelism and tensor parallelism divide activations across devices. Because activations depend on micro-batch size, you can also simply shrink the micro-batch and accumulate gradients, which costs throughput only when kernels become too small to saturate the GPU.
Temporary buffers and fragmentation
The remaining memory is the unglamorous part. NCCL needs communication buffers. FSDP needs room to materialize one layer’s full parameters at a time, plus prefetched layers. The CUDA caching allocator fragments under variable-sized allocations, so a job can run out of memory with gigabytes nominally free. Budget 10 to 20% headroom beyond the sum of the parts, and treat any plan that fits “exactly” as one that will fail at the first checkpoint save or the first long sequence.
Data Parallelism, ZeRO and FSDP: Trading Memory for Collectives
Direct answer: data parallelism replicates the model and all-reduces gradients each step. ZeRO and FSDP remove the replication by sharding optimizer state (stage 1), gradients (stage 2) and parameters (stage 3) across the data-parallel group, then all-gathering parameters just in time. Stage 3 cuts per-GPU state by the group size at roughly 1.5 times the communication volume of plain data parallelism.

Figure 2: ZeRO stages 1 to 3 (and FSDP full sharding) shown as what each of N data-parallel ranks stores, with the extra collective each stage adds.
Plain data parallelism and gradient all-reduce
In classic data parallelism (PyTorch DistributedDataParallel, DDP), each of N ranks runs forward and backward on its own micro-batch, producing a full gradient of Psi values. The ranks then average gradients with an all-reduce so that every replica applies the same update and stays identical. DDP buckets gradients in roughly 25 MB chunks by default and launches each bucket’s all-reduce as soon as its gradients are ready, overlapping communication with the remainder of the backward pass.
The communication cost is independent of N for large N. A ring all-reduce moves 2(N-1)/N times the buffer size through each link, which tends to 2 times. For bf16 gradients of a 7B model (14 GB) that is about 28 GB of traffic per GPU per step. Whether that hurts depends entirely on the step time and the link speed, which we quantify in the NCCL section below.
The memory cost is the problem: every rank holds all 16 Psi bytes. Data parallelism scales batch size and throughput, not model size.
ZeRO stages 1, 2 and 3
ZeRO (Zero Redundancy Optimizer) partitions the three state components progressively. With N data-parallel ranks, the per-GPU model state is:
| Stage | What is sharded | Per-GPU state (bytes per param) | Extra collective vs DDP |
|---|---|---|---|
| DDP (baseline) | nothing | 16 | none |
| ZeRO-1 | optimizer state | 4 + 12/N | none (all-gather of updated params replaces part of the work) |
| ZeRO-2 | + gradients | 2 + (2 + 12)/N | reduce-scatter replaces all-reduce |
| ZeRO-3 / FSDP | + parameters | 16/N | all-gather in forward and backward |
The formulas follow directly from the 2 + 2 + 12 split. Worked example (illustrative, 70B parameters, N = 64): DDP needs 1,120 GB per GPU, which is impossible. ZeRO-1 gives 4 x 70 + 12 x 70 / 64 = 280 + 13.1, about 293 GB. ZeRO-2 gives 140 + 14 x 70 / 64 = 140 + 15.3, about 155 GB. ZeRO-3 gives 1,120 / 64 = 17.5 GB. Only stage 3 fits, and it leaves about 60 GB for activations and buffers on an 80 GB card.
The communication story is subtler than the memory story. Stage 1 and stage 2 add no volume: a reduce-scatter of gradients followed by an all-gather of updated parameters costs the same bytes as one all-reduce, because a ring all-reduce is exactly a reduce-scatter followed by an all-gather. Stage 3 adds one extra all-gather, since parameters must be gathered in the forward pass and again in the backward pass. The total is therefore 3 Psi in volume versus 2 Psi for DDP, the 1.5 times figure reported in the ZeRO paper.
How FSDP implements stage 3
PyTorch FSDP is ZeRO-3 in spirit with a different implementation. The original FullyShardedDataParallel wrapper flattened groups of parameters into one FlatParameter per wrapped unit and sharded that tensor. FSDP2, exposed as torch.distributed.fsdp.fully_shard, replaces it. According to the PyTorch documentation, fully_shard shards each parameter individually on dimension 0 as a DTensor, converts parameters in place, and registers hooks that all-gather parameters before forward and backward and free them afterwards.
The documented consequences of per-parameter sharding matter in practice. Sharded state dicts need no communication. Memory management avoids Tensor.record_stream, so usage is more deterministic. Frozen parameters become easier to mix with trainable ones, which helps LoRA-style fine-tuning. The documentation also notes that FSDP2 has no direct full state dict, so you convert DTensor state dicts or use PyTorch Distributed Checkpoint.
Three arguments control the trade-offs. mesh takes a 1D device mesh for full sharding or a 2D mesh for hybrid sharding (HSDP), where you shard within a node and replicate across nodes. reshard_after_forward decides whether parameters are freed after the forward pass: True saves memory but pays a second all-gather in backward, False keeps them resident and skips it. mp_policy sets parameter, reduction and output dtypes. You can also pass an offload_policy to push parameters, gradients and optimizer state to CPU, which is a last resort given PCIe bandwidth.
FSDP vs ZeRO: what actually differs
The honest answer to “FSDP vs ZeRO” is that the algorithm is the same and the engineering differs. DeepSpeed ZeRO is a library with its own engine, configuration file, optional CPU and NVMe offload (ZeRO-Infinity) and tight integration with its own optimizers. FSDP is native to PyTorch, composes with DTensor, torch.compile and the device-mesh abstractions, and is what PyTorch-first stacks such as TorchTitan build on.
| Dimension | DeepSpeed ZeRO | PyTorch FSDP2 |
|---|---|---|
| Stage naming | stage 1, 2, 3 | reshard_after_forward and optimizer sharding; full shard is stage 3 |
| Integration | separate engine and JSON config | native module wrapping with fully_shard |
| Offload | CPU and NVMe | CPU via CPUOffloadPolicy |
| Composability with TP/PP | supported, more config surface | DTensor and device mesh make 2D/3D layouts first-class |
| Checkpointing | ZeRO checkpoint formats | Distributed Checkpoint with DTensor |
Pick by ecosystem, not by algorithm. If your stack is already DeepSpeed (for example with existing launcher scripts and NVMe offload needs), stay. If you want native composition with tensor parallelism through DTensor and torch.compile, FSDP2 is the lower-friction path.
A minimal FSDP2 training skeleton
The pattern below applies fully_shard bottom-up (each transformer block, then the root) and uses a device mesh, as the documentation prescribes. It is a skeleton to adapt, not a tuned recipe; the Block and Model classes are placeholders for your own modules.
import torch, torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
dist.init_process_group("nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank % torch.cuda.device_count())
mesh = init_device_mesh("cuda", (dist.get_world_size(),)) # 1D => full sharding
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
reduce_dtype=torch.float32) # fp32 grad reduce
model = Model().cuda() # placeholder: stack of Block modules
for block in model.layers: # bottom-up: blocks first ...
fully_shard(block, mesh=mesh, mp_policy=mp, reshard_after_forward=True)
fully_shard(model, mesh=mesh, mp_policy=mp) # ... then the root
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
for step, (x, y) in enumerate(loader):
loss = torch.nn.functional.cross_entropy(model(x.cuda()), y.cuda())
loss.backward() # reduce-scatter fires per block
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step(); opt.zero_grad(set_to_none=True)
Note the documented contract: call model(x), not model.forward(x), so the hooks run. Checkpoint with torch.distributed.checkpoint rather than gathering a full state dict on rank 0, which is how a 70B job ends up with an out-of-memory crash at save time.
Tensor, Pipeline, Sequence and Expert Parallelism
Direct answer: tensor parallelism splits individual matrix multiplies across GPUs and needs two all-reduces per transformer layer in each pass, so it belongs inside an NVLink domain. Pipeline parallelism assigns consecutive layers to stages and passes activations point to point, so it tolerates slower links but wastes a bubble of (p-1)/(m+p-1) of the time.
Tensor parallelism: the Megatron column/row split
Megatron-LM’s insight is that a transformer MLP block, Y = GeLU(XA) then Z = YB, can be split so that the nonlinearity needs no communication. Partition A by columns: A = [A1, A2]. Each GPU computes GeLU(X A_i) independently, because GeLU is applied element-wise and each column block yields complete output features. Partition B by rows to match: B = [B1; B2]. Each GPU then computes Y_i B_i, producing a partial sum of Z. A single all-reduce adds the partials. In the backward pass one all-reduce is needed at the mirrored point.
Self-attention splits the same way along the head dimension: each GPU owns a subset of heads, so Q, K and V projections are column-parallel, attention runs locally per head, and the output projection is row-parallel with one all-reduce. Per transformer layer that is two all-reduces in forward (attention and MLP) and two in backward, four in total.
Worked example (illustrative): hidden size h = 8,192, a micro-batch of 4,096 tokens, bf16. Each all-reduce moves a tensor of 4,096 x 8,192 x 2 bytes = 67 MB. With 80 layers and 4 all-reduces per layer that is 320 all-reduces per step, about 21.5 GB of tensor-sized payload per micro-batch, before the ring’s 2(t-1)/t factor. These transfers sit on the critical path because the next matmul needs the result immediately. That is why tensor parallelism degrees are almost always 2, 4 or 8, matching the GPUs in a single NVLink-connected node, and rarely cross nodes.
Sequence parallelism, in the Megatron sense of Korthikanti et al., removes a redundancy in this scheme. LayerNorm and dropout regions are replicated across tensor-parallel ranks and their activations duplicated. Splitting them along the sequence dimension replaces each all-reduce with a reduce-scatter plus an all-gather, which moves the same bytes but divides the activation memory of those regions by the tensor-parallel degree.
Pipeline parallelism, schedules and the bubble
Pipeline parallelism cuts the model into p stages of consecutive layers and splits each batch into m micro-batches that flow through the stages. Communication is only the activation tensor between adjacent stages, a small point-to-point send, so it tolerates InfiniBand across nodes. The cost is the pipeline bubble: at the start stages wait for the first micro-batch to arrive and at the end they wait for the last to drain.
GPipe (Huang et al., 2019) runs all forward micro-batches and then all backward ones. Its bubble fraction is (p-1)/(m+p-1) of total time, and it must hold activations for all m micro-batches at once. The 1F1B schedule (one forward, one backward) introduced in PipeDream-style systems keeps the same bubble but starts backward passes early, so a stage holds at most p in-flight micro-batches of activations rather than m. Interleaved 1F1B, from the PTD-P paper, assigns each GPU v non-contiguous model chunks, which shrinks the bubble to roughly (p-1)/(v m) at the price of v times more point-to-point messages.
Worked numbers (the formula is exact, the configuration is illustrative): with p = 8 stages and m = 32 micro-batches, the bubble is 7/39, or 17.9%. Raise m to 64 and it falls to 7/71 = 9.9%. With interleaving at v = 4 and m = 32, the approximate bubble is 7/128 = 5.5%. The levers are clear: more micro-batches, fewer stages, or interleaving. More micro-batches means a larger global batch, which collides with the critical batch size for optimization quality, so you cannot raise m indefinitely.
The practical headaches are balance and the ends. Stages must carry equal compute or the slowest sets the pace; the embedding and the loss head make the first and last stages unusually heavy, so layer assignment is often uneven. Weight tying between the input embedding and output head forces a gradient exchange between the first and last stage.
Sequence, context and expert parallelism
Context parallelism splits the sequence length across GPUs for attention, passing key/value blocks around a ring so no device stores the full sequence. It matters once sequences reach tens or hundreds of thousands of tokens, where activation memory per layer grows with s. The mechanism and its variants are covered in our post on ring attention for long-context training.
Expert parallelism places different experts of a mixture-of-experts layer on different GPUs. Tokens are routed by a gate and exchanged with an all-to-all collective before and after the expert computation. All-to-all is the most network-hostile collective: every rank talks to every other rank, load imbalance across experts becomes a straggler problem, and capacity limits can drop tokens. Details of routing and balancing losses are in our MoE architecture reference.
Composing the axes: the 3D and 4D layout
A cluster of W GPUs is factored as W = d x t x p (data, tensor, pipeline), with extra factors for context c and expert e where used. Frameworks express this as an N-dimensional device mesh, so each parallelism gets its own process groups.

Figure 3: A device mesh factoring GPUs into tensor-parallel groups on NVLink, pipeline stages across nodes, and data-parallel replicas that shard state with FSDP.
The placement rule follows the bandwidth hierarchy. The most communication-hungry axis, tensor parallelism, gets the fastest links, the 8 GPUs of one node. Pipeline stages, which exchange small activations, span nodes over InfiniBand. Data parallelism, whose gradient collectives overlap with backward compute and can be bucketed, takes the outermost dimension. Sharding with FSDP or ZeRO-1 usually applies along the data axis.
Worked example (illustrative): 1,024 GPUs as t = 8, p = 8, d = 16. Each pipeline replica uses 64 GPUs. Tensor parallelism divides the per-GPU parameter share by 8 and pipeline by another 8, so each GPU holds 1/64 of the weights before data-axis sharding. For a 70B model that is 140 GB / 64, about 2.2 GB of bf16 weights per GPU, and sharding the 12-byte optimizer state across d = 16 brings its share to 12 x 70e9 / (64 x 16) = 0.82 GB per GPU. State is no longer the constraint; communication and bubbles are. A model of this size would usually not need a pipeline at all, which is the central point of the next sections.
NCCL Collectives: The Bandwidth Math Behind Every Strategy
Direct answer: NCCL all-reduce on a ring of n GPUs moves 2(n-1)/n times the buffer size over each link, as a reduce-scatter followed by an all-gather. Time is roughly that volume divided by link bandwidth plus a per-hop latency term. Tree algorithms trade bandwidth for fewer hops, which wins for small messages and large node counts.

Figure 4: Ring all-reduce on four GPUs. Each chunk circulates n-1 steps to be reduced, then n-1 steps to be broadcast, so each GPU sends 2(n-1)/n of the buffer in total.
Ring all-reduce, step by step
Split the buffer of S bytes into n chunks. In the reduce-scatter phase, each GPU sends one chunk to its ring neighbor and adds the chunk it receives to its own copy. After n-1 steps, each GPU holds the fully reduced version of exactly one chunk. In the all-gather phase, those reduced chunks circulate for another n-1 steps until every GPU has all of them. Each step moves S/n bytes per link, so the total per GPU is 2(n-1) x S/n.
The beauty of the ring is that the volume is almost independent of n, and that every link is busy at every step. The cost is latency: 2(n-1) sequential steps. With per-step latency alpha, the time is approximately 2(n-1) alpha + 2(n-1)/n x S/B, where B is link bandwidth. For a 1 GB buffer on 8 GPUs at an effective 300 GB/s, the bandwidth term is 2 x 7/8 x 1/300, about 5.8 ms. For a 1 MB buffer the bandwidth term is microseconds and the 14 alpha latency term dominates, which is why small collectives run far below peak.
Tree algorithms and how NCCL chooses
For small messages or very many ranks, a ring’s 2(n-1) latency steps are too many. NCCL also implements tree algorithms in which data flows up a tree to be reduced and then back down, giving logarithmic depth, and NVIDIA has described double binary trees as a way to use both directions of each link. NCCL selects among algorithms and protocols (Simple, LL, LL128) using a built-in cost model based on topology, message size and rank count, and you can override with the NCCL_ALGO and NCCL_PROTO environment variables when diagnosing. In multi-node jobs, expect rings or trees within a node over NVLink and across nodes over InfiniBand or RoCE, with the inter-node link usually the bottleneck. NVIDIA’s own nccl-tests performance notes are the authoritative reference for measuring this.
Bus bandwidth: the number to compare
Raw algorithm bandwidth (algbw = S / time) is not comparable across collectives or rank counts. nccl-tests defines bus bandwidth (busbw) by multiplying algbw by a correction factor so that it reflects the hardware link utilization: 2(n-1)/n for all-reduce, (n-1)/n for reduce-scatter, all-gather and all-to-all, and 1 for broadcast and reduce. If busbw approaches the link’s physical bandwidth, you are hardware-bound. If it sits far below, look for topology, protocol or configuration problems.
Typical hardware tiers, from vendor specifications (nominal, not measured): an H100 SXM GPU exposes NVLink at 900 GB/s total bidirectional bandwidth, which is about 450 GB/s each way, while a common cluster network gives each GPU a 400 Gb/s NDR InfiniBand NIC, 50 GB/s each way. That roughly 9 to 1 gap is the entire reason for the placement rule in Figure 3.
Worked example (illustrative, 7B model, bf16 gradients, 14 GB, DDP across 64 GPUs on 50 GB/s per-GPU network): the all-reduce moves about 2 x 63/64 x 14 = 27.6 GB per GPU. At an optimistic 80% of line rate (40 GB/s) that is 0.69 s. If forward and backward take 3 s per step, and DDP overlaps most of the all-reduce with backward, the exposed communication may be well under 0.2 s. If the step is only 0.8 s, the exposed time is nearly the whole thing and scaling collapses. The ratio of compute time to communication time, not the absolute bytes, decides.
Overlapping communication with compute
Every strategy above has a trick for hiding collectives behind computation. DDP overlaps bucketed gradient all-reduce with the rest of backward. FSDP prefetches the next layer’s all-gather while the current layer computes; FSDP2 exposes set_modules_to_forward_prefetch and set_modules_to_backward_prefetch for explicit control, noting that passing multiple modules improves overlap at the cost of reserved memory. Tensor parallelism can split the matmul into chunks so that a chunk’s all-reduce overlaps the next chunk’s compute. Pipeline schedules overlap the sends with compute of other micro-batches.
Overlap has a ceiling. GPU streaming multiprocessors that run NCCL kernels are not available to your matmuls, and memory bandwidth is shared, so a perfectly overlapped job still loses a few percent. Overlap also depends on the communication stream not being blocked behind a long compute kernel, which is why kernel launch ordering and CUDA stream priorities matter more than they seem.
Measuring Efficiency: MFU, Step Time and a Worked Example
Model FLOPs utilization (MFU) is the fraction of peak hardware FLOPs your training actually converts into useful model math. The standard estimate uses 6 N floating-point operations per token for a dense transformer with N parameters (2N forward, 4N backward), ignoring the attention term and not counting recomputation as useful work. So:
MFU = (6 x N x tokens per second) / (number of GPUs x peak FLOPs per GPU)
Worked example (illustrative numbers): train a 70B model on 1,024 H100 GPUs at 989 TFLOPS peak dense bf16 per GPU (NVIDIA’s published figure; verify against your SKU). Suppose the job sustains 40% MFU. Aggregate peak is 1,024 x 989e12, about 1.01e18 FLOP/s, so useful throughput is 4.05e17 FLOP/s. Dividing by 6 x 70e9 = 4.2e11 FLOPs per token gives about 964,000 tokens per second. A 15-trillion-token run would take 1.5e13 / 9.64e5 = 1.56e7 seconds, about 180 days. At 30% MFU the same run takes 240 days. Each MFU point at this scale is worth about 4.5 days of wall-clock and a corresponding cloud bill.
For reference, the PTD-P paper reported 52% of theoretical peak per GPU at 1 trillion parameters on 3,072 GPUs, on A100 hardware and 2021 software. Published MFU for large production runs commonly falls in the 30 to 50% range, though exact values depend on model shape, sequence length, precision and how MFU is defined, so compare only like with like. Hardware FLOPs utilization (HFU) differs from MFU by counting recomputed activations; a job with activation checkpointing shows HFU roughly one third higher than MFU.
To debug a low MFU, decompose the step: time in matmul kernels, exposed NCCL time, host-side gaps, and bubble time. A profiler trace (PyTorch Profiler or Nsight Systems) showing NCCL kernels not overlapped with compute is a communication problem; many small kernels with gaps is a launch-overhead problem that torch.compile or CUDA graphs address; uneven step times across ranks is a straggler problem.
Choosing a Strategy by Model Size and Cluster
There is no universally best layout, but there is a sensible escalation ladder. Move to the next rung only when the previous one fails on memory or efficiency, because every added axis adds complexity, failure surface and tuning effort.
| Model size (dense) | Typical GPUs | Starting layout | Why |
|---|---|---|---|
| Under about 1B | 1 to 8 | DDP | State of 16 GB or less fits; all-reduce overlaps well |
| 1B to 13B | 8 to 128 | FSDP/ZeRO-3 full shard, activation checkpointing | State sharding alone fits; no tensor or pipeline complexity |
| 13B to 70B | 64 to 1,024 | HSDP or FSDP with TP of 2 to 8 inside the node | Sharding across many nodes makes all-gathers costly; TP shrinks the shard group |
| 70B to 400B+ | 1,024 and up | TP within node, FSDP or ZeRO-1 on data axis, add PP across nodes if needed | Memory and network hierarchy both bind |
| MoE models | varies | Add expert parallelism, usually with TP and FSDP | Parameters are large but active compute is small |
| Long context (very large s) | varies | Add context parallelism | Activation memory scales with sequence length |
These ranges are rules of thumb from published practice, not thresholds, and they shift with hardware generation. The reasoning matters more than the table.
Why FSDP-first is the modern default
The 3D-parallel recipe from 2021 was designed when GPUs had 40 GB and the network between nodes was the dominant constraint. Today an 80 GB or larger GPU plus FSDP2 plus per-layer prefetch can train models into the tens of billions of parameters with nothing but data-axis sharding. The reason is that FSDP communication, while 1.5 times DDP’s volume, is plain all-gather and reduce-scatter that overlaps well with compute and needs no model surgery. Tensor and pipeline parallelism require changes to model code, balanced stage assignment and careful micro-batch tuning.
The scaling limit of pure FSDP is the all-gather latency and bandwidth at large shard groups. A full-shard group of 1,024 ranks runs collectives across the entire cluster with thousands of small hops. HSDP fixes this by sharding within a smaller group (often one node or a few) and replicating across groups, with a gradient all-reduce only across replicas. That keeps the expensive all-gathers on fast links.
When tensor parallelism earns its complexity
Add tensor parallelism when a single layer’s activations or weights are too large for FSDP’s temporary materialization, when per-GPU batch size at full sharding is so small that kernels run inefficiently, or when you want to lower the number of data-parallel ranks to keep the global batch under the critical batch size. Keep it within one NVLink domain. Crossing nodes with tensor parallelism over InfiniBand typically destroys efficiency because of those four on-critical-path all-reduces per layer.
When pipeline parallelism is justified
Pipeline parallelism is the tool for clusters where inter-node bandwidth is poor, or models so deep that even tensor plus FSDP cannot fit, or when you want to avoid all-gathering enormous parameter shards across a slow network. If your fabric is strong (InfiniBand NDR, rail-optimized), many teams now avoid pipeline parallelism entirely, accepting FSDP plus tensor parallelism, since the bubble and the schedule-tuning effort are real costs. Where it is used, it is paired with interleaving and many micro-batches.
Trade-offs, Gotchas, and What Goes Wrong
Large-scale training fails in ways small-scale testing never reveals. The most important is that every collective is a barrier: if one rank is late or dead, all others wait.
Stragglers. One GPU with thermal throttling, a flaky NIC, or a noisy neighbor slows every synchronous step to its pace. At 1,024 GPUs a 5% slowdown on one device is a 5% slowdown on all of them. Log per-rank step time and compare to the median; automate replacement of persistent outliers.
NCCL hangs. The classic hang is mismatched collectives: one rank skips an all-reduce because of a data-dependent branch, an uneven final batch, or an exception caught on one rank. All others block until the NCCL timeout (default timeouts are configurable, check your version) and then the job dies with an opaque watchdog message. Set NCCL_DEBUG=INFO for topology and transport logs, use PyTorch’s flight recorder where available to see which collective each rank last entered, and make data loaders yield equal step counts per rank. Hangs can also come from network faults; InfiniBand errors show up as retries or timeouts well before a hard failure.
Out-of-memory at checkpoint time. Saving a full, unsharded state dict gathers the whole model onto one rank, which then exceeds memory exactly when training seemed stable. Use sharded checkpoints, and test save and restore, including resharding to a different world size, before the long run starts. Our guide to fault-tolerant checkpointing on preemptible GPUs covers cadence and async saves.
Numerical divergence from reduction dtype. Reducing gradients in bf16 across hundreds of ranks accumulates rounding error. The MixedPrecisionPolicy in FSDP2 lets you set reduce_dtype to fp32 for exactly this reason, at the cost of doubling gradient communication volume. Loss spikes that appear only at scale are a signal to check this first.
Nondeterminism and reproducibility. Different parallel layouts change the order of floating-point reductions, so the same seed produces different loss curves under different layouts. Do not expect bit-for-bit equivalence between a 64-GPU and a 512-GPU run; compare statistically.
Over-sharding. Sharding a small model across many GPUs makes each all-gather tiny and latency-bound. If per-layer parameter shards are under a few megabytes, collectives run at a fraction of line rate. Prefer HSDP or fewer shard ranks.
Hidden coupling. Tied embeddings in pipelines, MoE load imbalance, gradient clipping that needs a global norm (an extra all-reduce), and batch-norm-like statistics all introduce synchronization that does not appear in the simple cost models above. Profile rather than trust the model.
Cost of elasticity. A synchronous job on N ranks cannot lose one rank and continue. Recovery means restart from the last checkpoint, so the real efficiency of a run is MFU multiplied by the fraction of time not lost to failures and restarts. At several thousand GPUs, hardware failure rates make that multiplier a first-order concern.
Practical Recommendations
Start from measurement, not from architecture diagrams. Establish single-GPU throughput and MFU, then scale one axis at a time and track MFU at each step. A layout that looks elegant but delivers 20% MFU loses to a plain FSDP job at 38%.
Use the ladder: DDP if the 16-byte state fits, FSDP2 if it does not, HSDP when full-shard collectives span too many nodes, tensor parallelism inside a node when a layer or batch-size constraint demands it, and pipeline parallelism only when the network forces it. Add context and expert parallelism for the workloads that need them. Always enable activation checkpointing selectively (per block) rather than globally, and prefer fused attention kernels.
Place parallelism axes by bandwidth: tensor inside NVLink, pipeline across nodes, data outermost. Validate the device mesh mapping against real topology with nccl-tests before launching: run all-reduce and all-gather sweeps at the sizes your job uses and compare busbw to expectations.
Plan for failure from day one. Checkpoint to sharded storage, test restart and resharding, and monitor per-rank step time and GPU health.
Checklist before a long run:
- [ ] Computed the 16 x Psi state size and the per-GPU share under the chosen layout, plus 10 to 20% headroom
- [ ] Mesh placement matches topology: tensor inside node, pipeline across nodes, data outermost
- [ ]
nccl-testsbusbw measured for all-reduce, all-gather and reduce-scatter at job-relevant sizes - [ ] Bubble fraction computed for pipeline layouts and micro-batch count chosen against it
- [ ] Gradient reduction dtype set deliberately (fp32 for large rank counts)
- [ ] MFU baseline recorded on a small run, with a profiler trace showing overlap
- [ ] Sharded checkpoint save and restore, including a different world size, tested end to end
- [ ] Equal step counts per rank guaranteed to avoid mismatched collectives
- [ ] NCCL timeout, debug logging and hang diagnostics configured; straggler alarm in place
Frequently Asked Questions
What is the difference between data, tensor and pipeline parallelism?
Data parallelism gives each GPU a full model copy and a different slice of the batch, synchronizing gradients with an all-reduce. Tensor parallelism splits individual matrix multiplications across GPUs inside a layer and needs frequent all-reduces, so it stays within a node. Pipeline parallelism splits the model by layers into stages and streams micro-batches through them, sending only activations between stages but losing some time to the pipeline bubble. Large jobs combine all three.
How much GPU memory does training a model need?
With mixed-precision Adam, plan on about 16 bytes per parameter for model state: 2 for weights, 2 for gradients and 12 for fp32 master weights and two Adam moments. A 7B model needs about 112 GB, and a 70B model about 1.1 TB, before activations. Activations add memory proportional to batch size, sequence length and depth, and checkpointing can reduce them. Sharding with FSDP or ZeRO divides the state across the data-parallel group.
Is FSDP the same as ZeRO stage 3?
Conceptually yes: both shard parameters, gradients and optimizer state across data-parallel ranks and all-gather parameters just in time. Implementation differs. DeepSpeed ZeRO is a separate engine with its own configuration and NVMe offload, while PyTorch FSDP is native, and FSDP2 shards each parameter as a DTensor on dimension 0. Both cost about 1.5 times the communication volume of plain data parallelism. Choose based on your existing stack and composition needs.
What is the pipeline bubble and how do I reduce it?
The bubble is idle time at the start and end of each batch while stages wait for micro-batches to fill or drain the pipeline. For p stages and m micro-batches, it is (p-1)/(m+p-1) of total time under GPipe or standard 1F1B schedules. With 8 stages and 32 micro-batches that is about 18%. Reduce it with more micro-batches, fewer stages or interleaved schedules that give each GPU several model chunks, at the cost of more messages.
How does NCCL all-reduce work?
NCCL all-reduce combines a buffer across GPUs so each ends with the sum. In the ring algorithm it runs a reduce-scatter, where each GPU ends up owning one fully reduced chunk, followed by an all-gather that circulates the chunks. Each GPU sends 2(n-1)/n of the buffer, nearly independent of rank count. NCCL also has tree algorithms with fewer latency steps, and it picks algorithm and protocol based on topology and message size.
What is a good MFU for large language model training?
MFU is useful FLOPs (about 6 x parameters x tokens per second for dense models) divided by aggregate peak FLOPs. Published large-scale runs commonly report roughly 30 to 50%, with the PTD-P paper reporting 52% of peak per GPU on A100s at 1 trillion parameters. Treat that range as a rough guide, because definitions vary and results depend on model shape, sequence length and hardware. A first target on a new stack is above 35%.
Further Reading
- Ring attention and context parallelism for long-context LLM training
- Spot GPU training: checkpointing and preemption for fault-tolerant LLM runs
- Mixture-of-experts LLM architecture and expert parallelism
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al.)
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al.)
- PyTorch FSDP2 fully_shard documentation
- NVIDIA nccl-tests performance notes
By Riju — about
