LLM System InterviewHao Hoang · First edition, 2026Home

Chapter 12 Parallelism Strategies

12.1 Data Parallelism and ZeRO

12.1.1 Naive DDP and Its Memory Blow-Up

Data parallelism is the first strategy any interviewer reaches for when probing multi-GPU training, and the expected answer at the mid-level is straightforward: split the batch, sync the gradients, repeat. What separates a strong answer from a mediocre one is the memory accounting. Candidates who cannot explain why naive DDP requires 16 bytes per parameter-and why that number dwarfs the model weights themselves-have not thought seriously about scaling. This topic surfaces in every distributed training conversation, and the memory breakdown is the part that gets people.

Distributed Data Parallel training is a parallelization of minibatch gradient descent along the batch dimension. Given a global batch of size BB and MM GPU workers, each worker receives B/MB / M examples, computes a local gradient, and then synchronizes with every other worker before the optimizer step. The parameter update rule is unchanged:

θt+1=θt−ηB∑i=1B∇iℓ(θt)\theta_{t + 1} = \theta_{t} - \frac{\eta}{B} \sum_{i = 1}^{B} \nabla_{i} \ell ( \theta_{t} )

The synchronization step is an AllReduce over the gradient vectors. Each worker holds a locally computed partial sum gk=1B/M∑i∈shardk∇iℓg_{k} = \frac{1}{B / M} \sum_{i \in \text{shard}_{k}} \nabla_{i} \ell, and AllReduce produces the global average on every worker simultaneously. Ring-AllReduce, the standard implementation, moves a total of 2∣θ∣2 | \theta | bytes across the network per step (each element travels around the ring once in the scatter phase and once in the gather phase).

The memory concern is where the analysis becomes uncomfortable. In BF16 training, the naive accounting goes:

Component Dtype Bytes/param
Model weights BF16 2
Gradients BF16 2
Master weights (FP32) FP32 4
Adam first moment (mm) FP32 4
Adam second moment (vv) FP32 4
Total 16

Every GPU in a naive DDP run carries all 16 bytes per parameter. Adding more GPUs does not reduce this burden; it just means more machines each pinning 16×\times the parameter count in HBM.

In naive DDP, the optimizer states-not the model weights-are the dominant memory cost, and they do not shrink as you add more GPUs.

On the compute side, DDP scales well. Each GPU processes B/MB / M independent examples, and as long as the micro-batch is large enough to saturate the tensor cores, adding more GPUs gives near-linear throughput scaling. The AllReduce cost is 2∣θ∣2 | \theta | bytes per step, but this is paid once per batch regardless of MM in the ring implementation, so larger batches amortize it effectively.

The memory side tells the opposite story. Every worker is a full replica: parameters, gradients, master weights, and both Adam moments are each copied MM times across the cluster. Total cluster memory grows as 16∣θ∣⋅M16 | \theta | \cdot M-linear in MM but with no improvement per device. Contrast this with the ideal: if memory scaled inversely with MM, a cluster twice as large would need half the per-GPU memory. Naive DDP never achieves that.

The practical implication is that batch size becomes the lever that makes DDP viable. Large batches mean the AllReduce overhead is small relative to useful compute, and they improve GPU utilization. But batch size is a resource with diminishing returns: beyond what researchers call the critical batch size, increasing BB reduces gradient noise without meaningfully accelerating convergence, because the training is no longer noise-limited but step-count-limited. Naive DDP therefore hits a wall: it cannot scale memory, and the batch size resource needed to keep it efficient is finite.

A Worked Example

Consider a 7.5 billion parameter model trained on a cluster of 64 A100 80 GB GPUs with the AdamW optimizer in BF16 mixed precision.

Memory per GPU under naive DDP:

7.5×109params×16bytes/param=120GB per GPU7 . 5 \times 10^{9} \text{params} \times 16 \text{bytes/param} = 120 \text{GB per GPU}

An A100 has 80 GB of HBM. This does not fit. Even ignoring activations (which add another ∼\sim10-30 GB depending on batch size and sequence length), the model state alone exceeds the device by 50%.

Communication per step:

The AllReduce moves 2∣θ∣2 | \theta | bytes. With BF16 gradients:

2×7.5×109×2bytes=30GB2 \times 7 . 5 \times 10^{9} \times 2 \text{bytes} = 30 \text{GB}

Over an NVLink-connected 8-GPU node (effective bisection bandwidth ∼\sim600 GB/s aggregate), a single AllReduce on gradients takes roughly 50 ms. Over InfiniBand at ∼\sim200 GB/s, the same exchange takes ∼\sim150 ms-longer than a typical forward-backward pass on a short sequence. Without large batch sizes to hide that latency, DDP over slow interconnects is communication-bound.

Breakdown of the 120 GB:

  • Model weights (BF16): 15 GB

  • Gradients (BF16): 15 GB

  • Master weights (FP32): 30 GB

  • Adam mm (FP32): 30 GB

  • Adam vv (FP32): 30 GB

The optimizer states (master weights + moments) account for 90 GB out of 120 GB-75% of the total. This is the observation that motivates ZeRO: most of the per-GPU memory is redundant replication of optimizer state that only one GPU needs at a time.

What You Actually Decide in Practice

  • Use naive DDP when the model fits. If 16∣θ∣<16 | \theta | < HBM with headroom for activations, naive DDP is the simplest path. For a 1B parameter model on A100 80 GB, the model state is 16 GB-fine. For anything above roughly 4-5B parameters on 80 GB devices, you are already tight before counting activations.

  • AllReduce bandwidth determines your minimum batch size. The rule of thumb: the compute time for one step should be at least 2-3×\times the AllReduce time. If you cannot achieve that with your hardware bandwidth and model size, DDP over slow links will be communication-bound no matter how many GPUs you add.

  • Gradients dominate peak memory mid-backward. The peak memory during a training step occurs partway through the backward pass, when accumulated gradients and surviving activations overlap. Profiling this rather than just the steady-state model size reveals the true memory ceiling.

  • The critical batch size is model- and task-dependent, typically 10510^{5}-10610^{6} tokens. OpenAI’s batch size scaling analysis showed that past the critical batch size, each additional example contributes negligibly to gradient signal. If your global batch already saturates this threshold, adding more data-parallel workers buys compute but wastes gradient diversity.

  • Gradient accumulation is not free. Accumulating gradients over multiple micro-batches to simulate a larger effective batch does not reduce peak gradient memory-each micro-batch’s gradients still accumulate into a full-size gradient tensor. It reduces communication frequency but not per-step memory overhead.

How This Shows Up in Interviews

  • “Walk me through how DDP works and what its limitations are.” A strong answer covers: batch splitting, local forward-backward, AllReduce cost (2∣θ∣2 | \theta |), and the memory accounting that gives 16 bytes per parameter. The limitation is that memory does not scale with GPU count.

  • “You have a 13B parameter model and 8 ×\times A100 80 GB GPUs. Can you train it with naive DDP?” Calculate: 13×109×16=20813 \times 10^{9} \times 16 = 208 GB per GPU. No. Even before activations this is 2.6×2 . 6 \times over budget. The answer should then pivot to what you would do instead (ZeRO stage 3 / FSDP).

  • “Why does Adam make memory worse than SGD?” SGD requires parameters + gradients = 4 bytes/- param (BF16). Adam adds first moment (4 bytes FP32) + second moment (4 bytes FP32) + master weights (4 bytes FP32) = 12 additional bytes. Total becomes 16 vs. 4-a 4×4 \times blow-up attributable entirely to the optimizer.

  • Senior follow-up: “At what point does the AllReduce become the bottleneck, and how would you restructure training to hide it?” The senior answer identifies the compute-to-communication ratio: when one forward-backward pass takes less time than the AllReduce, you are communication-bound. Remedies include gradient compression (lossy, hurts convergence), asynchronous SGD (staleness issues), or switching to a ZeRO stage that interleaves communication with computation layer by layer-which is exactly what ZeRO Stage 2 does.

Key Takeaways

  • Naive DDP uses 16 bytes per parameter: 2 for weights, 2 for gradients, and 12 for Adam optimizer state (FP32 master weights, first moment, second moment).

  • AllReduce moves 2∣θ∣2 | \theta | bytes per step regardless of the number of workers in a ring implementation.

  • Memory does not scale with GPU count in naive DDP-every device carries the full 16 bytes per parameter.

  • Optimizer states (master weights + Adam moments) account for roughly 75% of per-GPU memory in mixed-precision training.

  • Batch size is a finite resource: past the critical batch size (∼105\sim 10^{5}-10610^{6} tokens for most LLMs), adding data-parallel workers does not improve sample efficiency.

12.1.2 ZeRO Stages 1, 2, and 3 (FSDP)

The moment you see that 75% of a DDP node’s memory is optimizer state that every GPU duplicates identically, the fix becomes obvious: shard it. ZeRO-the Zero Redundancy Optimizer-is that fix, structured into three progressive stages that trade implementation complexity for decreasing per-device memory. Interviewers at AI labs ask about ZeRO not because it is obscure, but because understanding it requires you to reason simultaneously about collective communication primitives, memory layout, and backward-pass scheduling. If you can explain all three stages and the identity that makes Stage 1 cost-free, you have demonstrated the kind of systems fluency the job requires.

ZeRO begins with one algebraic observation: AllReduce is equivalent to ReduceScatter followed by AllGather.

In a ring-AllReduce over MM workers, each element of the gradient travels around the ring in a scatter phase (each worker accumulates 1/M1 / M of the total) and then in a gather phase (each worker distributes its accumulated slice to every other). The total data movement is 2∣θ∣2 | \theta | bytes. ReduceScatter alone moves ∣θ∣| \theta | bytes and leaves worker kk holding the fully reduced sum for the kk-th shard. AllGather moves another ∣θ∣| \theta | bytes to propagate those shards everywhere. The bandwidth budget is identical, but the intermediate state after ReduceScatter-where each worker holds only 1/M1 / M of the gradients in reduced form-is the hook ZeRO exploits.

Stage 1 (optimizer state sharding). After the forward and backward passes, every worker has a full gradient vector. A ReduceScatter distributes the summed gradients so that worker kk holds the reduced gradient for parameter shard kk. Each worker then applies the Adam update to its own shard using the optimizer state it owns locally (first moment, second moment, master weights). An AllGather propagates the updated parameters to all workers. The result is a complete synchronized parameter vector on every device, achieved with the same 2∣θ∣2 | \theta | communication as naive DDP-but optimizer state is divided by MM.

