Distributed Training & GPU Systems

A frontier model doesn't fit on one GPU, and it wouldn't finish training on one even if it did. This page covers the systems side of training: where the bytes go, how numeric formats trade accuracy for speed, how GPUs talk to each other, and how data, tensor, pipeline, context and expert parallelism fit together into the "3D/4D/5D" setups used on clusters of tens of thousands of GPUs. Interviewers ask about this because it tells them whether you can reason about cost, memory and throughput from first principles instead of treating the training stack as a black box.

The mental model: compute, memory, communication

Every distributed training decision trades off three resources:

• Compute: matmul throughput (H100 SXM peak dense BF16 is about 989 TFLOP/s NVIDIA H100). Training a dense transformer costs about \(6N\) FLOPs per token (2N forward, 4N backward), where N is the parameter count Kaplan+ 2020.
• Memory: HBM capacity (80 GB on H100, more on H200/B200) and HBM bandwidth (about 3.35 TB/s on H100 SXM).
• Communication: bandwidth and latency of the links between GPUs. These are tiered: NVLink inside a node, InfiniBand/RoCE between nodes.

Every parallelism strategy is a way of cutting the computation graph across devices. Each one cuts along a different axis (batch, hidden dimension, layers, sequence, experts), so each one saves a different kind of memory and adds a different communication pattern. If you can say what each strategy shards, which collective it adds, how many bytes that collective moves, and whether it sits on the critical path, you can answer almost any interview question in this area.

Memory accounting for training

GPU memory during training holds four things: parameters, gradients, optimizer state and activations. A fifth category covers temporary buffers, communication buffers and allocator fragmentation, which together often cost 5–15% extra.

Model state: the 16-bytes-per-parameter rule

The standard recipe is mixed-precision Adam(W) Micikevicius+ 2017. Forward and backward run in BF16 (or FP16), and the optimizer keeps an FP32 master copy of the weights, because small updates added to a BF16 weight would be rounded away. Per parameter:

Variations come up a lot. Many frameworks accumulate gradients in FP32 for stability, especially across gradient-accumulation micro-steps, which gives 18 bytes. Some drop the separate BF16 weight copy and cast on the fly. 8-bit optimizers or factored optimizers (Adafactor-style) shrink the 12. Muon-style optimizers keep one momentum buffer rather than two. The 16–18 figure is the default to quote, and it's worth saying out loud which variant you're assuming.

Worked example: 7B and 70B

The ZeRO paper's reference example, 7.5B parameters on 64-way DP, gives 120 GB per GPU for plain DP, 31.4 GB with stage 1, 16.6 GB with stage 2 and 1.9 GB with stage 3 Rajbhandari+ 2019. It's worth memorizing the shape of that curve.

Activations

Activations are the tensors saved during forward so that backward can compute gradients. For a standard GPT-style layer with hidden size \(h\), sequence length \(s\), micro-batch \(b\), \(a\) attention heads and 16-bit activations, Korthikanti et al. derive the per-layer activation memory without any parallelism as Korthikanti+ 2022:

The \(34\,sbh\) term covers the linear-layer inputs, GeLU input, LayerNorm inputs and dropout masks. The \(5as^2b\) term is the attention score matrix (softmax output, dropout mask), which is quadratic in sequence length. FlashAttention removes most of that quadratic term because it never materializes the \(s\times s\) matrix: it recomputes it in tiles during backward and saves only the output and the softmax statistics. With TP degree \(t\) and sequence parallelism, everything is divided by \(t\).

Worked example (7B, Llama-like: \(h=4096\), \(L=32\), \(s=4096\), \(b=1\)): \(sbh = 16.8\)M elements. With FlashAttention, about \(34 \times 16.8\text{M} \approx 0.57\) GB per layer, so roughly 18 GB for 32 layers. Without FlashAttention, the score term adds \(5 \cdot 32 \cdot 4096/4096 = 160\) to the 34, about 3.3 GB per layer and over 100 GB in total. That's why naive attention made long-context training impractical. Activation memory grows linearly with micro-batch and (with FlashAttention) linearly with sequence length, so at 32K–128K context activations dwarf model state, and you need context parallelism or recomputation.

Activation checkpointing (recomputation)

