Learn: Model Parallelism

What you're learning

How to split the computation itself — not just the optimizer state — across GPUs, when a single device can no longer hold or feed a layer. By the end you should be able to (1) derive why Megatron’s column→row GEMM pairing needs only one all-reduce per block, (2) derive the pipeline bubble fraction $\frac{p-1}{m+p-1}$ from scratch, (3) explain sequence and context parallelism as sequence-dimension splits, (4) reason about the MoE all-to-all, and (5) lay out a full DP × TP × PP (× EP × CP) plan onto a real cluster topology.

This lesson assumes lessons 08–09: 08-data-parallelism-and-collectives (all-reduce, ring/tree collectives, NVLink vs InfiniBand) and 09-zero-and-fsdp (sharding optimizer state, gradients, params with ZeRO/FSDP). The crucial hand-off from lesson 09: ZeRO/FSDP shard state but every GPU still runs the same computation on its own data. Here we do the opposite — we split the computation.

---

1. Learning map

graph TD
    Z["Lesson 09:<br/>ZeRO/FSDP shard STATE<br/>(each GPU: same compute)"] --> M["Why model parallelism:<br/>split the COMPUTE itself"]
    C["Lesson 08:<br/>collectives — all-reduce,<br/>all-gather, reduce-scatter,<br/>all-to-all, NVLink vs IB"] --> M

    M --> TP["Tensor parallelism (TP)<br/>split each matmul<br/>Megatron col→row"]
    M --> PP["Pipeline parallelism (PP)<br/>split LAYERS into stages"]

    TP --> SP["Sequence parallelism (SP)<br/>split LN/dropout/residual<br/>along sequence"]
    TP --> EP["Expert parallelism (EP)<br/>MoE: split experts,<br/>all-to-all dispatch/combine"]

    PP --> BUB["Pipeline bubble<br/>GPipe → 1F1B →<br/>interleaved 1F1B"]

    M --> CP["Context parallelism (CP)<br/>Ring Attention:<br/>split sequence, overlap KV"]

    TP --> COMP["Composition:<br/>3D/4D/5D parallelism<br/>DP × TP × PP × EP × CP"]
    PP --> COMP
    SP --> COMP
    EP --> COMP
    CP --> COMP

    COMP --> TOPO["Map to topology:<br/>TP inside node (NVLink),<br/>PP/DP across nodes (IB)"]
    TOPO --> N["Lesson 11:<br/>putting it together at scale"]

    style M fill:#44a,color:#fff
    style COMP fill:#4a4,color:#fff
    style TOPO fill:#4a4,color:#fff

Figure 1 — Dependency graph. Everything branches from one decision: once state-sharding (lesson 09) is exhausted, you split the compute. The four axes (TP, PP, EP, CP) each attack a different bottleneck, and the final skill is composing them onto a heterogeneous interconnect.

---

2. Why model parallelism — when sharding state isn’t enough

Lesson 09 gave you a powerful lever: ZeRO-3 / FSDP shard parameters, gradients, and optimizer state across $N$ data-parallel ranks, so per-GPU state memory scales as $O(1/N)$. That solves the static memory problem. But three things it does not solve force us to split the computation:

1. A single layer’s activation working set can exceed one GPU. FSDP all-gathers a layer’s params, computes on the full local batch, then frees them. During that compute, the activations of that layer live entirely on one GPU. At large hidden size $h$, long sequence $s$, and micro-batch $b$, the activation tensor $b\cdot s\cdot h$ — and the attention score matrix $b\cdot n_{\text{head}}\cdot s^2$ — can simply not fit, no matter how you shard the weights.¹

2. You want to cut activation memory without shrinking the batch. Activation memory is what dominates at scale (lesson 06). Recomputation (checkpointing) trades it for FLOPs, but there’s a floor. Splitting the tensors across GPUs cuts activation memory proportionally — a lever orthogonal to ZeRO.

3. You want to keep per-GPU work high while growing the model. With pure data parallelism, making the model bigger means each GPU does more work and holds more state until it OOMs. Splitting layers/matmuls lets you grow model size by adding GPUs to the model dimension, keeping each GPU’s slice tractable.

The one distinction that organizes this whole lesson

ZeRO/FSDP shard state; model parallelism splits computation. Under ZeRO, every GPU eventually runs the same matmul (after gathering the weights). Under tensor parallelism, each GPU runs a different piece of the matmul and they combine results with a collective. This is why model parallelism adds communication inside the forward/backward of a single layer, whereas DP/ZeRO communication happens at gradient-sync / layer boundaries. The former is far more latency-sensitive — which is exactly why it must live on the fastest interconnect you have.

The four axes we’ll build:

Axis	Splits…	To reduce…	Comms primitive	Where it lives
TP (tensor)	each matmul (hidden dim)	param + activation mem, per-GPU FLOPs	all-reduce (or RS+AG)	inside a node (NVLink)
SP (sequence)	LN/dropout/residual (seq dim)	activation mem	all-gather / reduce-scatter	fused with TP, inside node
PP (pipeline)	layers into stages	param mem	point-to-point send/recv	across nodes (IB)
EP (expert)	MoE experts	expert param mem	all-to-all	node / across nodes
CP (context)	the sequence, for attention	activation mem for long $s$	ring send/recv (P2P)	across a ring of GPUs

---

3. Tensor parallelism — splitting the matmul itself

Tensor parallelism (TP), introduced by Megatron-LM (Shoeybi et al., 2019),² is intra-layer model parallelism: it partitions the weight matrix of a single GEMM across $t$ GPUs so each holds a slice, computes a partial result, and the group combines them with one collective. The elegance is entirely in how you pair the two GEMMs of a block so that only one all-reduce is needed.