Stage 2 (gradient sharding). The problem with Stage 1 is that the full gradient vector must be materialized before ReduceScatter. For large models this is still a memory spike. Stage 2 eliminates it by issuing the reduce operation layer by layer during the backward pass: as soon as the gradient for a layer is computed, it is immediately reduced to the responsible worker and freed on all others. Peak gradient memory becomes ∣θ∣/M| \theta | / M rather than ∣θ∣| \theta |.

Stage 3 (parameter sharding / FSDP). Stages 1 and 2 still replicate the parameters on every GPU. Stage 3 shards the parameters themselves: each worker owns 1/M1 / M of the model weights at rest. During the forward pass, the worker gathers the parameters it needs for each layer (AllGather), computes the layer, then immediately frees the gathered weights. During the backward pass, the same AllGather-compute-free pattern applies, followed by a ReduceScatter on gradients. PyTorch’s FullyShardedDataParallel (FSDP) is a production implementation of Stage 3.

ZeRO Stage 1 is algebraically free relative to naive DDP; Stages 2 and 3 pay modest overhead in exchange for memory that scales inversely with the number of workers.

The key insight is that the optimizer update is embarrassingly parallel across parameter shards. Worker kk does not need workers 0...k−10 . . . k - 1 or k+1...M−1k + 1 . . . M - 1 to be finished with their Adam steps in order to compute its own. All it needs is the fully reduced gradient for its shard (delivered by ReduceScatter) and its own optimizer state. This local update is independent, so the communication required to distribute and collect shards matches exactly what a naive AllReduce would have spent-no extra bandwidth for Stage 1.

Stages 2 and 3 do add overhead. Stage 2 requires synchronizing gradients layer by layer during the backward pass, introducing fine-grained reduce operations that cannot be batched into one large AllReduce. The total bandwidth is still 2∣θ∣2 | \theta | but the latency profile is less regular. Stage 3 is more expensive: the forward pass AllGathers add another ∣θ∣| \theta | bytes, pushing total communication to 3∣θ∣3 | \theta |. In return, per-device memory drops to roughly 16∣θ∣/M16 | \theta | / M for model state, making it possible to fit models that would be impossible otherwise.

Compared to naive DDP, Stage 3 sacrifices 50% more network bandwidth to achieve memory that scales linearly with the number of GPUs-a compelling trade whenever memory is the binding constraint.

A Worked Example

Consider a 7.5 billion parameter model trained with AdamW in BF16 mixed precision, distributed across M=64M = 64 workers.

Naive DDP baseline:

7.5×109×16bytes=120GB per GPU7 . 5 \times 10^{9} \times 16 \text{bytes} = 120 \text{GB per GPU}

ZeRO Stage 1 shards the 12 optimizer-state bytes (FP32 master weights + mm + vv):

2+2⏟params + grads (BF16)+4+4+464 bytes/param≈4.2 bytes/param⇒≈31 GB per GPU\underbrace{2+2}_{\text{params + grads (BF16)}}+\frac{4+4+4}{64}\text{ bytes/param}\approx4.2\text{ bytes/param}\Rightarrow\approx31\text{ GB per GPU}

Communication cost: 2∣θ∣2 | \theta | (identical to DDP).

ZeRO Stage 2 additionally shards gradients:

2⏟params (BF16)+2+4+4+464 bytes/param≈2.2 bytes/param⇒≈16.6 GB per GPU\underbrace{2}_{\text{params (BF16)}}+\frac{2+4+4+4}{64}\text{ bytes/param}\approx2.2\text{ bytes/param}\Rightarrow\approx16.6\text{ GB per GPU}

Communication cost: 2∣θ∣2 | \theta |, but paid in per-layer increments.

ZeRO Stage 3 (FSDP) shards everything:

2+2+4+4+464bytes/param=1664=0.25bytes/param⇒≈1.9GB per GPU\frac{2 + 2 + 4 + 4 + 4}{64} \text{bytes/param} = \frac{16}{64} = 0 . 25 \text{bytes/param} \Rightarrow \approx 1 . 9 \text{GB per GPU}

Communication cost: 3∣θ∣3 | \theta | per step (AllGather in forward, AllGather in backward, ReduceScatter in backward).

Maximum trainable model size on 8 ×\times A100 80 GB:

  • Naive DDP: ⌊(80\lfloor ( 80 GB −- activations)/16⌋≈4) / 16 \rfloor \approx 4-66 billion parameters.

  • FSDP (Stage 3): ⌊8×80\lfloor 8 \times 80 GB/16⌋≈40/ 16 \rfloor \approx 40-5050 billion parameters. FSDP enables roughly 8×8 \times larger models on the same hardware, at the cost of 3×3 \times versus 2×2 \times network utilization.

Overlapping communication and computation in FSDP:

FSDP achieves acceptable efficiency by prefetching. While GPU kk is computing layer ll’s forward pass, the AllGather for layer l+1l + 1’s weights is already in flight. By the time the compute for layer ll finishes, the weights for l+1l + 1 are available without stalling. The same prefetch applies to the backward pass. The communication bubbles that remain are small relative to compute time, provided the AllGather latency is less than the per-layer compute time-a condition that holds on NVLink-connected intra-node workers but can fail on slow inter-node links.

What You Actually Decide in Practice

  • Always use Stage 1. It costs nothing in bandwidth and cuts optimizer state memory by MM. There is no reason to use bare DDP over Stage 1 except tooling inertia.

  • Prefer Stage 2 over Stage 3 when the model fits. Stage 2 keeps parameters on every device, which means each layer’s forward pass requires no AllGather. The per-layer gradient reductions add overhead, but it is a small constant penalty. Stage 3 adds ∣θ∣| \theta | bytes of network traffic per step and introduces AllGather bubbles in every layer. Use Stage 3 only when the model state does not fit even with Stage 2.

  • Stage 3 (FSDP) requires fast intra-node interconnects. The AllGather per layer is latency-sensitive. On NVLink (600 GB/s aggregate across 8 GPUs on an H100 node), the overhead is acceptable. Across InfiniBand (200 GB/s), prefetching can fail to fully hide the gather latency, and GPU utilization drops. The practical rule: use FSDP within a node, and consider pipeline or tensor parallelism for cross-node sharding of very large models.

  • Activation memory is not reduced by any ZeRO stage. ZeRO 1, 2, and 3 operate on model state. The activation memory footprint-which scales as O(s⋅b⋅d⋅L)O ( s \cdot b \cdot d \cdot L ) for sequence length ss, batch size bb, hidden dimension dd, and layers LL-remains fully replicated on each data-parallel worker. For long-context training or very large hidden dimensions, activation recomputation (gradient checkpointing) or sequence parallelism is necessary on top of ZeRO.

  • Gradient accumulation interacts with Stage 2. With Stage 2, per-layer gradient reductions happen during the backward pass. If you accumulate gradients over multiple micro-batches before applying the optimizer, you must ensure that the per-layer reduces are deferred to the end of the accumulation windowotherwise you reduce partial gradients prematurely. PyTorch FSDP handles this with a no_sync() context manager that suppresses gradient synchronization during accumulation steps.

How This Shows Up in Interviews

  • “Explain ZeRO Stage 1 and why it doesn’t cost more bandwidth than naive DDP.” The answer turns on the AllReduce decomposition: ReduceScatter (∣θ∣| \theta | bytes) delivers per-shard reduced gradients; each worker runs its Adam update locally; AllGather (∣θ∣| \theta | bytes) redistributes updated parameters. Total: 2∣θ∣2 | \theta |- same as a ring-AllReduce. The insight is that the optimizer step fits cleanly between the two halves of the AllReduce without adding communication.

  • “What is FSDP and when would you use it instead of DDP?” FSDP is ZeRO Stage 3: parameters, gradients, and optimizer state are all sharded. Use it when the model state (16∣θ∣16 | \theta | bytes) exceeds available HBM even after stages 1 and 2. The cost is 3×3 \times network utilization versus 2×2 \times for DDP and potential compute bubbles if AllGather latency is not hidden by prefetching.

  • “Does ZeRO reduce activation memory?” No. ZeRO only shards model state (parameters, gradients, optimizer state). Activations are local to each data-parallel worker and scale with batch size and sequence length. The interviewer is probing whether you conflate model state with activation memory-a common mistake.

  • Senior follow-up: “At 1024 GPUs, would you use FSDP alone or combine it with tensor parallelism? Why?” The senior answer is: combine them. FSDP over 1024 GPUs means every AllGather for a layer’s parameters crosses the full cluster, which at inter-node InfiniBand speeds is brutally slow. Instead, apply tensor parallelism within each 8-GPU node (using NVLink) and FSDP or pipeline parallelism across nodes (using InfiniBand). This matches the communication pattern to the available bandwidth at each level of the hierarchy-a principle that underpins every practical large-scale training configuration.

Key Takeaways

  • AllReduce == ReduceScatter ++ AllGather; this identity is why ZeRO Stage 1 adds zero communication overhead relative to naive DDP.

  • Stage 1 shards optimizer state (÷M\div M); Stage 2 additionally shards gradients; Stage 3 (FSDP) shards everything including parameters, reducing per-device model state to 16∣θ∣/M16 | \theta | / M.

  • FSDP increases total communication from 2∣θ∣2 | \theta | to 3∣θ∣3 | \theta | per step; it masks this overhead by prefetching parameter AllGathers while the previous layer computes.

  • ZeRO does not reduce activation memory; gradient checkpointing and sequence parallelism handle that separately.

  • The practical rule: use Stage 1 unconditionally, Stage 2 when gradients pressure memory, Stage 3 only when the model itself does not fit-and pair Stage 3 with tensor parallelism on fast intra-node links rather than running it flat across the whole cluster.

12.2 Model Parallelism

12.2.1 Tensor Parallelism (Megatron-Style)

When a model’s weight matrices individually exceed a single GPU’s HBM-and ZeRO-style sharding is not enough because gathering those weights per layer costs too much bandwidth-you need to split the matrix multiplications themselves across GPUs. Tensor parallelism does exactly this. It shows up at senior-level interviews the moment the conversation turns to training 70B or larger models on a single node, and the question that separates candidates is whether they can name the two synchronization points per transformer layer, explain why the column-parallel/row-parallel split requires no communication at the first projection, and give the bandwidth rule that caps tensor parallel degree at eight GPUs.

The foundation is that matrix-multiply outputs are sums of partial products, and those partial products can live on different devices. Consider the dense MLP block: Y=Y = Dropout(GeLU(XA)B)( \text{GeLU} ( X A ) B ), where X∈Rs×hX \in \mathbb{R}^{s \times h}, A∈Rh×4hA \in \mathbb{R}^{h \times 4 h}, and B∈R4h×hB \in \mathbb{R}^{4 h \times h}.