Activation checkpointing (also called gradient checkpointing or rematerialization) stores only a subset of activations during forward, usually each transformer block's input, and recomputes the rest during backward Chen+ 2016. Chen et al. showed that checkpointing every \(\sqrt{L}\) layers gives \(O(\sqrt{L})\) memory for one extra forward pass. In transformers, the usual practice is to checkpoint at block boundaries.

• Full recomputation: save only each layer's input (\(2sbh\) bytes per layer in BF16). Cost is about one extra forward pass, so per-token compute goes from \(6N\) to roughly \(8N\), about 33% more FLOPs.
• Selective recomputation: recompute only the parts that are cheap in FLOPs but expensive in memory, mainly the attention softmax/dropout region. Korthikanti et al. report a ~5× activation memory reduction at about 2–4% compute overhead for large GPT models Korthikanti+ 2022. FlashAttention already does this for attention internally, which is why "selective" in modern stacks often means picking which MLP or norm outputs to keep.
• Offloading: move saved activations to CPU memory over PCIe and fetch them back for backward. This helps when PCIe bandwidth can be overlapped with compute. It's increasingly attractive on Grace-Hopper/Grace-Blackwell systems with fast CPU–GPU links.

Recomputation is also why MFU and HFU differ (see MFU). The recomputed forward does real work on the hardware but doesn't count as useful model FLOPs.

Mixed precision: BF16, FP16, FP8, MX formats

The formats

The FP8 E4M3/E5M2 split was standardized by NVIDIA, Arm and Intel Micikevicius+ 2022. The MX formats were defined in the OCP Microscaling spec Rouhani+ 2023. NVFP4's 16-element blocks with E4M3 scales are described in NVIDIA's format introduction NVIDIA 2025.

FP16 and loss scaling vs BF16

FP16 has only 5 exponent bits, so small gradient values (often below \(2^{-24}\)) underflow to zero. The fix is loss scaling: multiply the loss by a factor \(S\) (e.g. \(2^{16}\)) before backward, which scales every gradient by \(S\) through the chain rule, then divide by \(S\) before the optimizer step Micikevicius+ 2017. Dynamic loss scaling raises \(S\) periodically and halves it (skipping that step) whenever an inf/NaN appears.

BF16 keeps FP32's 8 exponent bits, so its range matches FP32 and underflow is rarely a problem. It needs no loss scaling. The cost is precision: 7 mantissa bits means about 2–3 significant decimal digits, so you need FP32 master weights (a weight of 1.0 plus an update of 1e-4 rounds back to 1.0 in BF16) and FP32 accumulation inside matmuls and reductions. Since A100, BF16 has been the default for LLM training.

FP8 training: scaling granularity is everything

FP8 roughly doubles tensor-core throughput over BF16 on Hopper and Blackwell and halves the memory traffic for the quantized tensors. The hard part is range: E4M3 only spans about 18 binades, so each tensor must be multiplied by a scale that maps its values into representable range. The design choices:

• Per-tensor delayed scaling: the scale comes from an amax history over the last N steps. It's fast, because there's no extra pass over the data, but a sudden outlier gets clipped. This was the original Transformer Engine recipe on H100 TE docs.
• Per-tensor current scaling: compute amax of the current tensor first. It's more robust, at the cost of an extra reduction.
• Fine-grained block scaling: DeepSeek-V3 used tile-wise 1×128 groups for activations (per token, per 128 channels) and 128×128 blocks for weights, with E4M3 everywhere. It promoted partial sums to FP32 on CUDA cores every 128 elements, because Hopper's FP8 tensor-core accumulation has limited precision DeepSeek-AI 2024. Outliers then only damage their own block.
• Hardware microscaling (MXFP8): Blackwell tensor cores natively consume per-32-element E8M0 scales, so fine-grained scaling costs almost nothing. NVIDIA reports MXFP8 pretraining of an 8B model on 15T tokens matching BF16 once the scale-factor rounding scheme is changed from the OCP-suggested one, which diverged (see the paper for the exact rule) Mishra+ 2025.

What stays in higher precision even in "FP8 training": master weights and optimizer state (FP32, sometimes BF16 moments), embeddings and the output head, normalization layers, the attention softmax, and often the last few layers. Only the big GEMMs run in FP8.

FP4 and the state of play in 2026