3.1 Two ways to split a GEMM

Consider $Y=XA$ with input $X\in {\mathbb{R}}^{(b\cdot s)\times k}$ and weight $A\in {\mathbb{R}}^{k\times n}$.

Column-parallel: split $A$ along its columns (output dim): $A=[A_1,A_2,\dots ,A_t]$. Each GPU $i$ holds $A_i\in {\mathbb{R}}^{k\times n/t}$ and computes $Y_i=X{A}_i$. No communication is needed to produce $Y_i$ — each GPU independently owns a slice of the output columns. $X$ must be replicated on every GPU.

$Y=X[A_1|\cdots |A_t]=[X{A}_1|\cdots |X{A}_t]=[Y_1|\cdots |Y_t]$

Row-parallel: split $B\in {\mathbb{R}}^{n\times k}$ along its rows (input dim): $B=[B_1;B_2;\dots ;B_t]$, and split the input $Y$ along columns to match: $Y=[Y_1|\cdots |Y_t]$. Each GPU computes a partial sum $Z_i=Y_i{B}_i$, and the true output is the sum:

$Z=YB={\sum }_{i=1}^t{Y}_i{B}_i\Rightarrow \text{all-reduce over }t\text{ GPUs.}$

3.2 The column→row pairing (the key trick)

The MLP block is $Z=\text{GeLU}(XA)B$ with $A$ the up-projection ($h\to 4h$) and $B$ the down-projection ($4h\to h$). Megatron makes $A$ column-parallel and $B$ row-parallel:

graph LR
    X["X (replicated)<br/>b·s × h"] --> A1["GPU0: A₁ (h × 4h/t)"]
    X --> A2["GPU1: A₂ (h × 4h/t)"]
    A1 --> G1["GeLU(XA₁)<br/>b·s × 4h/t"]
    A2 --> G2["GeLU(XA₂)<br/>b·s × 4h/t"]
    G1 --> B1["GPU0: B₁ (4h/t × h)"]
    G2 --> B2["GPU1: B₂ (4h/t × h)"]
    B1 --> S["all-reduce (g)<br/>Σ Bᵢ partials"]
    B2 --> S
    S --> Z["Z<br/>b·s × h"]

    style S fill:#a44,color:#fff

Figure 2 — Tensor-parallel MLP, $t=2$. The up-proj $A$ is column-split so each GPU owns a slice of the 4h intermediate; GeLU is elementwise so it applies independently per slice — no comms. The down-proj $B$ is row-split, so each GPU produces a partial sum of the output, resolved by a single all-reduce. One collective for the whole two-GEMM block.

Why does this need only one all-reduce? Because the column-parallel first GEMM produces exactly the column-sharded layout that the row-parallel second GEMM consumes as input. The intermediate $\text{GeLU}(X{A}_i)$ never needs to be gathered — it stays sharded and flows straight into $B_i$. The nonlinearity is the reason the split must be this way round: GeLU is elementwise, so $\text{GeLU}([Y_1|Y_2])=[\text{GeLU}(Y_1)|\text{GeLU}(Y_2)]$ commutes with column-sharding. If you did row-parallel first, you’d need to all-reduce before the nonlinearity (to get the correct pre-activation), then again after the second GEMM — two collectives.²

Conjugate operators $f$ and $g$

Megatron formalizes this as two conjugate operators. $f$ is identity in forward / all-reduce in backward; $g$ is all-reduce in forward / identity in backward. The MLP block is $Z=g(\text{GeLU}(f(X)A)B)$. In forward you all-reduce once (at $g$); in backward you all-reduce once (at $f$). Same for attention. Total: 2 all-reduces forward + 2 backward per transformer layer (one per block).²³

3.3 Attention is parallel over heads

Multi-head attention is already embarrassingly parallel across heads — each head computes on its own $d_{\text{head}}$-dim subspace and never interacts with other heads until the output projection. So Megatron assigns whole heads to GPUs: $Q/K/V$ projections are column-parallel (each GPU produces $Q,K,V$ for its heads), the per-head attention + softmax runs entirely locally (no comms — this is why FlashAttention composes cleanly with TP), and the output projection $W_O$ is row-parallel, closing with the single all-reduce.² Constraint: num_attention_heads must be divisible by $t$, and hidden_size divisible by $t$.⁴

3.4 Communication cost — and why TP stays inside a node

Each all-reduce is over an activation tensor of $b\cdot s\cdot h$ elements. A ring all-reduce (lesson 08) moves $2\frac{t-1}{t}$ times the tensor size in bytes per GPU (reduce-scatter + all-gather). Per transformer layer, with 2 all-reduces forward + 2 backward, in bf16 (2 bytes/elem):

$V_{\text{TP/layer}}=4\times 2\frac{t-1}{t}(bsh)\times 2\text{bytes}=16\frac{t-1}{t}bsh\text{bytes.}$

The critical property: this volume is independent of $t$ for large $t$ (the factor $\frac{t-1}{t}\to 1$) and — worse — it is exposed on the critical path of every layer, with no useful compute to hide behind it (the all-reduce sits between two dependent GEMMs). This is the entire reason TP must run over the highest-bandwidth, lowest-latency link available:

TP does not scale across nodes