Column-parallel split for AA. Partition AA along its column dimension into tt equal shards: A=[A1∣⋅⋅⋅∣At]A = [ A_{1} | \cdot \cdot \cdot | A_{t} ]. GPU kk holds Ak∈Rh×4h/tA_{k} \in \mathbb{R}^{h \times 4 h / t} and computes Yk=GeLU(XAk)Y_{k} = \text{GeLU} ( X A_{k} ). The input XX is replicated on every GPU-either via a broadcast at the start of the forward pass, or simply because the preceding layer’s AllReduce already produced a full copy everywhere. No communication is needed to produce YkY_{k}.

Row-parallel split for BB. Partition BB along its row dimension: B=[B1⊤;⋅⋅⋅;Bt⊤]⊤B = [ B_{1}^{\top} ; \cdot \cdot \cdot ; B_{t}^{\top} ]^{\top}, with Bk∈R4h/t×hB_{k} \in \mathbb{R}^{4 h / t \times h}. GPU kk computes Zk=YkBkZ_{k} = Y_{k} B_{k}, yielding a partial sum. The full output Z=∑kZkZ = \sum_{k} Z_{k} is obtained with a single AllReduce across all tt GPUs. This is synchronization point ff in the forward pass.

The attention block follows the same template: split the QQ, KK, VV projection matrices column-parallel so GPU kk handles h/th / t attention heads independently; split the output projection row-parallel; AllReduce to merge heads.

Per transformer layer this yields one AllReduce in the forward pass (after each output projection) and one AllReduce in the backward pass. Both have the same tensor size: s×b×hs \times b \times h elements, where ss is sequence length and bb is batch size.

Tensor parallelism splits weight matrices width-wise and communicates only activations at layer boundariesparameters never move.

Unlike data parallelism, tensor parallelism does not consume batch size. Every tensor-parallel GPU processes the full batch; parallelism comes from reducing the size of each GPU’s weight matrices, not from distributing data samples. This makes it the first tool to reach for when memory is the binding constraint but the global batch size is already at or near its optimal value.

Unlike pipeline parallelism, tensor parallelism introduces no pipeline bubble. All tt GPUs execute every layer in lockstep, just on different parameter shards. There is no stage sequencing, no micro-batch scheduling, and no idle time waiting for upstream activations. The forward and backward passes look like single-GPU execution to the optimizer.

The cost is raw bandwidth. Each AllReduce sends 2⋅(1−1/t)⋅s⋅b⋅h2 \cdot ( 1 - 1 / t ) \cdot s \cdot b \cdot h bytes per device. For a 70B model with h=8192h = 8192, batch size 32, and sequence length 2048:

s×b×h×2bytes (BF16)=2048×32×8192×2≈1.07GBs \times b \times h \times 2 \text{bytes (BF16)} = 2048 \times 32 \times 8192 \times 2 \approx 1 . 07 \text{GB}

Effective transfer per device in an 8-way AllReduce: 2×1.07×(7/8)≈1.872 \times 1 . 07 \times ( 7 / 8 ) \approx 1 . 87 GB.

At NVLink effective AllReduce bandwidth of ∼\sim600 GB/s: 1.87/600≈3.11 . 87 / 600 \approx 3 . 1 ms per AllReduce. A single MLP layer’s forward GEMMs on one GPU (processing 1/81 / 8 of the FFN width): roughly 2×32×2048×8192×(28672/8)≈30.72 \times 32 \times 2048 \times 8192 \times ( 28672 / 8 ) \approx 30 . 7 TFlop, which at H100’s 989 TF/s (BF16) takes ∼\sim31 ms. Communication (3.1 ms) is 10% of compute-consistent with the observed 10-12% throughput penalty at intra-node tensor parallel degree 8.

Cross the node boundary to InfiniBand (25 GB/s): the same 1.87 GB AllReduce takes ∼\sim75 ms, which is twice the compute time. That is why tensor parallelism must stay intra-node, and why empirically going from degree 8 to degree 16 causes a ∼\sim42% throughput drop, and degree 32 drops ∼\sim65%.

A Worked Example

A 70B transformer: h=8192h = 8192, FFN intermediate 28672, 80 layers, 64 attention heads. Tensor parallel degree t=8t = 8 on one 8×\timesH100 NVLink node.

Parameter memory per device. Total model state in BF16: 70×109×270 \times 10^{9} \times 2 bytes =140= 140 GB. With t=8t = 8:

140GB8=17.5GB per GPU\frac{140 \text{GB}}{8} = 17 . 5 \text{GB per GPU}

The remaining 62.5 GB of the 80 GB HBM budget is available for optimizer state (under ZeRO Stage 1), activations, and any KV cache.

Head count check. 64 attention heads, t=8t = 8: each GPU handles exactly 64/8=864 / 8 = 8 heads. Valid. If a model had 40 heads and you tried t=8t = 8, the split would be non-integer-a design error that forces rounding and unbalanced load.

Communication per training step. With batch size 32, sequence 2048:

  • AllReduce tensor size per layer: 32×2048×8192×232 \times 2048 \times 8192 \times 2 bytes =1.07= 1 . 07 GB

  • AllReduces per layer: 2 in forward (MLP out, attention out), 2 in backward-4 total

  • 80 layers: 80×4×1.0780 \times 4 \times 1 . 07 GB =342= 342 GB per forward+backward step

  • At 600 GB/s NVLink effective bandwidth: 342/600≈0.57342 / 600 \approx 0 . 57 s of communication This sounds alarming until you recognize the communication is overlapped with compute: each AllReduce fires while the CPU is dispatching the next layer’s kernel. The practical throughput penalty is 10-12%, not 37%.

What You Actually Decide in Practice

  • Cap tensor parallel degree at 8 (or the NVLink-connected GPU count for your node). Beyond that, you are issuing AllReduces over InfiniBand, and throughput collapses. The 10-12% intra-node overhead is worth the memory savings; the 40-65% inter-node overhead is not. Even if a single node has 16 GPUs, benchmark at t=8t = 8 versus t=16t = 16 before committing-some NVSwitch fabrics degrade at full all-to-all traffic.

  • Attention heads and FFN width must be divisible by tt. When designing a model for t=8t = 8, set head count to a multiple of 8 and FFN intermediate dimension to a multiple of 512. Misaligned shapes produce suboptimal GEMM tiling and can add 5-15% overhead on top of the communication cost.

  • Tensor parallelism does not consume batch size. This distinguishes it from data parallelism and pipeline parallelism, both of which require large batches to hide their respective overheads. If your global batch size is already at the critical batch size threshold and you cannot increase it, tensor parallelism is the correct memory-scaling tool.

  • Pair with sequence parallelism to eliminate the residual activation memory. Tensor parallelism leaves LayerNorm and Dropout unsharded-they operate on the full s×b×hs \times b \times h activation. Sequence parallelism distributes these operations along the sequence dimension and replaces the MLP/attention AllReduce with AllGather and ReduceScatter pairs of equal total cost. The combination reduces the non-matmul activation memory by a factor of tt.

  • Do not extend tensor parallelism cross-node to compensate for a model that is too wide. If TP-8 within a node cannot fit the model, the answer is to combine TP-8 intra-node with pipeline parallelism inter-node-not to stretch tensor parallelism across InfiniBand. The activation AllReduce pattern of tensor parallelism is uniquely hostile to high-latency links.

How This Shows Up in Interviews

  • “How would you fit a 70B model on a single 8×\timesH100 node?” Use tensor parallelism with t=8t = 8: each GPU holds 140140 GB/8=17.5/ 8 = 17 . 5 GB of parameters. Split each weight matrix column-parallel (first projection) and row-parallel (second projection), AllReduce the output activations. One AllReduce per transformer layer per direction. Model fits; communication overhead on NVLink is around 10%.

  • “Why is tensor parallel degree capped at 8 in most production systems?” Beyond 8 GPUs on a standard NVIDIA node, AllReduces cross InfiniBand. For a 70B model with batch 32, sequence 2048, each AllReduce is 1.87 GB; at 25 GB/s InfiniBand that is 75 ms versus 31 ms of compute-the GPU stalls more than it runs. Empirically: 10-12% overhead at degree 8 (intra-node NVLink), 42% at degree 16, 65% at degree 32.

  • “What determines the AllReduce size in tensor parallelism?” The output activation of each projection: s×b×hs \times b \times h elements per layer, not the weight matrix size. This is why tensor parallelism traffic scales with sequence length and batch size, not with model depth.

  • Senior follow-up: “If you had to scale to 1024 GPUs for a single model, how would you combine tensor parallelism with other strategies?” Apply TP-8 within each 8-GPU NVLink node, consuming the fast intra-node bandwidth for the activation AllReduces. Then apply pipeline parallelism across nodes to extend model depth coverage, with InfiniBand carrying only the small activation handoffs between pipeline stages. Data parallelism handles the remaining parallelism once the model fits. Never extend TP beyond the NVLink boundary. The senior insight is that each interconnect tier dictates which parallelism strategy uses it.

Key Takeaways

  • Tensor parallelism splits weight matrices width-wise: column-parallel for the first projection (input is replicated; no upfront communication), row-parallel for the second projection (partial sums merged by one AllReduce).

  • Cost is two AllReduces per transformer layer (forward and backward), each of size s×b×hs \times b \times h activationsweights never move.

  • The hard rule: TP degree ≤8\leq 8 on NVLink; going inter-node over InfiniBand causes 40-65% throughput collapse.

  • TP does not consume batch size and adds no pipeline bubble, making it the primary memory-scaling tool when fast intra-node interconnects are available.

  • Head count and FFN width must be divisible by tt; design models with multiples of 8 or 16 to keep GEMM tile shapes optimal.

12.2.2 Pipeline Parallelism and Bubble Reduction

Pipeline parallelism becomes necessary when a model is too large to fit even after slicing all weight matrices across the GPUs within a node using tensor parallelism, or when you must extend to a second node and cannot afford tensor parallelism’s AllReduce traffic over InfiniBand. Rather than splitting matrices width-wise, pipeline parallelism cuts the network depth-wise: each device owns a contiguous group of layers and passes activations forward to the next stage in sequence. Every serious interview on large-scale training eventually asks about the pipeline bubble, and the answer that separates mid-level from senior is knowing the formula, knowing how to suppress the bubble with micro-batches, and knowing why zero-bubble scheduling is hard enough that two people at a frontier lab may be the only ones who understand their own implementation.