NVIDIA trained a 12B model on 10T tokens in NVFP4 with loss and downstream accuracy comparable to an FP8 baseline. It took four tricks: random Hadamard transforms to spread outliers, 2D (block) weight quantization so forward and backward see consistent values, stochastic rounding on gradients, and keeping a few sensitive layers in higher precision NVIDIA 2025.

The interconnect hierarchy

A modern training cluster is a hierarchy of increasingly slow links. The rough per-GPU numbers below are for orientation. Exact figures depend on SKU and configuration, so check the vendor spec before quoting them as fact.

NVIDIA lists 900 GB/s per GPU for Hopper NVLink, 1.8 TB/s for Blackwell, and a preliminary figure for Vera Rubin's NVLink 6 NVIDIA NVLink. Llama 3 405B trained over RoCE at 400 Gb/s between GPUs, while smaller Llama 3 models used InfiniBand Llama Team 2024. DeepSeek-V3 quotes about 160 GB/s NVLink versus 50 GB/s InfiniBand on its export-restricted H800 cluster, roughly 3.2×, and designed its MoE routing around that gap DeepSeek-AI 2024.

The key ratio: intra-node bandwidth is about 10× inter-node bandwidth per GPU on a standard H100 HGX system. Rack-scale NVLink (GB200/GB300 NVL72) moves that cliff from 8 GPUs to 72, which matters a lot for TP and EP placement. Scale-out fabrics are usually rail-optimized fat trees: GPU i of every node connects to the same leaf "rail" switch, so collectives among same-index GPUs across nodes take one hop.

Collective communication

The primitives

All-reduce = reduce-scatter followed by all-gather. This identity is the foundation of ZeRO: if you only need your shard of the summed gradient, which you do once the optimizer state is sharded, you can stop after the reduce-scatter and save half the traffic for that phase.

Cost model: ring vs tree

Use the α–β model: sending \(m\) bytes costs \(\alpha + m/B\), where \(\alpha\) is per-message latency (a few µs) and \(B\) is link bandwidth. In a ring all-reduce, the buffer is split into N chunks. In N−1 reduce-scatter steps, each rank sends one chunk to its neighbor and adds the chunk it receives. Then N−1 all-gather steps circulate the finished chunks:

The bandwidth term approaches \(2S/B\) and doesn't grow with N. That's bandwidth-optimal: every byte has to leave and arrive somewhere. The latency term grows linearly with N, which hurts at thousands of ranks and for small messages. Reduce-scatter and all-gather alone each cost half: \((N-1)\alpha + \frac{N-1}{N}\frac{S}{B}\).

Tree algorithms (NCCL uses a double binary tree, so both halves of the bandwidth are used) get latency of \(O(\log N)\) with near-optimal bandwidth for large messages. NCCL picks ring, tree or other algorithms (such as NVLink SHARP on NVSwitch systems, which does the reduction inside the switch) based on message size and topology.

nccl-tests reports algorithm bandwidth (S / time) and bus bandwidth (algBW × \(2(N-1)/N\) for all-reduce). Bus bandwidth is comparable to raw link speed regardless of N, so it's the number to compare against hardware peak nccl-tests docs.

Back-of-envelope: DDP gradient all-reduce for 7B

BF16 gradients: \(S = 14\) GB. Ring traffic per GPU is about \(2S = 28\) GB. Over InfiniBand at 50 GB/s per GPU that's ≈0.56 s per step if nothing overlaps. Inside one node over NVLink at about 450 GB/s, it's ≈0.06 s. If a step's compute takes 2–5 s, the cross-node all-reduce can be hidden behind backward. If the step is short (small per-GPU batch), it can't, and you start needing bigger batches, gradient accumulation or hierarchical reduction.

NCCL and friends

NCCL is NVIDIA's collective library. It detects the topology (NVLink, PCIe, NICs), builds rings and trees, and runs collectives as CUDA kernels on dedicated streams, which is why communication competes with compute for SMs. Equivalents: RCCL (AMD), Gloo (CPU), and XLA collectives over ICI on TPUs. PyTorch exposes them through torch.distributed process groups NCCL GitHub.

Data parallelism (DDP)

Every rank holds a full replica of the model and processes a different slice of the global batch. After backward, gradients are averaged with all-reduce, so every replica applies the same update. The global batch is micro-batch × gradient-accumulation steps × DP degree.

PyTorch DDP makes this efficient in three ways Li+ 2020:

• Gradient bucketing: instead of one all-reduce per parameter (latency-bound) or one giant all-reduce at the end (no overlap), gradients are packed into buckets of about 25 MB by default. A bucket's all-reduce fires as soon as all its gradients are ready.
• Compute/communication overlap: backward produces gradients from the last layer to the first, and buckets are ordered roughly in reverse, so the all-reduce of late layers runs while early layers are still computing gradients. Only the last bucket's communication is exposed.
• no_sync during gradient accumulation: skip the all-reduce on intermediate micro-batches and sync only on the last one.

DDP's limit is memory: every GPU needs the full 16 bytes/param, so DDP alone tops out at roughly a 1–3B model on an 80 GB GPU once activations are included. Its strength is simplicity and the lowest communication volume per step of any approach that does something useful.

ZeRO and FSDP: sharded data parallelism

Plain DP stores N identical copies of the optimizer state, gradients and parameters. ZeRO (Zero Redundancy Optimizer) removes that redundancy in three stages Rajbhandari+ 2019. With \(\Psi\) parameters, DP degree \(N_d\) and \(K=12\):

How ZeRO-3 / FSDP runs a step

Prefetching the next unit's all-gather while the current one computes hides most of the cost when per-GPU compute is large enough. Keeping parameters gathered between forward and backward (reshard_after_forward=False, effectively ZeRO-2 for parameters) removes the second all-gather at the cost of memory. Llama 3 did exactly that: shard optimizer state and gradients, but don't reshard parameters after forward Llama Team 2024.

FSDP1 vs FSDP2, HSDP, and ZeRO variants

• FSDP (v1) flattened each wrapped module's parameters into one FlatParameter and sharded that. It was efficient but awkward with per-parameter features (mixed requires_grad, per-param optimizers) and with torch.compile Zhao+ 2023.
• FSDP2 (fully_shard) shards each parameter individually along dim 0 as a DTensor. It composes cleanly with TP on a 2D device mesh, with torch.compile, and with FP8 all-gathers, and it's the recommended API in current PyTorch and torchtitan PyTorch fully_shard Liang+ 2024.
• HSDP (hybrid sharding): shard within a node (or a group of nodes) and replicate across groups. Parameter all-gathers stay on fast links, and only a gradient all-reduce crosses the slow tier. This is the usual answer to "ZeRO-3 is too chatty across 1000 nodes."
• ZeRO++ cuts ZeRO-3 cross-node traffic with quantized weight all-gathers, a secondary in-node parameter copy, and quantized gradient reduce-scatter Wang+ 2023. ZeRO-Offload/Infinity pushes state to CPU or NVMe Rajbhandari+ 2021.

Tensor parallelism (Megatron-style)

Tensor parallelism splits individual weight matrices across GPUs, so each GPU computes a slice of every layer Shoeybi+ 2019. The trick is to arrange the splits so that communication is needed only once per sub-block.

The MLP

The MLP computes \(Z = \mathrm{GeLU}(XA)\,B\). Split \(A\) by columns, \(A=[A_1, A_2]\). Each GPU computes \(Y_i = \mathrm{GeLU}(XA_i)\) independently, because GeLU is elementwise and the column split keeps whole output features together. No communication is needed here. (A row split of A would need a sum before the nonlinearity, since GeLU(a+b) ≠ GeLU(a)+GeLU(b).) Then split \(B\) by rows, \(B = [B_1; B_2]\), so \(Z = Y_1B_1 + Y_2B_2\). Each GPU computes its partial \(Y_iB_i\), and one all-reduce sums them.

Attention

Attention parallelizes naturally over heads. The Q, K, V projections are column-parallel (each GPU gets \(a/t\) heads), attention runs locally per head with no communication, and the output projection is row-parallel, followed by one all-reduce. With GQA, the KV heads are also divided across GPUs, which caps TP at the number of KV heads unless they're replicated.

Communication cost and why TP stays inside a node