On a DGX/HGX H100 node, all-reduce runs over NVLink/NVSwitch (~900 GB/s bidirectional per GPU);⁵ across nodes you fall to InfiniBand (~400 Gb/s = 50 GB/s per direction) — roughly an order of magnitude slower, and higher latency. Because TP’s all-reduces are on the per-layer critical path, crossing the node boundary tanks throughput. Rule of thumb: keep $\text{TP}\le 8$ (one node’s NVLink domain). This limit is echoed across the field — NVIDIA’s own guidance is “max TP within the node, up to TP=8.”⁴³

Quiz: Why does the column→row pairing need only one all-reduce, and what breaks if you swap the order (row-parallel first)?

Answer

The column-parallel first GEMM leaves the intermediate activation sharded along the 4h dimension, exactly the input layout the row-parallel second GEMM expects. The elementwise GeLU commutes with that sharding — $\text{GeLU}$ of a column-slice equals the column-slice of $\text{GeLU}$ — so the nonlinearity applies locally with no communication, and the intermediate never has to be gathered. Only the row-parallel output is a partial sum, needing one all-reduce.

If you went row-parallel first: each GPU would produce a partial sum of the pre-activation $XA$. But GeLU is nonlinear, so $\text{GeLU}({\sum }_i{\text{partial}}_i)\ne {\sum }_i\text{GeLU}({\text{partial}}_i)$ — you’d have to all-reduce before the GeLU to form the correct pre-activation, then all-reduce again after the second GEMM. Two collectives instead of one. The pairing is chosen so the sharded layouts of consecutive GEMMs match and the nonlinearity sits entirely inside a shard.

---

4. Sequence parallelism — cut activation memory in the non-TP regions

TP shards the matmul activations, but look at what it doesn’t touch: the LayerNorm, dropout, and residual add around each block. In vanilla Megatron TP these operate on the full $b\cdot s\cdot h$ activation, replicated identically on all $t$ TP ranks — pure redundancy that dominates activation memory at scale.¹

Sequence parallelism (SP) (Korthikanti et al., 2022) splits these regions along the sequence dimension: rank $i$ owns tokens $s\cdot i/t:s\cdot (i+1)/t$. LayerNorm and dropout are per-token operations, so they parallelize trivially over the sequence — no communication within them. Activation memory for these regions drops by $t\times$.¹

The beautiful part is how SP splices into TP at the block boundaries. The TP block wants its input replicated (sharded over hidden); the SP region has it sharded over sequence. Converting between the two:

• Entering the TP block (the $f$ operator): sequence-sharded → hidden-replicated is an all-gather (forward) / reduce-scatter (backward).
• Leaving the TP block (the $g$ operator, previously an all-reduce): hidden-partial → sequence-sharded is a reduce-scatter (forward) / all-gather (backward).

SP is "free" — same comms volume as plain TP

A ring all-reduce is a reduce-scatter followed by an all-gather. SP replaces the single all-reduce ($g$) with a reduce-scatter, and replaces the identity ($f$) with an all-gather. Total communication volume is identical to TP-only — you’ve just re-associated the same bytes so that between the reduce-scatter and the next all-gather each GPU holds only $1/t$ of the activation. You get the $t\times$ activation-memory reduction on the LN/dropout regions at zero extra bandwidth.¹ This is why modern Megatron enables SP whenever TP > 1.⁴

Quiz: SP replaces TP's single all-reduce with a reduce-scatter + a later all-gather. Why is this "free" in bandwidth, yet a memory win?

Answer

A ring all-reduce is literally implemented as a reduce-scatter followed by an all-gather — same bytes on the wire. SP doesn’t add a collective; it re-associates the two halves so the reduce-scatter happens at the block exit ($g$) and the all-gather at the next block entry ($f$), with the LayerNorm/dropout/residual computation sitting in between. During that window each GPU holds only $1/t$ of the sequence, so those regions’ activations shrink $t\times$. Bandwidth is unchanged (it was already RS+AG under the hood); the win is that the fully-replicated activation now exists only as a $1/t$ shard for the duration of the non-TP work.

---

5. Pipeline parallelism — splitting layers into stages

TP splits within a layer and is bandwidth-hungry; pipeline parallelism (PP) splits across layers — assign layers $\mathrm{1..}L/p$ to stage 0, the next chunk to stage 1, etc. Communication is only the activation at stage boundaries, sent GPU→GPU with cheap point-to-point send/recv (no collective) — so PP tolerates the slower inter-node IB link. The cost is a new enemy: the pipeline bubble.

5.1 The bubble

Naively, stage $i$ can’t start until stage $i-1$ finishes forwarding the whole batch, so only one GPU works at a time — $p\times$ slowdown. The fix is micro-batching (GPipe, Huang et al., 2018):⁶ split the global batch into $m$ micro-batches and stream them through, so stages overlap. But the pipe still has to fill (stages $\mathrm{2..}p$ idle while microbatch 1 propagates) and drain (stages $\mathrm{1..}p-1$ idle at the end). That idle time is the bubble.

graph TD
    subgraph GPipe["GPipe: all forwards, then all backwards (high activation mem)"]
        direction LR
        g["fill: p-1 idle → m fwd → m bwd → drain: p-1 idle<br/>each stage stores activations for ALL m microbatches → O(m) memory"]
    end
    subgraph F1["1F1B: interleave 1 fwd / 1 bwd in steady state (same bubble, less memory)"]
        direction LR
        f["after warmup, each stage alternates F,B,F,B...<br/>a microbatch's activations freed right after its backward → O(p) memory"]
    end
    GPipe --> F1

Figure 3 — GPipe vs 1F1B. Both have the identical bubble fraction; the difference is when backward passes happen, which sets peak activation memory. See the space-time diagram below.