Partition LL transformer layers into pp equal pipeline stages, each assigned to one device. Device 0 runs layers 0...L/p−10 . . . L / p - 1; device 1 runs layers L/p...2L/p−1L / p . . . 2 L / p - 1; and so on. The forward pass is strictly sequential: device 0 finishes its layers for one micro-batch, sends the output activations to device 1 via a point-to-point send, then immediately begins the next micro-batch. The backward pass flows in reverse-gradients of activations travel from device p−1p - 1 back to device 0.

The pipeline bubble is the idle time that arises during startup and teardown. In the startup phase, devices 1...p−11 . . . p - 1 are idle while device 0 processes the first micro-batch. In the teardown phase, device 0 is idle while device p−1p - 1 finishes. The ratio of wasted time to useful time is:

bubble fraction=p−1m\text{bubble fraction} = \frac{p - 1}{m}

where mm is the number of micro-batches that partition the global batch. With p=8p = 8 stages and m=1m = 1, the bubble fraction is 7-seven eighths of the wall-clock step is wasted. With p=8p = 8 and m=32m = 32, it drops to 7/32≈22%7 / 32 \approx 22 \%-still significant but workable.

The standard production schedule is 1F1B (one-forward-one-backward): each device interleaves one forward micro-batch with one backward micro-batch. This keeps the number of live activation tensors bounded at pp regardless of mm, preventing activation memory from growing linearly with the micro-batch count.

Pipeline parallelism tiles the model across depth; the bubble fraction (p−1)/m( p - 1 ) / m is the price, and microbatch count is the only handle that reduces it without adding complexity.

The decisive advantage over tensor parallelism is communication volume. Tensor parallelism AllReduces the full activation tensor at every layer boundary-proportional to s×b×hs \times b \times h per layer per direction. Pipeline parallelism transmits activations only at stage transitions: a single point-to-point send of size sμ×bμ×hs_{\mu} \times b_{\mu} \times h per micro-batch, where bμb_{\mu} is the micro-batch size. For a large model with small micro-batches this is one to two orders of magnitude less inter-device traffic per step, and critically, this communication is a one-time handoff-not repeated per layer.

This makes pipeline parallelism viable over InfiniBand where tensor parallelism fails. A 201 MB activation handoff at 25 GB/s takes about 8 ms. A single pipeline stage’s compute for a 175B model takes hundreds of milliseconds. The communication cost is well under 5% of compute time, compared to tensor parallelism’s 40-65% overhead at the same link speed.

Pipeline parallelism also shards activation memory in proportion to the number of stages. Each device holds activations only for its L/pL / p layers rather than the full model. This is something ZeRO stages 1-3 cannot providethey shard model state but leave activations fully on each data-parallel worker. For long-sequence training where activations dominate, pipeline parallelism combined with micro-batch 1F1B scheduling is the primary lever.

A Worked Example

Consider a 175B transformer (h=12288h = 12288, 96 layers, FFN intermediate =49152= 49152) distributed across 64 A100 80 GB GPUs using p=8p = 8 pipeline stages and 8-way data parallelism.

Layer assignment and memory: Each stage handles 96/8=1296 / 8 = 12 layers. Parameter memory per device in BF16:

175×109×2bytes8=43.75GB\frac{175 \times 10^{9} \times 2 \text{bytes}}{8} = 43 . 75 \text{GB}

leaving roughly 36 GB per device for optimizer state, activations, and buffers.

Bubble suppression: Global batch 1024. Each data-parallel replica processes 1024/8=1281024 / 8 = 128 samples. With micro-batch size 4: m=128/4=32m = 128 / 4 = 32 micro-batches. Bubble fraction:

8−132=732≈21.9%\frac{8 - 1}{32} = \frac{7}{32} \approx 21 . 9 \%

GPU utilization: ∼\sim78%. Compare to m=4m = 4: bubble =7/4=175%= 7 / 4 = 175 \% of useful compute-completely unusable without larger batches.

Communication per stage transition: Each activation handoff is the output of one micro-batch after one stage:

bμ×s×h×2bytes=4×2048×12288×2≈201MBb_{\mu} \times s \times h \times 2 \text{bytes} = 4 \times 2048 \times 12288 \times 2 \approx 201 \text{MB}

At InfiniBand 25 GB/s: 201201 MB/25/ 25 GB/s ≈8\approx 8 ms.

Compute per stage: A single layer’s MLP forward pass, per GPU:

2×4×2048×12288×49152≈9.88TFlop2 \times 4 \times 2048 \times 12288 \times 49152 \approx 9 . 88 \text{TFlop}

At 312 TF/s (A100 BF16): ≈32\approx 32 ms per MLP layer. With 12 layers per stage: ∼\sim384 ms of compute per microbatch step. The 8 ms handoff is ∼\sim2% of stage compute-negligible. This is precisely why pipeline parallelism survives on InfiniBand.

What You Actually Decide in Practice

  • Target micro-batch count m≥4(p−1)m \geq 4 ( p - 1 ) to keep bubble below 25%. With p=8p = 8, this means m≥28m \geq 28. If your global batch cannot support this micro-batch count, pipeline parallelism is imposing a large and unrecoverable compute penalty. Either reduce pipeline depth or switch to FSDP if the model fits.

  • Zero-bubble scheduling (DualPipe) fills bubble slots with weight-gradient work. The standard backward pass has two independent phases: computing activation gradients (BB, which is a serial dependencyearlier layers need it before they can proceed) and computing weight gradients (WW, which has no serial dependency and can run any time after the corresponding forward pass). By deferring WW computations into the idle bubble slots, you approach 100% GPU utilization. DeepSeek-V3 uses a variant of this called DualPipe in production. The catch: interposing on autograd to separate the BB and WW phases requires custom gradient hooks, a scheduling queue for WW operations, and careful management of intermediate activations that must be kept alive between the BB and WW steps. The anecdote I keep hearing from engineers at frontier labs is that two people in the training infrastructure group understand their pipeline parallel code-and when one of them leaves, the codebase becomes a liability.

  • Align pipeline stage boundaries with physical node boundaries. An activation handoff that crosses a single InfiniBand hop costs 8 ms for a 200 MB tensor. One that crosses two hops (leaf switch to spine switch) can cost 30+ ms. Design stages so that each stage lives entirely within one node, and the handoffs stay on the lowest-latency inter-node link available.

  • Batch size is a shared resource between data parallelism and pipeline parallelism. If you run 8-way data parallel with p=4p = 4 pipeline stages and need m=12m = 12 micro-batches per replica, the effective per-replica batch is bμ×12b_{\mu} \times 12. Verify that your global batch size supports this without pushing below the critical batch size threshold-beyond which gradient noise no longer benefits from larger batches.

  • Prefer FSDP over pipeline parallelism whenever the model fits. FSDP is far simpler to implement, has no bubble, and is architecture-agnostic. Pipeline parallelism is the right answer only when (a) the model cannot fit per-node even with FSDP and TP-8, or (b) inter-node bandwidth is too low for FSDP’s 3∣θ∣3 | \theta | per-step communication, or (c) you have a very large global batch that can suppress the bubble to an acceptable level.

How This Shows Up in Interviews

  • “Explain the pipeline bubble and how you would reduce it.” The bubble arises because during startup, later stages idle while device 0 fills them; during teardown, earlier stages idle while the last stage drains. Fraction =(p−1)/m= ( p - 1 ) / m. To reduce: increase mm (more micro-batches), or use 1F1B scheduling to interleave forward and backward passes while bounding activation memory at pp live tensors regardless of mm. Cite a concrete example: p=8p = 8, m=32m = 32, bubble =7/32≈22%= 7 / 32 \approx 22 \%, utilization ≈78%\approx 78 \%.

  • “When would you choose pipeline parallelism over FSDP?” When the model must scale across nodes and the inter-node bandwidth is too low for FSDP’s AllGather-per-layer pattern (which costs 3∣θ∣3 | \theta | per step), or when the model is large enough that even FSDP cannot fit it per node. Pipeline parallelism’s activation handoffs are one to two orders of magnitude smaller than FSDP’s weight AllGathers at the same model size, making it viable over 25 GB/s InfiniBand links where FSDP stalls.

  • Senior follow-up: “Describe zero-bubble pipelining and what it costs to implement.” The key insight is that the backward pass has two independent phases: BB (activation gradients, required by earlier stages) and WW (weight gradients, locally executable any time after the corresponding forward). The WW computations can be deferred to slots where the pipeline would otherwise be idle. To implement this, you must intercept autograd at the layer level to separate BB and WW, maintain a deferred queue for WW work, and track which intermediate activations need to stay alive across the separation. The memory overhead for deferred activations can partially eat back the efficiency gains, and the scheduling logic is complex enough that teams at production-scale labs have had serious bus-factor problems with it. A strong senior answer names the autograd intervention, the memory tracking complexity, and gives a real example (DualPipe in DeepSeek-V3).

Key Takeaways

  • Pipeline parallelism partitions the model by layers: device kk owns L/pL / p consecutive layers and communicates only output activations-not weights-to the next stage via point-to-point sends.

  • Bubble fraction =(p−1)/m= ( p - 1 ) / m; with p=8p = 8 and m=32m = 32, GPU utilization is ∼\sim78%. Target m≥4(p−1)m \geq 4 ( p - 1 ) to keep the bubble under 25%.

  • Activation handoffs are small (one tensor per micro-batch per stage transition), making pipeline parallelism viable over InfiniBand where tensor parallelism’s per-layer AllReduces are not.

  • Zero-bubble scheduling splits the backward pass into BB (serial activation gradients) and WW (local weight gradients), filling bubble slots with WW work-at the cost of significant autograd infrastructure complexity.

  • The production rule: TP-8 intra-node, pipeline parallel inter-node, data parallel for the remaining GPU count; set m≥4(p−1)m \geq 4 ( p - 1 ); use FSDP instead of pipeline parallel whenever the model fits per node.

12.3 Activation and Sequence Parallelism

12.3.1 Sequence Parallelism for LayerNorm and Dropout

After you’ve walked an interviewer through Megatron-style tensor parallelism-column-parallel first projection, row-parallel second projection, one AllReduce per layer boundary-the senior follow-up is almost always: “Those AllReduces shard the matmul outputs across GPUs. But what about LayerNorm and Dropout? Those operators see the full activation, right?” Most candidates who prepped TP can’t answer this without hesitation. The unsharded pointwise operations are the hole in the tensor-parallel memory story, and at sequence lengths above 4096 they can consume more HBM than the weight matrices you just worked so hard to shard. Sequence parallelism, introduced alongside the Megatron-LM tensor-parallel formulation, closes the gap at exactly zero additional bandwidth cost.