Per transformer layer: 2 all-reduces in forward (one after attention, one after MLP) and 2 in backward, each of size \(b\cdot s\cdot h\) activations. These are synchronous and on the critical path: the next operation needs the result, so overlapping them is hard. (Recent work does overlap them by chunking the GEMM, but it's not free.) Rough ratio: a layer does about \(24\,bsh^2\) forward FLOPs and moves about \(2 \times 2 \times bsh\) bytes per all-reduce pair, so compute per byte scales with \(h\). With only about 50 GB/s per GPU across nodes, the all-reduces would take as long as the math. Over 450+ GB/s NVLink they're manageable. Hence the rule: TP ≤ GPUs per NVLink domain, typically 8 on HGX, potentially larger on NVL72.

TP's benefits: it shards parameters, gradients, optimizer state and activations (within the parallel regions) by \(t\), with no pipeline bubble. Its costs: the communication above, and smaller per-GPU GEMMs that run at lower efficiency as \(t\) grows. Megatron's PTD-P paper found TP within nodes plus PP across nodes the best combination, reaching about 52% of peak on 3072 A100s for a 1T model Narayanan+ 2021.

Sequence parallelism (with TP)

Plain TP leaves some activations replicated: LayerNorm, dropout and residual adds operate on the full \(b\cdot s\cdot h\) tensor on every TP rank. That's the \(10\,sbh\)-ish part of the 34 that TP doesn't divide. Megatron sequence parallelism shards these regions along the sequence dimension instead Korthikanti+ 2022.

Each TP all-reduce becomes a reduce-scatter (into sequence shards) and an all-gather (before the next TP region). Since all-reduce = RS + AG, communication volume is unchanged, and all activations are now divided by \(t\): per-layer memory becomes \(\frac{sbh}{t}\left(34 + \frac{5as}{h}\right)\). In practice SP is always turned on with TP.

Pipeline parallelism

Pipeline parallelism assigns contiguous blocks of layers to different GPUs ("stages"). Stage \(i\) sends its output activations to stage \(i+1\) in forward and receives gradients back in backward. It's point-to-point communication of only \(b\cdot s\cdot h\) per micro-batch per boundary, which is tiny compared with TP's traffic. That's why PP is the parallelism you can stretch across slow inter-node links.

The problem is the bubble. With one batch, only one stage works at a time. GPipe splits the batch into \(m\) micro-batches and pipelines them Huang+ 2018. With \(p\) stages and per-micro-batch forward+backward time \(t\) per stage, the ideal time is \(m\,t\) and the fill/drain overhead is \((p-1)\,t\):

So you want \(m \gg p\), typically \(m \ge 4p\), to keep the bubble under about 20%. But \(m\) is bounded by global batch ÷ (DP × micro-batch size), which ties PP to the batch size. That's one reason very large DP plus deep PP becomes a problem at fixed global batch.

Schedules

Practical PP headaches: stage balancing (the embedding layer and the large-vocabulary output head and loss make the first and last stages heavier, so frameworks put fewer transformer blocks there), the dependence on enough micro-batches, more complex code (schedules, P2P deadlocks), and the interaction with gradient accumulation and FSDP (re-gathering parameters for every micro-batch is wasteful, so PP is usually combined with ZeRO-1 rather than ZeRO-3).

Context parallelism and ring attention

At 128K+ tokens, even one sequence's activations overflow a GPU after TP+SP. Context parallelism (CP) shards the sequence across \(c\) GPUs for the entire layer, attention included. Everything except attention is token-local (MLP, norms, projections), so it needs no communication. Attention needs every query to see every earlier key and value.

Ring attention

Each GPU holds its query chunk \(Q_i\) and its K/V chunk. Over \(c\) steps, K/V chunks rotate around a ring. At each step a GPU computes blockwise attention of \(Q_i\) against the K/V chunk it currently holds and merges the result using FlashAttention's online-softmax rescaling (running max and normalizer). Sending the next K/V chunk overlaps with computing the current one Liu+ 2023. Communication is hidden as long as computing one block (∝ \((s/c)^2 h\)) takes longer than transferring a K/V chunk (∝ \((s/c)\,h_{kv}\)), which gets easier as chunks grow.

Causal load imbalance: with a causal mask, the GPU holding the last chunk attends to everything, while the first attends almost nothing. The standard fix is to split the sequence into \(2c\) pieces and give rank \(i\) pieces \(i\) and \(2c-1-i\) (zig-zag), so every rank does the same work. Llama 3 used this kind of 2×CP chunking, but chose an all-gather of K/V instead of a ring. With GQA, K/V are much smaller than Q, and all-gather makes document masks easier to support Llama Team 2024.

DeepSpeed-Ulysses (all-to-all sequence parallelism)

An alternative: keep activations sequence-sharded outside attention, then use an all-to-all to re-shard from "all heads, s/c tokens" to "a/c heads, all tokens", run ordinary attention per head group, and all-to-all back Jacobs+ 2023. Per-GPU communication stays constant when \(s\) and \(c\) grow together. The degree is limited by the number of (KV) heads. Hybrids (Ulysses inside a node, ring across nodes) are common.

Expert parallelism for MoE

In a Mixture-of-Experts layer, a router sends each token to its top-\(k\) of \(E\) expert MLPs. Expert parallelism places different experts on different GPUs Lepikhin+ 2020. Each MoE layer then needs:

Per-GPU volume per MoE layer is roughly tokens × k × h × bytes, each way. Unlike all-reduce, all-to-all doesn't benefit from ring pipelining. Every pair talks, so it's sensitive to bisection bandwidth and to imbalance. Key issues:

• Load imbalance: a popular expert becomes a straggler. Fixes include auxiliary load-balancing losses, capacity factors that drop overflow tokens (Switch Transformer Fedus+ 2021), and DeepSeek-V3's auxiliary-loss-free bias adjustment.
• Topology-aware routing: DeepSeek-V3 limits each token to at most 4 nodes, sends cross-node traffic over IB once per node, then forwards over NVLink. It used 64-way EP across 8 nodes and overlapped all-to-all with compute using DualPipe DeepSeek-AI 2024.
• Composition: attention layers are dense and typically use DP/TP, while MoE layers use EP. Often the EP group is carved out of the DP group, so the same GPUs act as DP replicas for attention and as EP shards for experts.

Composing 3D/4D/5D parallelism

Real runs combine several dimensions on a device mesh: a logical N-dimensional array of ranks where each axis is one parallelism type. The total GPU count is the product: \(N = d_{dp} \times d_{pp} \times d_{cp} \times d_{tp}\) (× EP folded into DP for MoE layers). The placement rule: the most communication-intensive dimension goes innermost, on the fastest links.

Case study: Llama 3 405B

Meta trained Llama 3 405B with 4D parallelism ordered [TP, CP, PP, DP] from innermost to outermost, with DP implemented as FSDP that shards optimizer state and gradients but doesn't reshard parameters after forward Llama Team 2024:

Things to notice. TP=8 matches the 8-GPU NVLink node. PP=16 spans nodes. Going from 8K to 16K GPUs at a fixed 16M-token batch halved the per-DP-rank batch and cost about 2 points of MFU. For the long-context phase, CP=16 was carved out of DP (DP dropped from 128 to 8) to keep the global batch in tokens constant. Sanity check (using the reported ~15.6T training tokens): \(6 \times 405\text{B} \times 15.6\text{T} \approx 3.8\times10^{25}\) FLOPs; at about 400 TFLOP/s × 16,384 GPUs that's about \(5.8\times10^6\) s, roughly 67 days of perfect running. That's the right order of magnitude for a frontier pretraining run.

Case study: DeepSeek-V3 (671B MoE)

DeepSeek-V3 is a useful contrast: no tensor parallelism at all. It used 16-way PP, 64-way EP across 8 nodes and ZeRO-1 DP on H800s. Heavy engineering on all-to-all kernels and DualPipe made TP unnecessary, and the whole run took 2.788M H800 GPU-hours DeepSeek-AI 2024. The lesson: the "right" mesh depends on the model (dense vs MoE), the interconnect (H800's reduced NVLink) and engineering investment, not on a fixed formula.

How to choose a parallelism strategy

A practical decision procedure, roughly following the Ultra-Scale Playbook's advice HF Playbook 2025:

1. Fit one replica. Can model state plus activations at micro-batch 1 fit on one GPU? If yes, use DDP (or ZeRO-1/2 to free memory for bigger micro-batches). If no, go to step 2.
2. Shard state first. Use FSDP/ZeRO-3, or HSDP across many nodes. It's the least invasive option, with no model code changes.
3. If a single layer or the activations still don't fit, or FSDP communication is exposed at large scale, add TP+SP inside the node (≤8 on HGX).
4. If the model still doesn't fit across one node's TP group, or DP has grown so large that per-rank batch is tiny, add PP across nodes.
5. For long sequences, add CP. For MoE, add EP.
6. Then tune for throughput: micro-batch size, recompute policy, bucket sizes, overlap. Measure MFU and exposed communication with a profiler.

MFU and HFU: measuring efficiency

Model FLOPs Utilization (MFU), popularized by PaLM, is the ratio of the FLOPs the model needs per second at your observed throughput to the hardware's peak Chowdhery+ 2022:

The \(6N\) term covers the parameter matmuls (forward + backward). The second term approximates attention's \(QK^\top\) and \(AV\) FLOPs, which grow with sequence length. Some definitions halve it for causal masking, so state which one you use. Hardware FLOPs Utilization (HFU) also counts FLOPs actually executed, such as recomputation, so HFU ≥ MFU. With full recomputation, HFU is about 8/6 of MFU. MFU is the honest metric, because it can't be inflated by doing redundant work.

Where the missing 50–60% goes: exposed communication (TP all-reduces, FSDP gathers, the last DDP bucket), pipeline bubbles, memory-bound kernels (norms, activations, optimizer step, attention softmax), kernel launch and Python overhead, stragglers and synchronization, data loading, and the fact that marketed peak FLOP/s (and the power and thermal limits behind them) aren't sustainable on real GEMM shapes. Note that MFU computed against FP8 peak looks worse than against BF16 peak for the same wall-clock speedup, so always state the reference precision.

Fault tolerance at scale

At 16K GPUs, failures are routine. During a 54-day snapshot of Llama 3 405B pretraining, Meta saw 466 job interruptions, 419 of them unexpected, with about 78% of the unexpected ones attributed to confirmed hardware issues, faulty GPUs and HBM being the largest categories. Even so, it kept effective training time above 90% Llama Team 2024. That's roughly one interruption every 3 hours. Because synchronous training is all-or-nothing, one bad GPU stops all 16K.

Failure modes

• Hard failures: GPU falls off the bus, HBM ECC errors, NIC/link flaps, host crashes. Usually surfaced as NCCL timeouts or hangs.
• Stragglers: a slow GPU (thermal throttling, a bad link) slows everyone, because every collective waits for the slowest rank.
• Silent data corruption (SDC): wrong arithmetic with no error. Detected via loss spikes, gradient-norm anomalies or cross-replica checksums.
• Numerical instability: loss spikes and divergence. Typical response is to roll back to a checkpoint and skip or reorder data batches.

Checkpointing

The optimal checkpoint interval follows the classic Young/Daly approximation: \(\tau^* \approx \sqrt{2\,C\,M}\), where \(C\) is the time to write a checkpoint and \(M\) is the job's mean time between failures. Example: \(C = 1\) min and \(M = 3\) h gives \(\tau^* \approx \sqrt{2 \cdot 1 \cdot 180} \approx 19\) min. Expected loss per failure is about half an interval plus restart time. Engineering moves:

• Sharded (distributed) checkpoints: each rank writes only its shard in parallel, with metadata that allows resharding on load to a different mesh (PyTorch DCP) PyTorch DCP. A 405B checkpoint with optimizer state is several TB.
• Asynchronous checkpointing: copy state to CPU memory quickly (training blocks only for the device-to-host copy), then persist to storage in the background PyTorch blog.
• In-memory / peer checkpoints: keep recent snapshots in other hosts' RAM for near-instant recovery, with periodic durable checkpoints to object storage.
• Fast restart: hot-spare nodes, health checks before job start, automated bad-node exclusion, and cached NCCL setup. Restart overhead (often minutes) can matter more than checkpoint write time.

Elastic and fault-tolerant training

Elastic training lets the job continue with fewer or more workers instead of waiting for a replacement. torchrun's elastic mode re-forms the process group within a min:max node range torchrun docs. Newer approaches such as torchft tolerate failures at step granularity by treating DP replica groups as independently recoverable: a failed replica drops out of the gradient all-reduce and rejoins later, while the others keep training torchft. The catch is that changing DP size changes the effective batch, and model-parallel groups (TP/PP) still can't lose a member without reconfiguring.

Frameworks

In JAX, you write a single-device program, place arrays on a named mesh (e.g. axes ('data', 'fsdp', 'model')) with PartitionSpecs, and the XLA SPMD partitioner propagates shardings and inserts all-gathers, reduce-scatters and all-reduces. That's the same set of collectives you'd write by hand in Megatron. shard_map gives manual per-device control when the compiler's choices aren't good enough.

Interview question bank