Here is the space-time schedule that makes the bubble concrete (rows = GPUs/stages, columns = time; F${}_i$ = forward of microbatch $i$, B${}_i$ = backward). $p=4$, $m=4$:

Figure 4 — GPipe space-time diagram, $p=m=4$. Green = forward, orange = backward, hatched red = bubble. The bubble is the $(p-1)$ steps of fill plus $(p-1)$ of drain that every stage cannot avoid.

5.2 Deriving the bubble fraction

Let $t_f$, $t_b$ be the per-microbatch forward/backward time on one stage. Total useful work along the critical path (steady state, once the pipe is full) is $m$ microbatches each costing $(t_f+t_b)$:

$T_{\text{useful}}=m(t_f+t_b).$

The bubble is the fill + drain: the last stage waits $p-1$ forward-steps to receive microbatch 1, and the first stage waits $p-1$ backward-steps at the end:

$T_{\text{bubble}}=(p-1)(t_f+t_b).$

$\boxed{\text{bubble fraction}=\frac{T_{\text{bubble}}}{T_{\text{bubble}}+T_{\text{useful}}}=\frac{(p-1)(t_f+t_b)}{(m+p-1)(t_f+t_b)}=\frac{p-1}{m+p-1}}$

For $m\gg p$ this is $\approx \frac{p-1}{m}$. The ratio is what matters, not absolute size: $p=8,m=16$ and $p=16,m=32$ both give $\frac{7}{23}=\frac{15}{47}\approx$ 30–32%.³ To make the bubble small you push $m$ up — but $m$ is capped by memory (GPipe holds activations for all $m$) and by the global batch size.

5.3 1F1B and interleaved 1F1B

PipeDream 1F1B (Narayanan et al., 2018/2021)⁷ keeps the same bubble fraction as GPipe but slashes memory: after a warmup of $p-1$ forwards, each stage settles into one-forward-one-backward in steady state, so a microbatch’s activations are freed as soon as its backward completes. Peak activation memory drops from $O(m)$ (GPipe) to $O(p)$ — this is what made deep pipelines feasible at all.³

Common misconception: "1F1B reduces the bubble"

It does not. GPipe and standard 1F1B have the identical bubble $\frac{p-1}{m+p-1}$. 1F1B is purely a memory optimization (reorder backwards earlier → free activations sooner). The thing that actually reduces the bubble is interleaving.³

Interleaved 1F1B (Megatron)³ gives each device $v$ non-contiguous chunks of layers (“virtual pipeline stages”) instead of one contiguous block. Each chunk is $1/v$ the size, so the fill/drain cost per chunk shrinks and the bubble drops by a factor of $v$:

$\text{bubble fraction (interleaved)}=\frac{1}{v}\cdot \frac{p-1}{m}.$

The price: $v\times$ more point-to-point communication (each microbatch now crosses stage boundaries $v$ times as often). Megatron makes this affordable with a scatter/gather optimization that does the extra tensor-reassembly over fast intra-node NVLink.³ So interleaving pays off when your interconnect has headroom, and stops paying when comms becomes the bottleneck. Constraint: $m$ must be a multiple of $p$ for the interleaved schedule.³

The 2024+ frontier: zero-bubble

Zero-Bubble Pipeline Parallelism (Qi et al., 2024) splits the backward pass into B (input-gradient) and W (weight-gradient) sub-passes; since W has no downstream dependency it can be slotted into the bubble, driving it toward zero under a synchronous schedule.⁸ [recent]

Quiz: Derive the bubble fraction for $p=4$ stages, $m=8$ microbatches under GPipe. Then say what interleaving with $v=2$ does to it, and what it costs.

Answer

GPipe/1F1B: $\frac{p-1}{m+p-1}=\frac{3}{8+3}=\frac{3}{11}\approx 27\%$.

Interleaved with $v=2$: $\frac{1}{v}\cdot \frac{p-1}{m}=\frac{1}{2}\cdot \frac{3}{8}=\frac{3}{16}\approx 19\%$ (using the large-$m$ form; the exact schedule is a bit better). Each device now holds 2 non-contiguous layer chunks.

Cost: roughly $2\times$ the point-to-point send/recv volume, because each microbatch crosses stage boundaries twice as often. Worth it only if the interconnect (ideally NVLink via Megatron’s scatter/gather trick) can absorb the extra traffic; otherwise the added comms eats the bubble savings.

---

6. Context parallelism — Ring Attention for long sequences

TP/SP/PP all leave one thing on a single GPU: the attention over the full sequence for a given set of heads. At million-token context, the KV cache and the $O(s^2)$ attention working set blow up regardless of how you shard weights. Context parallelism (CP) partitions the sequence itself across GPUs, so each holds only $s/N$ query/key/value tokens.

The problem: attention is not token-local — query block $i$ must attend to all key/value blocks, which live on other GPUs. Ring Attention (Liu, Zaharia, Abbeel, 2023)⁹ solves this by arranging the CP GPUs in a ring: each GPU computes attention of its local queries against its local KV block, then passes its KV block to the next GPU in the ring while simultaneously computing on the KV block it just received. Over $N$ steps every query has seen every KV block. Because attention is computed blockwise with online softmax (FlashAttention, lesson 06), partial results accumulate correctly without ever materializing the full score matrix.⁹

The key win: communication is fully hidden

The KV send/recv for step $k+1$ is launched while the blockwise attention for step $k$ computes. As long as (compute per block) ≥ (time to ship one KV block), the ring communication is completely overlapped with compute — Ring Attention adds no wall-clock overhead versus single-device, up to $N\times$ longer sequences.⁹ This is the same overlap philosophy as FSDP prefetch (lesson 09), applied to the sequence dimension.