LayerNorm and Dropout are embarrassingly parallel along the sequence dimension. Each position normalizes independently of every other position; Dropout applies a per-element mask with no cross-position dependency. This means we can partition the ss-length sequence across the same tt tensor-parallel GPUs and run these pointwise ops on an s/ts / t-token slice, exactly as if we had cut the batch into micro-batches but along the sequence axis instead.

The wrinkle is that the column-parallel matmul that follows each LayerNorm needs the full-sequence activation on every GPU. The synchronization point must therefore change. Standard tensor parallelism issues a single AllReduce after the row-parallel matmul to merge partial sums and replicate the result everywhere. Sequence parallelism replaces this with a two-part operation:

  1. ReduceScatter after the row-parallel matmul: partial sums are reduced across the group, and the result is scattered so that GPU kk holds the fully-reduced rows for sequence positions [k⋅s/t,(k+1)⋅s/t)[ k \cdot s / t , ( k + 1 ) \cdot s / t ). No other GPU holds those rows.

  2. LayerNorm, Dropout, and any other pointwise op run on the local s/ts / t-token slice.

  3. AllGather before the next column-parallel matmul: each GPU’s slice is concatenated so that every GPU sees the full ss-length sequence for the subsequent projection.

By the identity AllReduce ≡\equiv ReduceScatter ++ AllGather in the bandwidth-limited regime, the total bytes transferred are unchanged. The ReduceScatter ships the same data as the first half of an AllReduce ring; the AllGather ships the same data as the second half. Sequence parallelism is literally a rearrangement of when the data moves, not an addition of new communication. The only change visible in the memory profile is that the activation tensor between the ReduceScatter and the AllGather has shape (s/t)×b×h( s / t ) \times b \times h instead of s×b×hs \times b \times h-a factor-of-tt reduction for every byte spent on LayerNorm input, LayerNorm output, Dropout, and the residual add.

The net activation memory per transformer layer, after applying both tensor parallelism and sequence parallelism and using FlashAttention to eliminate the quadratic softmax term, is:

Activation memory per layer=sbh×34tbytes (BF16)\text{Activation memory per layer} = \frac{s b h \times 34}{t} \text{bytes (BF16)}

Without sequence parallelism the formula was sbh×34/t+sbh×10s b h \times 34 / t + s b h \times 10: the second term, representing LayerNorm, Dropout, and input embeddings, stayed constant regardless of tt. Sequence parallelism removes that stranded term.

Sequence parallelism converts the one part of the activation that tensor parallelism could not touch into a term that scales with tt-at zero bandwidth cost.

The AllReduce in vanilla tensor parallelism does the work of reduce-and-scatter and then immediately throws away the scattered structure by copying the result to all GPUs. Sequence parallelism simply stops discarding that structure. The ReduceScatter leaves each GPU holding exactly the rows it needs to run LayerNorm on its sequence slice; the AllGather restores the full tensor before the next matmul requires it. The data travels the same physical distance over the same wires; the only difference is how the receive buffer is allocated.

The alternative, running LayerNorm on the full sequence on every GPU and then AllReducing, is what you get by default with TP but without SP. It is not incorrect-the gradients compute cleanly-but it wastes sbh×10s b h \times 10 bytes per layer per GPU, regardless of how aggressively you parallelized the weight matrices. For a 13B model at sequence 4096 and t=8t = 8, that stranded term is over 60 GB of HBM. Sequence parallelism costs nothing to recover it.

A second advantage is composability. Because TP and SP operate on the same communicator group with the same degree, adding SP requires no new topology decisions, no new process groups, and no change to the pipeline or data-parallel configuration. In Megatron-LM, enabling SP is a single flag. Frameworks that expose TP without SP are leaving memory on the table.

A Worked Example

A 13B-parameter transformer: h=5120h = 5120, 40 layers, 80 attention heads, FFN width 20480. Training configuration: batch b=4b = 4, sequence s=4096s = 4096, tensor and sequence parallel degree t=8t = 8, BF16 weights.

Without sequence parallelism. After tensor parallelism, the matmul activations (FFN intermediates, attention projection outputs) divide by tt. The pointwise activation per layer does not:

sbh×10×2bytes=4096×4×5120×10×2=1,678MB≈1.56GB per layers b h \times 10 \times 2 \text{bytes} = 4096 \times 4 \times 5120 \times 10 \times 2 = 1 , 678 \text{MB} \approx 1 . 56 \text{GB per layer}

Over 40 layers: 40×1.56=62.440 \times 1 . 56 = 62 . 4 GB. This occupies 78% of an H100’s 80 GB HBM before a single weight matrix or optimizer state is counted.

With sequence parallelism (t=8t = 8):

sbh×10×2bytest=1.56GB8≈195MB per layer\frac{s b h \times 10 \times 2 \text{bytes}}{t} = \frac{1 . 56 \text{GB}}{8} \approx 195 \text{MB per layer}

Over 40 layers: ≈7.8\approx 7 . 8 GB. Savings: 62.4−7.8=54.662 . 4 - 7 . 8 = 54 . 6 GB per GPU-nearly a full H100 freed.

Communication cost check. The AllReduce in vanilla TP transfers per device:

2×sbh×(1−1t)×2bytes=2×4096×4×5120×78×2≈1,468MB2 \times s b h \times \bigl( 1 - \frac{1}{t} \bigr) \times 2 \text{bytes} = 2 \times 4096 \times 4 \times 5120 \times \frac{7}{8} \times 2 \approx 1 , 468 \text{MB}

The ReduceScatter + AllGather pair transfers the same total bytes per device by the AllReduce identity. Communication overhead: zero.

Minimum total activation per layer after TP + SP + FlashAttention:

sbh×34t×2bytes=4096×4×5120×348×2≈2,863MB≈2.80GB per layer\frac{s b h \times 34}{t} \times 2 \text{bytes} = \frac{4096 \times 4 \times 5120 \times 34}{8} \times 2 \approx 2 , 863 \text{MB} \approx 2 . 80 \text{GB per layer}

This is the floor achievable without gradient checkpointing.

What You Actually Decide in Practice

  • Always enable SP whenever you enable TP. There is no scenario where you want tensor parallelism without sequence parallelism: SP is strictly better for memory and has identical communication cost. Frameworks that separate the two flags default SP to off; always check.

  • SP degree equals TP degree-they share a communicator. If TP is 8, SP is 8. They are not independently configurable, and there is no benefit to mismatching them.

  • Sequence length must be divisible by tt. A 4096-token sequence with t=8t = 8 gives exactly 512 tokens per GPU. If your sequence length is not a multiple of tt, pad to the next multiple. Uneven partitions require masked gather operations and add significant implementation complexity.

  • The stranded sbh×10s b h \times 10 term is the right diagnostic when TP memory savings disappoint. If you apply TP-8 and memory usage drops by less than 8×8 \times, the culprit is almost always the unsharded pointwise activations. Confirm by enabling SP and measuring again.

  • After TP+SP+FlashAttention, further memory reduction requires activation recomputation. The minimum activation footprint is sbh×34/ts b h \times 34 / t per layer. Reducing below this floor means discarding and recomputing activations during the backward pass at a cost of roughly 33% more forward-pass FLOPs. That trade is often worth it when the alternative is a smaller batch size that hurts data-parallel efficiency.

How This Shows Up in Interviews

  • “Tensor parallelism divides your parameter memory by tt. Does it also divide your activation memory by tt?” Strong answer: not entirely. The matmul activations (FFN intermediate, attention projections) divide by tt. The pointwise activations-LayerNorm input and output, Dropout-stay full-size on every GPU in vanilla TP. Sequence parallelism fixes this by replacing the row-parallel AllReduce with ReduceScatter + AllGather; GPU kk then holds only s/ts / t rows through the pointwise ops, and the memory for those layers divides by tt as well. Total communication cost is unchanged.

  • “Why is sequence parallelism said to be free?” The AllReduce and the ReduceScatter + AllGather pair are bandwidth-equivalent in the ring-based implementation: both move the same number of bytes per device. SP rearranges when the scatter happens but does not increase the total data transferred.

  • Senior follow-up: “Walk me through the exact communication ops for one transformer layer under TP+SP.” Strong answer: (1) Input arrives as s/ts / t rows on each GPU (scattered from the previous layer’s ReduceScatter). (2) AllGather before the column-parallel QKV or FFN projection: each GPU now holds the full s×hs \times h activation. (3) Column-parallel matmul: each GPU computes its 4h/t4 h / t columns, no communication. (4) Row-parallel matmul: each GPU holds a partial sum over the full hh-dim output. (5) ReduceScatter: partial sums merged, GPU kk keeps its s/ts / t rows. (6) LayerNorm and Dropout on s/ts / t rows. (7) Residual add with the corresponding s/ts / t-row slice held from before step (2). Repeat per layer. Two collective ops-AllGather and ReduceScatter-per MLP and per attention block, identical total bytes to the AllReduce they replace.

Key Takeaways

  • Tensor parallelism leaves LayerNorm and Dropout unsharded: their per-layer activation cost is sbh×10×2s b h \times 10 \times 2 bytes per GPU, independent of tt.

  • Sequence parallelism distributes these pointwise ops along the sequence dimension by replacing the TP AllReduce with ReduceScatter + AllGather-same bandwidth, 1/t1 / t activation memory for the pointwise ops.

  • After TP + SP + FlashAttention, the minimum per-layer activation footprint is sbh×34/ts b h \times 34 / t bytes per GPU.

  • SP degree always equals TP degree; sequence length must be divisible by tt.

  • For a 13B model at s=4096s = 4096, b=4b = 4, t=8t = 8, enabling SP recovers ≈55\approx 55 GB per GPU-the difference between fitting comfortably on an H100 and running out of memory before the optimizer state is even allocated.

12.3.2 Context Parallelism and Ring Attention

Sequence parallelism cuts the LayerNorm and Dropout memory by tt, and FlashAttention removes the quadratic softmax matrix from the per-GPU footprint entirely-for sequences that fit within a single GPU’s SRAM tiling budget. At s=4096s = 4096 that budget is comfortable. At s=128,000s = 128 , 000, the sequence LLaMA-3 trained its long-context variant at, the FlashAttention tile loop itself requires KK and VV tensors of 128K×h128 K \times h elements to be addressable from a single device’s HBM, which at h=8192h = 8192 and BF16 is 128,000×8192×2×2128 , 000 \times 8192 \times 2 \times 2 bytes ≈4.2\approx 4 . 2 GB per layer just for the KV pair-before any attention scores, before the 80 remaining layers, before gradients. Context parallelism, also called ring attention, distributes the attention computation itself across pp devices and is the last activation-memory tool you need to understand before discussing long-context training at the designinterview level.