Causal masking breaks the load balance

With a causal mask, later query blocks attend to more KV blocks than earlier ones, so a naive ring leaves some GPUs idle. Striped Attention (Brandon et al., 2023) permutes tokens so each GPU gets a balanced mix of early/late positions, restoring even load.¹⁰ Production CP implementations (Megatron-Core, Transformer Engine) use this kind of balancing. [recent]

CP composes with SP (both split the sequence) and is the standard axis for long-context training today.⁴

---

7. Expert parallelism — the MoE all-to-all

Mixture-of-Experts (MoE) replaces a dense FFN with $E$ expert FFNs and a router that sends each token to its top-$k$ experts (typically $k=1$ or $2$; DeepSeek-V3 uses fine-grained routing with many small experts).¹¹ The whole point is conditional compute: parameters grow with $E$ but per-token FLOPs stay ~constant. But now the experts — which hold most of the parameters — are too big to replicate. Expert parallelism (EP) places different experts on different GPUs.

This creates a communication pattern unlike anything above. A token arrives on the GPU that holds its sequence position, but must be computed by the GPU holding its chosen expert — an arbitrary all-to-all shuffle:

graph LR
    subgraph before["tokens by position"]
        T0["GPU0<br/>tok a,b,c,d"]
        T1["GPU1<br/>tok e,f,g,h"]
    end
    R["router: top-k<br/>expert assignment"]
    subgraph after["tokens by expert (after dispatch)"]
        E0["GPU0: Expert 0<br/>gets a,e,..."]
        E1["GPU1: Expert 1<br/>gets c,g,..."]
    end
    T0 --> R
    T1 --> R
    R -->|"all-to-all<br/>DISPATCH"| E0
    R -->|"all-to-all<br/>DISPATCH"| E1
    E0 -->|"all-to-all<br/>COMBINE"| out["tokens back<br/>to original position"]
    E1 -->|"all-to-all<br/>COMBINE"| out

Figure 5 — MoE expert parallelism. Two all-to-all collectives per MoE layer: dispatch (route tokens to the GPU holding their expert) and combine (send expert outputs back to the token’s home position).

7.1 Why the all-to-all is the hard part

Unlike all-reduce, whose volume is independent of the parallelism degree, all-to-all volume is proportional to token count × hidden dim, and as EP grows it pushes traffic cross-node where bandwidth is scarce.¹² For fine-grained MoEs the pain compounds: many experts ⇒ more dispatch/combine ops per layer (DeepSeek-V3 has ~58 MoE layers, each with two all-to-alls), large EP spanning nodes ⇒ inter-node latency dominates, and small experts ⇒ low arithmetic intensity, leaving little compute to hide the comms behind. Megatron’s report notes comms can exceed 50% of training time in DeepSeek-V3-scale MoE without optimization.¹²

Mitigations in current practice (2025–26): DeepEP (DeepSeek’s NVSHMEM-based dispatch/combine kernels) and NVIDIA’s HybridEP overlap all-to-all with expert GEMMs and hit near-hardware-limit bandwidth; overlapping the shared-expert MLP with the dispatch is standard; and “1F1B A2A overlap” hides EP comms inside the pipeline schedule.¹²

7.2 Load balancing

The router can collapse onto a few popular experts, starving the rest and creating stragglers (the all-to-all waits for the slowest expert). Two levers:¹²

• Auxiliary load-balancing loss — penalize uneven expert assignment (Shazeer et al., 2017);¹³ DeepSeek-V3 notably moves to an auxiliary-loss-free bias-adjustment scheme to avoid the aux-loss quality tax.¹¹
• Capacity factor + token dropping — cap tokens per expert; overflow tokens are dropped (skip the expert), bounding the worst-case imbalance.

EP is orthogonal to TP/DP — it's a different axis

Experts are replicated across the DP dimension and split across the EP dimension. You can (and do) combine EP with TP inside each expert. Megatron exposes expert-model-parallel-size and expert-tensor-parallel-size separately from the dense TP/DP; a common tuning move is to shift the split from EP toward expert-TP to rebalance load across ranks.¹²

Quiz: Why is the MoE all-to-all fundamentally harder to scale than the TP all-reduce, and what makes it worse for fine-grained MoEs like DeepSeek-V3?

Answer

An all-reduce’s per-GPU volume is independent of the group size (it’s $\approx 2\times$ the tensor, forever) — you can grow TP without the comms bill exploding. An all-to-all’s volume scales with tokens × hidden, and as you grow EP the shuffle inevitably spills across nodes onto the slow IB link, so cost grows with parallelism degree and with model activity.

Fine-grained MoE compounds three ways: (1) frequency — many experts + many MoE layers means two all-to-alls (dispatch + combine) per layer, dozens of times per forward; (2) cross-node bottleneck — large EP spans nodes, and inter-node all-to-all latency dominates; (3) low arithmetic intensity — each expert is small, so its GEMM finishes fast and leaves little compute to overlap the comms behind. Megatron reports comms can exceed 50% of DeepSeek-V3-scale training time without dedicated overlap (DeepEP/HybridEP, shared-expert overlap, 1F1B A2A overlap).

---

8. Composing everything — 3D/4D/5D parallelism

No single axis suffices at frontier scale. The full configuration is a product:

$\text{total GPUs}=\underset{\text{DP}}{\underbrace{d}}\times \underset{\text{TP}}{\underbrace{t}}\times \underset{\text{PP}}{\underbrace{p}}[\times \underset{\text{EP}}{\underbrace{e}}\times \underset{\text{CP}}{\underbrace{c}}]$