In context parallelism, the sequence is partitioned across pp GPUs. GPU kk permanently holds the query, key, and value tokens for positions [k⋅s/p,(k+1)⋅s/p)[ k \cdot s / p , ( k + 1 ) \cdot s / p ). To produce the correct attention output for its query slice, GPU kk must compute inner products with all ss keys-including the s(1−1/p)s ( 1 - 1 / p ) keys it does not own. Ring attention solves this by rotating the K and V chunks around a logical ring of pp GPUs while each GPU keeps its own Q shard fixed.

In each of pp ring steps, GPU kk:

  1. Receives the K, V chunk from its left neighbor.

  2. Computes the partial attention scores QkKjT/dheadQ_{k} K_{j}^{T} / \sqrt{d_{\text{head}}} for the received chunk jj and accumulates the weighted values into a running output buffer.

  3. Updates its running online-softmax statistics: the current maximum logit mm and the running log-sum-exp denominator ℓ\ell.

  4. Passes the K, V chunk to its right neighbor while simultaneously receiving the next chunk from the left.

After pp steps, GPU kk has seen all ss keys and values, and its output buffer holds the exact same attention result it would have computed on a single GPU with the full sequence in memory. The mathematical correctness follows from the associativity of the online-softmax merge rule used inside FlashAttention. When a GPU receives a new KV chunk, it merges the new partial output OnewO_{\text{new}} into the running OO using:

m′=max⁡(m,mnew),ℓ′=em−m′ℓ+emnew−m′ℓnew,O′=em−m′O+emnew−m′Onewℓ′m^{\prime} = \max ( m , m_{\text{new}} ) , \qquad \ell^{\prime} = e^{m - m^{\prime}} \ell + e^{m_{\text{new}} - m^{\prime}} \ell_{\text{new}} , \qquad O^{\prime} = \frac{e^{m - m^{\prime}} O + e^{m_{\text{new}} - m^{\prime}} O_{\text{new}}}{\ell^{\prime}}

This update is associative over chunks in any order, so ring rotation produces bit-identical output to sequential full-sequence attention.

Communication cost per layer per device: each ring step sends one chunk of K and one chunk of V, each of shape s/p×dhead×nheadss / p \times d_{\text{head}} \times n_{\text{heads}}, for p−1p - 1 steps:

(p−1)×sp×dhead×nheads×2×2bytes≈2⋅s⋅h⋅2bytes( p - 1 ) \times \frac{s}{p} \times d_{\text{head}} \times n_{\text{heads}} \times 2 \times 2 \text{bytes} \approx 2 \cdot s \cdot h \cdot 2 \text{bytes}

This is O(s)O ( s )-linear in sequence length, the same order as a single AllGather of the full K or V tensor-while the memory saving is O(s2/p)O ( s^{2} / p ).

Ring attention converts the O(s2)O ( s^{2} ) attention memory problem into an O(s)O ( s ) communication problem that is independent of pp, leaving the sequence-length scaling entirely in the network rather than in HBM.

FlashAttention eliminates the full s×ss \times s softmax matrix from SRAM by tiling the key-query computation in blocks, but it still requires all ss keys and values to be reachable from the GPU’s HBM within a single forward pass. At s=128Ks = 128 K, the K and V matrices per layer already consume several gigabytes per device, and they must be replicated for the backward pass. Ring attention removes this requirement: each GPU holds s/ps / p keys and values permanently and receives the remaining (p−1)/p( p - 1 ) / p fraction as a transient stream, one chunk at a time, overlapping communication with the attention tile computation for the chunk it just received. The peak local KV footprint is therefore s/ps / p, not ss.

The alternative-splitting the sequence with sequence parallelism and hoping FlashAttention’s block loop handles the rest-breaks down because the block loop still iterates over all ss key positions. Sequence parallelism distributes the query dimension but does not distribute the key-value dimension. Ring attention distributes both.

Compared to a naive all-to-all approach where every GPU first AllGathers the full KV and then runs local FlashAttention, ring attention is strictly more memory-efficient: the all-to-all approach requires each GPU to hold all ss keys in memory simultaneously, defeating the purpose. The ring rotates through chunks that can be freed as soon as the corresponding attention tile is accumulated.

A Worked Example

LLaMA-3 70B long-context training configuration: h=8192h = 8192, nheads=64n_{\text{heads}} = 64 (GQA with 8 KV heads but illustrated here for full attention), dhead=128d_{\text{head}} = 128, 80 layers, sequence s=131,072s = 131 , 072 (128K tokens), batch b=1b = 1, context parallel degree p=8p = 8, tensor parallel degree t=8t = 8.

KV memory without context parallelism (per GPU, per layer):

s×nheads×dhead×2×2bytes=131,072×64×128×4≈4.29GB per layers \times n_{\text{heads}} \times d_{\text{head}} \times 2 \times 2 \text{bytes} = 131 , 072 \times 64 \times 128 \times 4 \approx 4 . 29 \text{GB per layer}

Over 80 layers: ≈343\approx 343 GB-impossible on any single-node system.

With context parallelism (p=8p = 8): each GPU holds 131,072/8=16,384131 , 072 / 8 = 16 , 384 tokens permanently. Local KV per layer:

4.29GB8≈537MB per layer,80layers≈41.9GB\frac{4 . 29 \text{GB}}{8} \approx 537 \text{MB per layer} , 80 \text{layers} \approx 41 . 9 \text{GB}

Combined with t=8t = 8 TP (which further shards head groups in the KV projection), the actual per-device KV footprint is closer to 41.9/8≈5.241 . 9 / 8 \approx 5 . 2 GB, leaving most of the 80 GB H100 HBM for parameters, optimizer state, and gradient buffers.

Ring communication per layer per device:

(p−1)×sp×nheads×dhead×2×2bytes=7×16,384×64×128×4≈3.76GB( p - 1 ) \times \frac{s}{p} \times n_{\text{heads}} \times d_{\text{head}} \times 2 \times 2 \text{bytes} = 7 \times 16 , 384 \times 64 \times 128 \times 4 \approx 3 . 76 \text{GB}

At NVLink ring-step bandwidth ∼600\sim 600 GB/s: approximately 6.36 . 3 ms of communication per layer. Attention compute for 16K queries over 128K keys on one GPU is ∼\sim10×\times longer at typical H100 throughput, so communication is largely masked by overlapping ring sends with tile accumulation.

What You Actually Decide in Practice

  • Enable context parallelism only at s≳32Ks \gtrsim 32 K. Below that threshold, FlashAttention within a TP group handles the attention memory without the ring-communication overhead. LLaMA-3 used CP only in the final long-context fine-tuning phase; the main pre-training run did not need it.

  • Causal attention requires a striped sequence partition. In contiguous partitioning (GPU kk holds tokens [k⋅s/p,(k+1)⋅s/p)[ k \cdot s / p , ( k + 1 ) \cdot s / p )), GPU p−1p - 1 holds only the last block of tokens, which must attend to all earlier tokens held by all other GPUs. This creates severe load imbalance: GPU 0 computes a nearly zero attention matrix (only self-attention, all future tokens masked), while GPU p−1p - 1 computes a nearly full triangle. The standard fix is striped or interleaved partitioning: GPU kk holds every pp-th token, distributing causal load evenly. Stripe partitioning is the default in production ring-attention implementations.

  • Variable-length batches require padding or sequence packing. Ring attention assumes uniform sequence length within the CP group. If sequences vary, pad each to the next multiple of pp, or use sequence packing with separate attention masks per document. Packing is significantly more complex to implement correctly and requires masked attention kernels that are aware of document boundaries within the padded token stream.

  • CP and TP use separate process groups. The pp ring-rotation peers form a CP communicator; the tt tensor-parallel peers form a TP communicator. In LLaMA-3 with t=8t = 8 and p=8p = 8, a batch item involves 64 GPUs in combined TP×\timesCP parallelism. The two communicators do not interfere, but scheduling their collective operations without stalling requires careful overlap planning in the training loop.

  • Allocate two KV buffers per device. At any ring step, a GPU holds its own outgoing KV chunk and the incoming chunk from its neighbor simultaneously. Ping-pong buffering between two HBM allocationsone for send, one for receive-is standard. Forgetting this doubles the apparent per-GPU KV overhead in memory estimates.

How This Shows Up in Interviews

  • “How would you design a training system for 128K context?” Strong answer: within each 8-GPU NVLink node, apply TP=8 for weight sharding and SP=8 for LayerNorm activation sharding. For the attention computation, apply CP=8 across the same node group (TP and CP share the NVLink fabric via separate communicators), rotating KV in a ring while each GPU holds 16K queries and uses onlinesoftmax tile accumulation. Use pipeline parallelism across nodes for model depth. The CP ring adds roughly one AllGather-equivalent of communication per attention layer per device, maskable by overlapping ring sends with tile computation.

  • “Why doesn’t sequence parallelism alone solve the 128K context memory problem?” SP splits the sequence for LayerNorm and Dropout. Each GPU still needs to iterate over all ss keys during FlashAttention’s tile loop, so the KV tensors must be accessible from every device’s HBM-O(s)O ( s ) per tensor, O(s)O ( s ) per layer, blowing up at s=128Ks = 128 K. Ring attention extends the sequence split into the attention computation itself, so no GPU ever holds more than s/ps / p KV positions in its local memory at once.

  • Senior follow-up: “Why is ring attention mathematically exact?” Strong answer: the online-softmax accumulator maintains three statistics per attention head: running maximum logit mm, running log-sum-exp denominator ℓ\ell, and weighted output sum OO. When a new KV chunk arrives, these three values can be merged from two partial results using the update equations above, which are associative and commutative over any partitioning of the key dimension. Ring attention applies this merge protocol over pp ring steps instead of over SRAM tiles; since the same merge rule applies in both cases, the output is bit-identical to single-GPU full-sequence attention.

Key Takeaways

  • Context parallelism (ring attention) distributes the attention computation across pp GPUs by rotating K and V through a logical ring while each GPU keeps its own Q shard fixed.

  • The online-softmax merge rule (from FlashAttention) guarantees mathematical exactness: the accumulated output across pp ring steps is identical to single-GPU attention.

  • Local KV memory per GPU per layer scales as O(s/p)O ( s / p ); ring communication cost scales as O(s)O ( s ) bytes per device-linear, not quadratic.

  • Causal attention requires a striped (interleaved) sequence partition to avoid severe load imbalance across ring members.

  • CP stacks with TP and SP on separate process groups; LLaMA-3 uses p=8p = 8, t=8t = 8, engaging 64 GPUs in combined attention-and-weight parallelism for a single 128K-token batch item.