Each axis attacks a different bottleneck, and — crucially — each has a different communication profile, which dictates where on the hardware topology it must live:

graph TD
    subgraph node0["Node 0 (8× H100, NVLink/NVSwitch ~900 GB/s)"]
        g0["GPU0"] --- g1["GPU1"] --- g2["GPU2"] --- g3["GPU3"]
        g4["GPU4"] --- g5["GPU5"] --- g6["GPU6"] --- g7["GPU7"]
    end
    subgraph node1["Node 1 (8× H100, NVLink)"]
        h0["GPU0..7"]
    end
    node0 -. "InfiniBand ~50 GB/s/dir<br/>(PP send/recv, DP all-reduce)" .- node1

    TPlab["TP=8 + SP: inside a node<br/>(all-reduce every layer, needs NVLink)"] -.-> node0
    PPlab["PP: across nodes<br/>(P2P at stage boundary, IB-tolerant)"] -.-> node1
    DPlab["DP/ZeRO: outermost<br/>(grad all-reduce once/step, IB-tolerant)"] -.-> node1

Figure 6 — Topology mapping. TP’s per-layer all-reduces demand NVLink, so TP is confined to one node (≤8). PP’s boundary P2P and DP’s once-per-step gradient sync tolerate the ~10× slower InfiniBand, so they span nodes.

The mapping heuristics (established practice):³⁴

1. TP within a node, TP ≤ 8. Its all-reduces are on the per-layer critical path; keep them on NVLink. Enable SP alongside it for free activation savings.
2. PP across nodes. Only boundary activations cross the link, via cheap P2P; IB is fine. Size $p$ so that per-stage params fit, and pick $m$ (microbatches) large enough that $\frac{p-1}{m+p-1}$ is tolerable; interleave if NVLink allows.
3. DP / ZeRO outermost. Gradient all-reduce (or reduce-scatter+all-gather for FSDP) happens once per optimizer step and overlaps with backward — most latency-tolerant, so it spans the most nodes.
4. EP sized to fit experts; keep intra-node when possible because all-to-all is bandwidth-heavy. CP for long context, as a ring (often intra-node or over fast links).
5. The ordering (which axis is “innermost” in the GPU rank layout) is chosen so the chattiest collective lands on the fastest link. Megatron’s rank ordering is essentially TP-innermost, then CP/EP, then PP, then DP-outermost.

How to actually reason about a config

Work memory-first, then comms. (a) Compute per-GPU memory need (params/DP-or-ZeRO shard + optimizer + peak activations). (b) If params don’t fit → add PP (cheap comms) until they do. (c) If a layer’s activations don’t fit, or you want to cut activation memory / raise per-GPU FLOPs → add TP (≤8, NVLink) + SP. (d) MoE experts too big → EP. (e) Sequence too long → CP. (f) Fill remaining GPUs with DP for throughput, sizing microbatches to shrink the bubble. Always check: does each axis’s collective fit on the link you’ve placed it on?

Quiz: You have 512 H100s in 64 nodes (8 GPUs/node, NVLink intra-node, IB inter-node). You pick TP=8, PP=8, DP=8. Which axis goes inside a node, which spans nodes, and why?

Answer

TP=8 goes inside a node. TP does 2 all-reduces per layer (×2 for backward) on the critical path between dependent GEMMs — it needs the ~900 GB/s NVLink domain, and TP=8 exactly fills one 8-GPU node.

PP=8 and DP=8 span nodes (over InfiniBand). PP only sends the boundary activation P2P once per stage crossing — low volume, latency-tolerant. DP/ZeRO syncs gradients once per step, overlapped with backward — the most latency-tolerant axis of all. Both are fine on the ~50 GB/s IB link.

Layout: each node hosts one TP group (the 8 GPUs); the 64 nodes are organized into an 8×8 grid of (PP stage × DP replica). $8\times 8\times 8=512$. ✓

---

9. Practice problems

Practice problem 1: Fit a 70B dense model on 32 H100s (80 GB each)

Problem. A 70B-parameter dense transformer, bf16 training with the Adam optimizer. Rough memory model: parameters 2 B/param, gradients 2 B/param, Adam optimizer state 12 B/param (fp32 momentum+variance+master weights). Ignore activations for a first pass. You have 32× H100 (80 GB). Can plain data parallelism work? Design a parallel plan.

Worked solution.

• State per full replica: $70\text{B}\times (2+2+12)=70\text{B}\times 16=1120$ GB. That’s the static state alone — it does not fit on one 80 GB GPU, so pure DP (which replicates all state) is impossible regardless of GPU count.
• Option A — ZeRO/FSDP only (lesson 09). ZeRO-3 shards all state across $N=32$: $1120/32=35$ GB per GPU. Fits with ~45 GB left for activations + the transient all-gathered layer. This can work for a 70B dense model! ZeRO-3 across 32 GPUs is a legitimate answer — no model parallelism needed for state.
• Why still add TP? Activations. At long sequence / large micro-batch, a single layer’s activations (and the transient full-precision all-gathered weights) may not fit in the remaining ~45 GB, and per-GPU FLOPs may be low. Add TP=8 within each node: shards activations 8× and keeps GEMMs large. Then DP/ZeRO=4 across the 4 nodes. $8\times 4=32$. ✓
• Sanity on comms: TP=8 all-reduces ride NVLink (fine). ZeRO reduce-scatter/all-gather over 4 nodes on IB, overlapped with backward (fine).
• Takeaway: for a 70B dense model at 32 GPUs, ZeRO-3 alone handles memory; TP is added to control activation memory and keep MFU high — exactly the lesson-09-vs-lesson-10 distinction (state vs compute).