12.4 Putting It All Together: 3D and 4D Parallelism

12.4.1 Rules of Thumb for Choosing a Mix

Once you have all the parallelism strategies in hand-tensor, pipeline, data, and sequence-the question that separates a strong interview answer from a great one is: in what order do you apply them, and why? This composition question appears at senior and staff-level interviews at frontier labs, typically framed as “walk me through how you would set up training for a 100B model on a thousand GPUs.” The answer is not a free choice. It is constrained by hardware topology, model size, and a finite batch-size budget, and those constraints resolve into a repeatable recipe that practitioners call 3D parallelism-tensor, pipeline, and data-optionally extended to 4D by adding sequence or expert parallelism.

The central insight is that different parallelism strategies impose different communication patterns and tolerate different network bandwidths. A GPU cluster is not a flat network: GPUs within a single machine communicate via NVLink or NVSwitch at roughly 600 GB/s aggregate per GPU (on H100s), while GPUs on different machines communicate over InfiniBand at roughly 25-50 GB/s per GPU-an 8-20×\times gap. Applying a high-bandwidth-hungry strategy across the slower link wastes throughput; the converse is merely suboptimal.

Tensor parallelism requires an all-reduce per transformer layer, carrying activation tensors proportional to s×b×hs \times b \times h (sequence length, micro-batch size, hidden dimension). At hidden size 8192 and sequence length 4096, a single all-reduce moves hundreds of megabytes. NVLink saturates this in under a millisecond; InfiniBand would take tens of milliseconds, dominating compute time entirely. Tensor parallelism must stay inside a single NVSwitch domain, typically 8 GPUs.

Pipeline parallelism transfers activations point-to-point between consecutive stages-tensors of size s×b×hs \times b \times h per stage boundary-and tolerates InfiniBand latency reasonably well. Its cost comes from the pipeline bubble: the fraction of time GPUs sit idle waiting for earlier stages. The bubble ratio is (p−1)/m( p - 1 ) / m, where pp is the number of pipeline stages and mm is the number of micro-batches in a step. Large mm hides the bubble; small mm wastes compute.

Data parallelism synchronizes gradients once per training step. With ZeRO stage 1, that is a reduce-scatter plus an all-gather over the full parameter count PP (total cost: 2P2 P bytes); with ZeRO stage 3, it is 3P3 P bytes spread across many small operations that can be overlapped with computation. Amortized over a full forward-backward pass, this communication tolerates slow inter-node links-provided the batch is large enough for compute to dominate.

The composition follows directly:

  1. Tensor parallelism first, up to the GPU count per machine (typically 8), exploiting NVLink bandwidth.

  2. Pipeline parallelism or ZeRO-3 (FSDP) second, across nodes, until the model fits in aggregate GPU memory.

  3. Data parallelism third, across all remaining GPUs, scaling throughput until the batch-size budget is exhausted.

Adding a fourth dimension-sequence parallelism (SP), context parallelism (CP), or expert parallelism (EP)- handles long-context activations or sparse MoE routing within this framework.

Match the most communication-intensive strategy to the fastest link; let data parallelism absorb the slow tail of the network hierarchy.

The batch size as resource framing makes the trade-offs concrete. A fixed global batch size BB is dictated by optimization stability-the critical batch size beyond which per-step sample efficiency saturates rapidly. Every parallelism dimension that requires micro-batches competes for that budget.

Pipeline parallelism needs m≥4pm \geq 4 p micro-batches to keep the bubble below roughly 20%. With p=8p = 8 pipeline stages, you need m≥32m \geq 32 micro-batches per DP replica per step. If your global batch is 2048 sequences and DP degree is 64, each replica has 32 sequences-exactly 32 micro-batches of size 1. Increasing pp to 16 demands m≥64m \geq 64, but the batch budget may be exhausted.

Data parallelism has no bubble, but without ZeRO it does not reduce per-GPU parameter memory. Pure DP on a 100B model requires every GPU to hold the full model; at bf16 that is 200 GB, exceeding any current GPU’s HBM. ZeRO-3 shards parameters across all DP ranks but adds 3P3 P communication overhead. For large DP degree and sufficient compute time per step, this communication is overlapped and affordable.

A useful mental picture from the TPU literature: plot compute efficiency against the ratio of batch size to GPU count. At low ratios, any strategy is communication-bound. At high ratios, pure data parallelism reaches the compute-bound regime on its own. The mixed TP + PP + DP regime occupies the middle, where no single strategy suffices and composition is the only path to high utilization.

Compared to pipeline parallelism, ZeRO-3 (FSDP) offers a simpler trade-off: 3P3 P communication per step regardless of model depth, no bubble, but higher parameter communication volume. For models whose batch budget cannot sustain a large pipeline, ZeRO-3 is almost always the right choice over PP.

A Worked Example

Consider a 70B dense transformer (h=8192h = 8192, 80 layers, 64 attention heads) on 256 H100 80 GB GPUs arranged as 32 nodes of 8 GPUs each. Target global batch: 2048 sequences of 4096 tokens, bf16 mixed precision.

Step 1 - Tensor parallelism. Set TP =8= 8 within each node. Parameter memory per GPU drops from 7070B ×2\times 2 bytes =140= 140 GB to 140/8=17.5140 / 8 = 17 . 5 GB. The model fits on a single node after TP.

Step 2 - Pipeline check. With TP =8= 8 sufficient to fit parameters, pipeline parallelism is not required. Each 8-GPU node holds one TP-sharded model replica.

Step 3 - Data parallelism. With 32 nodes and one model replica per node, DP degree D=32D = 32. Each replica processes 2048/32=642048 / 32 = 64 sequences per step.

Step 4 - Optimizer state. With ZeRO stage 1, first and second moment tensors (fp32, ≈560\approx 560 GB total for the model) are sharded across 32 DP ranks: each rank stores ≈17.5\approx 17 . 5 GB of optimizer state. Per-GPU total: 17.517 . 5 (params) +17.5+ 17 . 5 (optimizer) +17.5+ 17 . 5 (gradients in bf16) ≈52.5\approx 52 . 5 GB-comfortably below 80 GB before activations.

Step 5 - Gradient sync cost. After the backward pass, a reduce-scatter plus all-gather over 140 GB of gradients (DP degree 32) moves 140/32≈4.4140 / 32 \approx 4 . 4 GB per rank. At 25 GB/s effective InfiniBand bandwidth, this completes in ≈175\approx 175 ms. A full forward-backward pass for a 70B model at this batch size runs on the order of 200-250 ms on H100s, so gradient synchronization is almost fully overlapped with the tail of the backward pass.

Now scale to 405B: parameters at bf16 are 810 GB. After TP =8= 8, each GPU still needs 810/8=101810 / 8 = 101 GB for parameters alone, exceeding HBM. Adding PP =8= 8 splits 80 layers into 8 pipeline stages of 10 layers each; per-GPU parameter memory drops to ≈12.6\approx 12 . 6 GB. Pipeline parallelism was not chosen-it was forced by memory.

What You Actually Decide in Practice

  • Cap TP at 8, always. Crossing a node boundary for tensor parallelism forces per-layer all-reduces over InfiniBand. Published benchmarks show 40%+ throughput degradation going from TP =8= 8 (intra-node) to TP =16= 16 (cross-node). This is a consequence of bandwidth physics, not software; it holds across GPU generations.

  • Size the pipeline to keep the bubble under 20%. Use m≥4pm \geq 4 p as a starting point. With PP =8= 8, target at least 32 micro-batches per DP replica per step. If your global batch cannot support this, prefer ZeRO-3 over pipeline parallelism.

  • Choose ZeRO-3 over PP when batch sizes are tight. ZeRO-3’s overhead (3P3 P bytes per step, partially overlapped with computation) is predictable and independent of model depth. Pipeline parallelism’s bubble grows with stage count regardless of tuning, and its scheduling complexity is non-trivial-production teams have described situations where a single engineer understands the pipeline scheduler.

  • Add sequence parallelism alongside TP at no extra bandwidth cost. SP replaces the TP all-reduce on activation tensors with a reduce-scatter/all-gather pair of equivalent total cost, while also partitioning LayerNorm and Dropout activations across the TP group. Minimum activation memory per GPU drops by a factor of the TP degree; use it whenever TP ≥4\geq 4.

  • Use gradient accumulation to rescue a tight micro-batch budget. If each DP replica can hold only micro-batch size 1 after TP and PP, accumulate kk backward passes before syncing. Effective micro-batch count becomes kk, shrinking the bubble from (p−1)/m( p - 1 ) / m to (p−1)/(km)( p - 1 ) / ( k m ) at the cost of kk-fold longer steps. This is preferable to reducing DP degree and cutting overall throughput.

How This Shows Up in Interviews

  • “Design the parallelism strategy for a 100B dense transformer on 512 H100s.” Strong answer: set TP =8= 8 within each 8-GPU node (64 nodes total, 64 TP groups). A 100B model at bf16 is 200 GB; after TP =8= 8, 25 GB per GPU-fits without PP. Set DP =64= 64. With global batch 2048, each replica gets 32 sequences. Add ZeRO-1 to shard optimizer state across DP ranks. Explain why TP stays at 8, why PP is unnecessary here, and how gradient sync overlaps with computation.

  • “Why can’t you just scale data parallelism to 10,000 GPUs?” Strong answer: DP alone does not reduce per-GPU parameter memory without ZeRO-3. Beyond the critical batch size, adding more DP replicas does not improve sample efficiency-you run out of useful batch-size budget before you run out of GPUs. ZeRO-3 adds 3P3 P synchronization cost that becomes harder to overlap at very high DP degree, and parameter communication across spine switches introduces topology-dependent latency.

  • Senior follow-up: “A 400B+ training run spans multiple rack spine switches. How does this affect your PP assignment?” Strong answer: assign consecutive pipeline stages to physically adjacent nodes within the same rack or pod. PP’s point-to-point activation transfers stay on high-bandwidth leaf switch fabric; DP’s gradient all-reduces must traverse spines but happen once per step and tolerate additional latency through asynchronous overlap. Modern cluster schedulers expose rack-locality hints; use them to constrain the PP group placement before setting DP. The failure to do this is one of the common causes of unexplained throughput variance at large scale.