Practice problem 2: Compute the comms bill of TP vs the bubble of PP

Problem. Transformer layer: hidden $h=8192$, micro-batch $b=1$, sequence $s=8192$, bf16. (a) For TP=8 with ring all-reduce, how many bytes does each GPU move per layer (fwd+bwd)? (b) For PP with $p=8$ stages, what $m$ do you need to hold the bubble under 15%?

Worked solution.

• (a) TP comms. Activation tensor per all-reduce: $bsh=1\cdot 8192\cdot 8192=6.7\times 10^7$ elems $=1.34\times 10^8$ bytes (bf16) $\approx 128$ MB. Ring all-reduce moves $2\frac{t-1}{t}\approx 2\cdot \frac{7}{8}=1.75\times$ that per GPU $=224$ MB per all-reduce. Per layer: 2 fwd + 2 bwd = 4 all-reduces $\Rightarrow \approx 896$ MB per GPU per layer. Over NVLink at ~900 GB/s that’s ~1.0 ms of exposed comms per layer — and it’s on the critical path, which is precisely why this must not cross to 50 GB/s IB (where it would be ~18 ms/layer). ✓
• (b) PP bubble. Want $\frac{p-1}{m+p-1}\le 0.15$ with $p=8$: $\frac{7}{m+7}\le 0.15\Rightarrow m+7\ge 46.7\Rightarrow m\ge 40$. So ≥40 microbatches. If memory or global-batch limits cap you below that, use interleaved 1F1B with $v=2$: bubble $\approx \frac{1}{2}\cdot \frac{7}{m}$, so $m\ge 24$ suffices — at the cost of ~2× the P2P traffic.
• Takeaway: TP buys you small per-GPU tensors but a heavy, critical-path comms bill (NVLink-only); PP buys you cheap P2P comms but a bubble you must amortize with many microbatches or interleaving.

---

10. Practice & resources

Hands-on (do these, don’t just read):

1. Write a two-GPU tensor-parallel MLP from scratch. Implement ColumnParallelLinear and RowParallelLinear with torch.distributed (all_reduce, all_gather) on 2 GPUs (or 2 processes with Gloo/NCCL). Verify the output equals the single-GPU MLP bit-for-bit (up to fp error). Then insert a GeLU and confirm why you can’t swap column/row order. This is the single most clarifying exercise for TP.
2. Run Megatron-LM / Megatron-Core. Clone NVIDIA/Megatron-LM, train a small GPT with --tensor-model-parallel-size, --pipeline-model-parallel-size, --sequence-parallel, and --num-layers-per-virtual-pipeline-stage (interleaving). Watch throughput and memory change as you vary TP/PP/$m$. Read megatron/core/pipeline_parallel/schedules.py to see 1F1B and interleaving in code.¹⁴
3. Instrument the bubble. With PP on, profile with Nsight Systems and see the fill/drain idle in the timeline. Vary $m$ and confirm the bubble fraction matches $\frac{p-1}{m+p-1}$.
4. Compute a parallelism plan. Pick a target (e.g. train Llama-3-70B, or DeepSeek-V3-scale MoE) on a stated cluster (e.g. 256× H100, 8/node, NVLink + 400G IB). Write out DP×TP×PP(×EP×CP), justify each axis by memory-then-comms, and place each on the topology. Sanity-check against a published config.

Reading (paper → docs → code):

• TP: Megatron-LM, arXiv:1909.08053² — read Section 3 (the MLP + attention parallelism).
• SP + selective recompute: Korthikanti et al., arXiv:2205.05198¹.
• PP schedules: GPipe arXiv:1811.06965⁶; PipeDream/1F1B arXiv:1806.03377⁷; interleaved 1F1B + the $\frac{p-1}{m+p-1}$ analysis in the Megatron scaling paper arXiv:2104.04473³; zero-bubble arXiv:2401.10241⁸.
• CP: Ring Attention arXiv:2310.01889⁹; Striped Attention arXiv:2311.09431¹⁰.
• EP / MoE: the original sparsely-gated MoE arXiv:1701.06538¹³; DeepSeek-V3 arXiv:2412.19437¹¹; Megatron-Core MoE technical report + docs¹².
• Docs: the NVIDIA Megatron-Bridge / Megatron-Core parallelism guides are the best current (2025–26) practical reference for how these compose and are configured.⁴

---

11. What’s next

Three threads, in order:

1. 11-sota-large-scale-training — putting all of this together at frontier scale: how a real training run (Llama-3, DeepSeek-V3) picks its DP×TP×PP×EP×CP config, overlaps everything, handles fault tolerance and MFU, and squeezes the last 10% out of the interconnect. This lesson is the toolbox; lesson 11 is the finished machine.
2. Back to 09-zero-and-fsdp — re-read with the state-vs-compute distinction sharp. ZeRO-3 and TP are complementary: ZeRO shards the state you replicate across DP; TP shards the compute within a node. Frontier stacks run both at once (the “ZeRO + 3D parallelism” combination).
3. Concept stubs to write: tensor-parallelism, pipeline-parallelism, sequence-parallelism, expert-parallelism, context-parallelism, 3d-parallelism.

---

Part of the gpu-systems-for-llms curriculum. Builds on 08-data-parallelism-and-collectives and 09-zero-and-fsdp; sets up 11-sota-large-scale-training. Filed: 2026-09-02

---

References

Footnotes

1. Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models,” arXiv:2205.05198 (2022). https://arxiv.org/abs/2205.05198 — sequence parallelism: split LayerNorm/dropout/residual along the sequence dim, interleave with TP via all-gather ($f$) / reduce-scatter ($g$); same total comms volume as plain-TP all-reduce, with $t\times$ activation-memory reduction on the non-TP regions. Also selective activation recomputation. [established] ↩ ↩² ↩³ ↩⁴ ↩⁵

2. Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,” arXiv:1909.08053 (2019). https://arxiv.org/abs/1909.08053 — introduces tensor (intra-layer) parallelism: column-parallel then row-parallel GEMM pairing for the MLP (so a single all-reduce per block), head-split multi-head attention, and the conjugate operators $f$/$g$. [established] ↩ ↩² ↩³ ↩⁴ ↩⁵

3. Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM,” arXiv:2104.04473 / SC’21 (2021). https://arxiv.org/abs/2104.04473 — 3D parallelism (DP×TP×PP), the pipeline bubble fraction $(p-1)/(m+p-1)$, interleaved 1F1B reducing the bubble by $\sim v$ at $v\times$ P2P cost, and the “TP within a node, PP across nodes” topology guidance. [established] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹ ↩¹⁰

4. NVIDIA, “Parallelisms Guide” (Megatron-Bridge / Megatron-Core), https://docs.nvidia.com/nemo/megatron-bridge/latest/parallelisms.html and the Megatron-Core parallelism guide. Current (2025–26) practical guidance: enable sequence parallelism whenever TP>1; keep TP within a node up to TP=8 (NVLink), PP across nodes (IB); use context parallelism for long-context training; TP requires hidden-size and num-attention-heads divisible by the TP degree. [recent] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶

5. NVIDIA H100 / DGX H100 specs: fourth-gen NVLink ~900 GB/s bidirectional per GPU via NVSwitch; typical node-to-node fabric is 8× 400 Gb/s InfiniBand (NDR). NVIDIA H100 datasheet, https://resources.nvidia.com/en-us-data-center-overview/hpc-datasheet-sc23-h100 and DGX H100 datasheet. [established] ↩

6. Huang et al., “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism,” arXiv:1811.06965 (2018). https://arxiv.org/abs/1811.06965 — micro-batching to fill the pipeline; all-forwards-then-all-backwards schedule ($O(m)$ activation memory). [established] ↩ ↩²

7. Narayanan et al., “PipeDream: Generalized Pipeline Parallelism for DNN Training,” arXiv:1806.03377 / SOSP’19 (2018). https://arxiv.org/abs/1806.03377 — the 1F1B schedule; reordering backward earlier drops peak activation memory from $O(m)$ to $O(p)$ at the same bubble fraction as GPipe. [established] ↩ ↩²

8. Qi et al., “Zero Bubble Pipeline Parallelism,” arXiv:2401.10241 / ICLR 2024. https://arxiv.org/abs/2401.10241 — splits backward into input-grad (B) and weight-grad (W); W has no downstream dependency and fills the bubble, approaching zero bubble under a synchronous schedule. [recent] ↩ ↩²

9. Liu, Zaharia, Abbeel, “Ring Attention with Blockwise Transformers for Near-Infinite Context,” arXiv:2310.01889 (2023). https://arxiv.org/abs/2310.01889 — context parallelism: distribute the sequence across a ring of devices; each passes its KV block to the next while computing blockwise attention (online softmax), fully overlapping KV communication with compute; enables $N\times$ longer context with no added overhead. [established] ↩ ↩² ↩³ ↩⁴

10. Brandon et al., “Striped Attention: Faster Ring Attention for Causal Transformers,” arXiv:2311.09431 (2023). https://arxiv.org/abs/2311.09431 — permutes tokens so each ring GPU gets a balanced early/late position mix, fixing the load imbalance Ring Attention suffers under causal masking. [recent] ↩ ↩²

11. DeepSeek-AI, “DeepSeek-V3 Technical Report,” arXiv:2412.19437 (2024). https://arxiv.org/abs/2412.19437 — fine-grained MoE with many small experts + shared experts, and an auxiliary-loss-free load-balancing strategy (bias adjustment) to avoid the aux-loss quality tax; a canonical large-scale EP training system. [recent] ↩ ↩² ↩³

12. “Scalable Training of Mixture-of-Experts Models with Megatron Core” (technical report, 2026, arXiv:2603.07685, https://arxiv.org/abs/2603.07685) and the Megatron-Core MoE docs (https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/moe.html). Four-stage MoE forward (route → dispatch → compute → combine); all-to-all volume scales with tokens×hidden and pushes cross-node at large EP; comms can exceed 50% of DeepSeek-V3-scale training time without overlap; DeepEP/HybridEP kernels, shared-expert overlap, 1F1B A2A overlap; load balancing via aux loss, capacity-factor token dropping, and shifting EP↔expert-TP. See also NVIDIA, “Optimizing Communication for Mixture-of-Experts Training with Hybrid Expert Parallel” (https://developer.nvidia.com/blog/optimizing-communication-for-mixture-of-experts-training-with-hybrid-expert-parallel/). [recent] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶

13. Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer,” arXiv:1701.06538 (2017). https://arxiv.org/abs/1701.06538 — top-k routing with an auxiliary load-balancing loss; the foundation of modern MoE. [established] ↩ ↩²

14. NVIDIA/Megatron-LM, megatron/core/pipeline_parallel/schedules.py. https://github.com/NVIDIA/Megatron-LM — reference implementation of GPipe, 1F1B, and interleaved 1F1B schedules; tensor_parallel/ holds ColumnParallelLinear / RowParallelLinear. ↩