Key Takeaways

  • Apply TP first (cap at 8, within a single NVSwitch domain), then PP or ZeRO-3 until the model fits in memory, then DP for remaining GPUs.

  • Batch size is a finite resource: micro-batches spent hiding the pipeline bubble are not available for data-parallel scaling.

  • Communication costs: DP ≈2P\approx 2 P per step, ZeRO-3 ≈3P\approx 3 P per step, TP ≈\approx activation-size all-reduce per layer-compare each to the step’s compute time to locate the bottleneck.

  • TP beyond 8 GPUs on current GPU hardware causes 40%+ throughput degradation and should not be used.

  • Extending to 4D by adding SP or CP alongside TP costs no additional bandwidth and is the default choice for sequence lengths above 4-8K tokens.

12.4.2 Case Studies: Llama, DeepSeek-V3, Gemma

Training reports from frontier labs are among the most information-dense documents in the field-each parallelism configuration table is the compressed output of thousands of hours of ablations. Three public examples span the design space usefully: Llama 3, which sets the canonical template for dense models on GPU clusters; DeepSeek-V3, which deliberately departs from that template for a large MoE model; and Gemma 2, which shows how TPU network topology reshapes the same trade-offs. Reading these side by side reveals how the rules of thumb hold across hardware and architecture-and exactly where they bend.

Every configuration in these reports is traceable to one of three constraints: model memory, network topology, or batch-size budget. The useful mental model is to read the TP setting as the intra-node memory reduction strategy, the PP setting as the cross-node model-fit strategy, and the DP setting as the throughput-scaling strategy. When those three are set, whatever remains-CP, EP, SP-is driven by architecture-specific demands: long sequences, sparse experts, quadratic attention memory.

The three case studies share a common thread. Each configuration minimizes idle time: PP bubble, communication stalls, and memory-bound micro-batch stalls are all forms of idle time. A configuration is correct when the matrix-multiply units are busy essentially continuously.

If you can read a parallelism table and explain why each dimension is set to its specific value, you can answer most senior-level design questions about distributed training.

The three examples collectively demonstrate the universality of the bandwidth-hierarchy argument even as specific choices differ. On GPU clusters, NVLink forces TP ≤8\leq 8. On TPU pods with toroidal-mesh networkingwhere every chip can reach every neighbor at similar bandwidth-model parallelism can scale further without hitting the same cliff. For MoE models, expert parallelism collapses the TP/PP distinction: routing tokens to expert GPUs is effectively tensor parallelism in the MLP block, implemented via all-to-all rather than all-reduce.

The GPU case also illustrates a hardware-topology argument that rarely appears in textbooks: pipeline parallelism is used not just to fit models, but to route around the slow spine switches that connect GPU racks above ∼256\sim 256 GPUs. PP’s point-to-point activation transfers are directional and predictable, degrading far more gracefully than all-to-all operations when crossing network hierarchy boundaries.

A Worked Example

Llama 3 (405B, dense, GPU) is the closest to the canonical recipe and the most instructive to decode step by step.

The model has ∼405\sim 405B parameters, 126 transformer layers, hidden dimension 16,384. Training used clusters of H100 80 GB GPUs. The parallelism configuration for the standard pre-training phase:

  • TP =8= 8: one full TP group per 8-GPU node (NVLink). At bf16, 405B parameters across 8 GPUs gives ≈101\approx 101 GB per GPU for parameters alone-still over the 80 GB HBM limit. TP alone cannot fit the model; PP is necessary.

  • PP =16= 16: 126 layers split across 16 pipeline stages of ∼8\sim 8 layers each. Per-GPU parameter memory drops to ≈12.6\approx 12 . 6 GB, well within budget. TP ×\times PP =128= 128 GPUs constitutes one full model instance.

  • DP == (total GPUs) / 128: on a cluster of roughly 16,000 GPUs, DP ≈125\approx 125 replicas. Global batch: 125×4125 \times 4 micro-batches ×\times 4096 tokens ≈2\approx 2M tokens per step.

  • CP =8= 8: added specifically for the long-context fine-tuning phase (128K token sequences). Attention computation distributes across 8 devices via ring-style KV passing, keeping the quadratic attention activation memory tractable. This dimension is disabled for standard pre-training.

The key lesson: PP was forced by memory, not preferred. The team would have avoided it if 80 GB HBM were sufficient after TP =8= 8. The long-context CP dimension illustrates the 4D pattern: a fifth dimension would be needed if the model were also MoE (expert parallelism), which it is not.

At this scale, Llama 3’s training report logged 148 interruptions from faulty GPUs and additional incidents from unplanned network maintenance. The parallelism configuration is only part of the engineering; checkpoint frequency, automatic restart logic, and silent data-corruption detection are equally critical once you cross ∼10\sim 10K GPUs.

DeepSeek-V3 (671B total, 37B active, MoE, GPU) is the most interesting departure from the canonical recipe. The model uses 256 routed experts with top-8 routing per token, so only 37B parameters are active for any given forward pass.

Published configuration:

  • PP =16= 16: 61 transformer layers split across 16 pipeline stages.

  • EP =64= 64: 64-way expert parallelism. Each GPU holds a subset of experts. Tokens are dispatched to the GPU holding the selected expert via all-to-all communication. This substitutes for the role TP plays in dense models.

  • ZeRO stage 1 for DP: optimizer states sharded across DP ranks. DP degree fills remaining GPU capacity.

  • No tensor parallelism.

The absence of TP is the key departure and worth understanding precisely. In a 671B MoE model, the expert MLP layers dominate parameter count. Expert parallelism distributes those parameters across 64 GPUs via token routing-achieving the same memory reduction that TP provides in dense models, but through a different communication primitive (all-to-all for routing versus all-reduce for TP). Adding TP on top of EP would introduce per-layer all-reduces on the attention and dense MLP blocks, which are relatively small in a MoE model, while providing little incremental memory benefit. The net effect would be increased communication overhead for negligible memory gain-a straightforward throughput loss.

The lesson: the canonical TP-first rule assumes a dense model. For MoE, EP substitutes for TP. The bandwidth-hierarchy argument still applies-EP’s all-to-all also benefits from fast intra-pod networking-but the primitive changes.

Gemma 2 (dense, TPU) illustrates how topology reshapes every decision. TPU pods use a toroidal mesh: each chip connects to a small number of spatial neighbors, but collective bandwidth across the full mesh is high and symmetric for any global operation. There is no hard bandwidth cliff at a node boundary analogous to the NVSwitch-to-InfiniBand transition on GPU clusters.

Published configuration: ZeRO-3 (FSDP) ++ model parallelism ++ data parallelism. Model parallelism can extend well beyond 8 devices on TPUs because all inter-chip links are comparable in bandwidth; there is no architectural reason to cap at 8. ZeRO-3 is the primary vehicle for memory efficiency, sharding parameters, gradients, and optimizer state across all model-parallel ranks. The report does not feature pipeline parallelism, for two reasons: the toroidal mesh has no slow inter-rack spine to route around, and ZeRO-3 handles the memory constraint without the scheduling complexity of PP.

What You Actually Decide in Practice

  • Dense GPU training (Llama-style): TP =8= 8, PP only when forced by memory (roughly above 100B at 80 GB HBM), DP for the rest. Add CP only for the long-context fine-tuning phase; it adds pipeline complexity and should be disabled for standard pre-training.

  • MoE GPU training (DeepSeek-style): replace TP with EP sized to the expert count. PP still handles the dense attention and embedding layers. Expect the all-to-all routing to be a communication bottleneck at high expert count; token dropping and expert-capacity buffers control variance but add algorithm complexity that must be tested carefully.

  • TPU training (Gemma-style): lean on ZeRO-3 and model parallelism beyond 8 chips. Pipeline parallelism is less commonly needed because there is no hard bandwidth cliff at a hardware boundary. Collective operations scale to the full pod size without topology-driven restrictions.

  • At production scale, fault tolerance is a first-class concern. Llama 3 counted 148 GPU failures during a single training run. The correct response is aggressive checkpointing (every 30-60 minutes at this scale), automatic job restart, and silent-error detection via periodic gradient norm checks and loss spike monitors. The parallelism configuration alone does not determine success.

How This Shows Up in Interviews

  • “DeepSeek-V3 has no tensor parallelism. Why not?” Strong answer: in a 671B MoE model, expert MLP layers dominate parameters. Expert parallelism distributes those parameters across 64 GPUs via token routing, providing the same memory reduction that TP gives in dense models. Adding TP would introduce per-layer all-reduces on the smaller attention and dense MLP layers with minimal memory benefit-net throughput loss.

  • “Why does Llama 3 use context parallelism only for the long-context phase?” Strong answer: CP distributes attention computation by passing KV blocks in a ring across multiple devices. It is necessary only when the quadratic s2hs^{2} h activation memory from attention exceeds per-GPU HBM. At standard pretraining sequence lengths (4096 tokens), TP ++ SP already keeps activation memory tractable. At 128K tokens, the quadratic term exceeds any mitigation from SP alone, making CP unavoidable.

  • Senior follow-up: “If you had to train a 1T dense model on GPUs today, what configuration would you propose, and what would you be most worried about?” Strong answer: TP =8= 8, PP sized to fit the model (likely PP =32= 32+), DP for remaining capacity. The primary concern is pipeline bubble: at PP =32= 32, the bubble ratio (31)/m( 31 ) / m demands m≥128m \geq 128 micro-batches per replica to stay under 20% idle time, which strains the global batch budget severely. Secondary concerns: checkpoint and restart time scales with the run duration; a single mid-step GPU failure at 32K GPUs wastes hours. The engineering response is zero-bubble pipeline scheduling (filling idle slots with weight-gradient computation) and a battle-tested fault-tolerant checkpoint infrastructure-both are as important as the parallelism configuration itself.

Key Takeaways

  • Dense GPU models follow the canonical recipe: TP =8= 8 within a node, PP only when forced by memory, DP for throughput. CP is a 4th dimension added only for long-context phases.

  • MoE models substitute EP for TP; the bandwidth-hierarchy argument still applies, but the communication primitive changes from all-reduce to all-to-all.

  • TPU topology removes the 8-device NVLink cliff, allowing model parallelism to scale further and reducing reliance on pipeline parallelism.

  • At production scale (>>10K GPUs), fault tolerance, silent-error detection, and checkpoint infrastructure are as critical as the parallelism configuration itself.

  • Reading a published parallelism table backward-inferring why each dimension is set to its specific value from memory, topology, and batch-size constraints-is the core skill this interview topic tests.