Tensor and pipeline parallelism, and choosing a layout

Before this lesson: lesson 6 (model states and activation memory), lesson 7 (all-reduce cost: each GPU sends $2\frac{t-1}{t}n$ bytes), lesson 8 (data parallelism, ZeRO, and the overlap condition $T \ge F/W$); block matrix products (linear algebra 1); the transformer block (attention and MLP). Time: about 85 minutes, plus the notebook.

---

Why this matters #

ZeRO-3 can shard the model states of any model, but data parallelism has two limits (lesson 8). Each GPU still runs whole layers on its own micro-batch, so the activations are not divided. And each data-parallel rank needs about $F/W$ tokens to hide its communication, so at a fixed global batch you cannot add GPUs forever. Tensor parallelism (TP) and pipeline parallelism (PP) split the model itself.

Large LLM runs combine all three. Llama 3 405B was trained with TP = 8 inside each 8-GPU server, PP = 16 and DP = 128, on 16,384 H100 GPUs (Llama 3 paper, 2024, §3.3.2 and Table 4). Megatron-LM, NeMo, DeepSpeed and torchtitan implement these methods, and vLLM uses tensor parallelism to serve models that do not fit on one GPU. Choosing the degrees is a standard design and interview question; the answer is FLOPs, bytes and link bandwidths.

1. Predict first #

2. The picture #

Notation for this lesson. $t$: TP degree (GPUs that share one layer). $p$: number of pipeline stages (in lesson 7, $p$ was the number of ranks of a collective; here it is the number of stages). $N_d$: DP degree. $m$: micro-batches per step. $v$: model chunks per GPU (interleaving). $b$: sequences per micro-batch, $s$: sequence length, $T = bs$: tokens per micro-batch. $h$: hidden size, $a$: attention heads, $L$: layers. $F$: FLOP/s one GPU achieves on the matmuls. $W$: bandwidth per GPU and direction, bytes/s. $t_f$, $t_b$: time of the forward and backward pass of one micro-batch on one stage. Activations are bf16 ($e = 2$ bytes).

One matmul on two GPUs #

Take $Y = XA$ with $X$ of size $T \times h$ and $A$ of size $h \times k$. There are two ways to split $A$ (block products, linear algebra 1):

• Column split $A = [A_1, A_2]$: GPU $i$ computes $XA_i$, which is the $i$-th block of columns of $Y$. Each GPU needs all of $X$ and sends nothing. The output is split by columns.
• Row split $A = \begin{bmatrix} A_1 \\ A_2 \end{bmatrix}$, with $X = [X_1, X_2]$ split by columns to match: GPU $i$ computes $X_iA_i$, a full-size $T \times k$ partial sum, and $Y = X_1A_1 + X_2A_2$. An all-reduce adds the two.

Chain the two splits. A column split produces an output split by columns, and a row split wants exactly that as its input. So "column split, then row split" needs no communication in the middle and one all-reduce at the end. An elementwise function such as GeLU can sit in the middle, because it acts on each column alone. This is the whole idea of Megatron tensor parallelism (Shoeybi et al. 2019).

Interactive figure: loading…

A pipeline is an assembly line #

Put layers 1–20 on GPU 0, layers 21–40 on GPU 1, and so on. If the whole batch flows through as one piece, only one GPU works at a time. Cut the batch into $m$ micro-batches; then the stages work at the same time on different micro-batches. But the line must fill at the start and drain at the end. The forward phase with $p = 3$ stages and $m = 4$ micro-batches (F0 is the forward pass of micro-batch 0, F1 that of micro-batch 1, and so on):

time slot	1	2	3	4	5	6
stage 0	F0	F1	F2	F3	·	·
stage 1	·	F0	F1	F2	F3	·
stage 2	·	·	F0	F1	F2	F3

The forward phase takes $m + p - 1 = 6$ slots, and each stage works in only 4 of them. The dots are the bubble: idle time. The backward phase is the mirror image.

Interactive figure: loading…

Why the two methods use different links. TP all-reduces activations twice per layer in the forward pass, and the next layer waits for the result. PP sends one activation tensor per micro-batch, point to point, to the next stage only. So TP needs the fastest links, and PP works well over the network.

3. Definitions and main results #

The Megatron MLP split #

Proposition (Megatron MLP split). Let $Z = \mathrm{GeLU}(XA)\,B$ with $X \in \mathbb{R}^{T \times h}$, $A \in \mathbb{R}^{h \times 4h}$ and $B \in \mathbb{R}^{4h \times h}$. Split $A = [A_1, \dots, A_t]$ by columns and $B$ by rows into $B_1, \dots, B_t$, in blocks of $4h/t$. Then the forward pass needs one all-reduce, the backward pass one all-reduce, and the weight gradients none.

Proof.

1. Block product: $XA = [XA_1, \dots, XA_t]$. Write $U_i = XA_i$.
2. GeLU acts on each element, so $\mathrm{GeLU}(XA) = [Y_1, \dots, Y_t]$ with $Y_i = \mathrm{GeLU}(U_i)$. GPU $i$ computes $Y_i$ alone.
3. Column blocks of $Y$ meet row blocks of $B$: $Z = \sum_i Y_iB_i$. GPU $i$ computes $Y_iB_i$, and one all-reduce gives $Z$ on every GPU. Megatron calls this operator $g$: all-reduce in the forward pass, identity in the backward pass.
4. Backward pass, with $\ell$ the loss and $G = \partial\ell/\partial Z$ on every GPU: $\partial\ell/\partial B_i = Y_i^\top G$, $\;\partial\ell/\partial Y_i = GB_i^\top$, $\;\partial\ell/\partial U_i = \partial\ell/\partial Y_i \odot \mathrm{GeLU}'(U_i)$ and $\partial\ell/\partial A_i = X^\top \partial\ell/\partial U_i$. All of them use only local data.
5. $X$ feeds all $t$ blocks, so $\partial\ell/\partial X = \sum_i (\partial\ell/\partial U_i)\,A_i^\top$: one all-reduce. This is operator $f$: identity in the forward pass, all-reduce in the backward pass. $\blacksquare$

$f$ and $g$ are conjugate: each is the identity in one direction and an all-reduce in the other. The reason is the chain rule. When one tensor is copied to $t$ GPUs (forward, $f$), its gradient is the sum of the $t$ gradients (backward). When $t$ partial results are summed (forward, $g$), each of them receives the same gradient (backward).

Why not split $A$ by rows? Then GPU $i$ holds $X_iA_i$, a partial sum, and $\mathrm{GeLU}(\sum_i X_iA_i) \ne \sum_i \mathrm{GeLU}(X_iA_i)$ because GeLU is not linear. The sum must be complete before the GeLU: an extra all-reduce, on a $T \times 4h$ tensor, 4× larger than $Z$.

Attention is split by heads. $W_Q$, $W_K$ and $W_V$ are column-parallel with $a/t$ heads per GPU (so $t$ must divide $a$; with grouped-query attention, the key/value heads are split too, or copied when there are fewer of them than $t$). Attention runs locally for each head, and $W_O$ is row-parallel. The same $f$ and $g$ appear. So per layer: 2 all-reduces in the forward pass and 2 in the backward pass, each on $T \cdot h$ values. Activation recomputation repeats the forward ones. The embedding is split along the vocabulary, and the output logits are fused with the cross-entropy, so only $b \times s$ numbers are communicated instead of $b \times s \times V$, with $V$ the vocabulary size (Megatron-LM §3).

What TP divides. The weights, gradients and optimizer states of the blocks are divided by $t$, and so are the activations inside the parallel regions. The LayerNorm and dropout regions stay replicated: the activation memory per layer is $sbh(10 + 24/t + 5as/(ht))$ bytes (lesson 6). Sequence parallelism (lesson 10) shards those regions too.

The TP communication-to-compute ratio #

Proposition (TP ratio). Per layer and per forward pass, ignoring the FLOPs of the attention scores,

Proof.

1. Matmul FLOPs per layer (2 FLOPs per multiply-add): QKV $2Th \cdot 3h = 6Th^2$, output projection $2Th^2$, MLP $2 \cdot 2Th \cdot 4h = 16Th^2$. Total $24Th^2$, so $24Th^2/t$ per GPU and time $24Th^2/(tF)$.
2. Two all-reduces of $2Th$ bytes each (bf16). A ring all-reduce of $n$ bytes sends $2\frac{t-1}{t}n$ bytes per GPU (lesson 7): $2 \cdot 2\frac{t-1}{t} \cdot 2Th = 8Th\frac{t-1}{t}$ bytes, time $8Th(t-1)/(tW)$.
3. Divide: $\dfrac{8Th(t-1)/(tW)}{24Th^2/(tF)} = \dfrac{(t-1)F}{3hW}$. $\blacksquare$

What the formula says:

• $T$ cancels. More tokens per micro-batch do not help: compute and bytes both grow like $T$.
• $t$ hurts. The compute per GPU falls as $1/t$, but the bytes per GPU stay near $2n$.
• $h$ helps. Compute grows as $h^2$, bytes only as $h$. Wide models tolerate TP better.
• The backward pass has twice the FLOPs and the same two all-reduces: half the ratio. Megatron overlaps the backward all-reduce of $\partial\ell/\partial X$ with the weight-gradient matmul, which does not need it. The forward all-reduces sit on the critical path.
• Narayanan et al. (2021, §3.2) count $8bsh\frac{t-1}{t}$ elements per layer for forward plus backward. In bf16, the forward alone is $8bsh\frac{t-1}{t}$ bytes. Same expression, different quantity.

Numbers (H100 SXM, $t = 8$, $F = 989$ TFLOP/s, the dense bf16 peak, as a worst case):

$h$	NVLink, $W = 450$ GB/s	InfiniBand, $W = 50$ GB/s
8192 (70B class)	0.63	5.6
4096 (7B class)	1.25	11.3

Across nodes, TP spends 5–11× more time communicating than computing. Inside a node it is affordable for wide layers. For a 7B model, $t = 8$ is already too much: use $t \le 4$, or no TP.

Interactive figure: loading…

Pipeline schedules #

A stage holds $L/p$ consecutive layers. Per micro-batch, each stage boundary carries $T \cdot h$ values forward (activations) and $T \cdot h$ values backward (their gradients), point to point. All three schedules below flush at the end of the step: every micro-batch finishes its backward pass before the optimizer step, so the optimizer gets exactly the gradient of the whole batch.

• GPipe (Huang et al. 2019): all $m$ forward passes, then all $m$ backward passes.
• 1F1B (PipeDream-Flush, Narayanan et al. 2021; Megatron's default): stage $s$ (counted from 0) runs $\min(p - s - 1, m)$ warm-up forward passes, then alternates one forward and one backward pass, then runs the remaining backward passes.
• Interleaved 1F1B (Narayanan et al. 2021): each GPU holds $v$ chunks of $L/(pv)$ layers that are not adjacent (GPU 0 holds chunks $0, p, 2p, \dots$), so the pipeline has $pv$ virtual stages. The schedule needs $m$ to be a multiple of $p$.

Theorem (pipeline bubble). Assume $p$ equal stages, times $t_f$ and $t_b$ for the forward and backward pass of one micro-batch on one stage, one operation at a time per GPU, no communication time, and a flush at the end of every step.

1. Any schedule with one contiguous stage per GPU takes at least $(m + p - 1)(t_f + t_b)$. GPipe and 1F1B reach this bound. So their bubble is $\frac{p-1}{m+p-1}$ of the total time, which is $\frac{p-1}{m}$ of the ideal time $m(t_f + t_b)$ (the form Narayanan et al. use).
2. Stage $s$ stores the activations of $m$ micro-batches in GPipe, and of at most $p - s$ in 1F1B.
3. With $v$ chunks per GPU and the interleaved schedule ($m$ a multiple of $p$), the bubble time falls to $(p-1)(t_f + t_b)/v$, a fraction $\frac{p-1}{vm + p - 1}$ of the total. The point-to-point traffic grows about $v$-fold.

Proof. Write $F(s, i)$ and $B(s, i)$ for the forward and backward pass of micro-batch $i$ on stage $s$, both counted from 0.

1. Lower bound. The last stage cannot start before micro-batch 0 has passed the other $p - 1$ stages: not before $(p-1)t_f$. Its last operation is some $B(p-1, j)$, and after it $B(p-2, j), \dots, B(0, j)$ still run: $(p-1)t_b$ more. In between, the last stage does $m(t_f + t_b)$ of work. So the step takes at least $(p-1)t_f + m(t_f + t_b) + (p-1)t_b$.
2. GPipe reaches it. $F(s, i)$ waits for $F(s-1, i)$ and for $F(s, i-1)$. By induction on $s + i$, it ends at $(s + i + 1)t_f$, because both predecessors end at $(s + i)t_f$. The last forward pass ends at $(m + p - 1)t_f$. The backward phase is the mirror image: $(m + p - 1)t_b$ more.
3. 1F1B reaches it too. One checks by induction over the operations of each stage that $B(s, i)$ ends at $(p-1)t_f + (i+1)(t_f + t_b) + (p - 1 - s)\,t_b$. For $s = 0$, $i = m - 1$ this is $(m + p - 1)(t_f + t_b)$. In words: once the last stage starts, it alternates F and B without waiting, because the stage before it delivers the next forward pass no later than the moment the last stage finishes a backward pass. (The simulator in the figure and in the notebook confirms the formula, also for $t_f > t_b$ and for $m < p$.)
4. Bubble. Each GPU works $m(t_f + t_b)$ out of $(m + p - 1)(t_f + t_b)$: an idle fraction $\frac{p-1}{m+p-1}$.
5. Memory. A micro-batch's activations stay on a stage from its forward to its backward pass. GPipe runs all $m$ forward passes first: $m$ in flight. In 1F1B, after $p - s - 1$ warm-up forward passes each new forward pass is followed by a backward pass, so at most $p - s$ are in flight, whatever $m$ is.
6. Interleaving (a sketch, not a full proof). A chunk's forward and backward pass take $t_f/v$ and $t_b/v$. In this schedule the ramp-up and the ramp-down of step 1 cost $p - 1$ chunk steps each, so the bubble time is $(p-1)(t_f + t_b)/v$; a simulation of the schedule gives exactly this total. The total time is $m(t_f + t_b) + (p-1)(t_f + t_b)/v$, and the fraction is $\frac{p-1}{vm + p - 1}$. Each micro-batch now crosses $pv - 1$ stage boundaries instead of $p - 1$. $\blacksquare$

Consequences:

• More micro-batches, smaller bubble. With $m \ge 4p$ the bubble is below 20%, because $(p-1)/(5p-1) < 1/5$. 1F1B keeps the memory bounded while $m$ grows; GPipe does not.
• 1F1B saves memory, not time. It has the same bubble as GPipe (point 1), but each stage keeps at most $p$ micro-batches instead of $m$ (point 2).
• The bubble is tied to the batch. With a global batch of $B$ sequences and DP degree $N_d$, $m = B/(bN_d)$. More DP at a fixed global batch means fewer micro-batches and a larger bubble.
• Unequal stages. The slowest stage sets the pace. The embedding and the output layer (with a large vocabulary) make the first and the last stages heavier, so frameworks give them fewer layers.

Placing TP, PP and DP on a cluster #

Compare the communication per unit of compute of each method:

Method	What moves	When	Comm / compute
TP	all-reduce of activations, 4 per layer	every layer, on the critical path	$(t-1)F/(3hW)$: 0.63 on NVLink, 5.6 on InfiniBand
PP	activations to the next stage	once per micro-batch and boundary	$pF/(12hLW)$: 0.02 on InfiniBand
DP	gradients (DDP) or weights and gradients (ZeRO-3)	once per step (ZeRO-3: per micro-batch); overlapped	hidden if each DP rank has $T \ge F/W$ tokens

The PP ratio. Per micro-batch, one GPU of a stage computes $24Th^2(L/p)/t$ forward FLOPs. With TP inside the stage, all $t$ GPUs hold the same output tensor, so each sends only $1/t$ of it over its own network card, and the receivers all-gather the pieces over NVLink (the scatter/gather optimization, Narayanan et al. §4.1): $2Th/t$ bytes per GPU. The ratio is $\frac{2Th/(tW)}{24Th^2L/(ptF)} = \frac{pF}{12hLW}$. The table uses $t = 8$, $p = 8$, $L = 80$, $h = 8192$ and $F = 989$ TFLOP/s.

The three rules (Narayanan et al. 2021, takeaways 1 and 2; the Llama 3 order [TP, CP, PP, DP]):

1. TP inside the NVLink domain: $t \le$ GPUs per node (8 on a DGX H100, 4 on a GH200 node with 4 GPUs), and smaller for narrow models.
2. PP across nodes, with $t \cdot p$ just large enough for the model to fit (including the $p$ micro-batches in flight of 1F1B), and $m \ge 4p$.
3. DP for the rest, outermost. It communicates least often and overlaps with compute, so it tolerates the slowest links.

The order [TP, CP, PP, DP] lists the dimensions from the fastest links to the slowest (CP is context parallelism, lesson 10): consecutive GPU ranks, which share a node, form the TP groups.

Interactive figure: loading…

Beyond 1F1B and 1D tensor parallelism

• Zero-bubble schedules (Qi et al., ICLR 2024) split the backward pass into the input-gradient part, which the previous stage waits for, and the weight-gradient part, which can wait. Scheduling the weight-gradient parts into the idle slots removes most of the bubble under synchronous training; the paper also reports schedules with no bubble at all.
• DualPipe (DeepSeek-V3 technical report, 2024) feeds micro-batches from both ends of the pipeline and overlaps computation with communication.
• 2D and 3D tensor parallelism (Optimus, Xu et al. 2021; Bian et al. 2021) split each weight matrix along two or three dimensions of a grid of GPUs, so that activations are also partitioned and the communication per GPU grows more slowly with the number of GPUs than in Megatron's 1D split.

Answers to section 1. (1) One: a column split for $A$, a row split for $B$, and one all-reduce at the end. (2) $\frac{p-1}{m+p-1} = 3/7 \approx 43\%$ of the step time. 1F1B has the same bubble; it saves memory, not time.

4. Worked example #

(a) TP on one layer #

$h = 8192$, one sequence of $s = 4096$ tokens ($T = 4096$), $t = 8$, $F = 989$ TFLOP/s.

1. FLOPs per GPU and layer in the forward pass: $24 \cdot 4096 \cdot 8192^2 / 8 = 8.25 \times 10^{11}$. Time: $8.25 \times 10^{11} / 989 \times 10^{12} = 0.83$ ms.
2. Bytes per GPU: $8 \cdot 4096 \cdot 8192 \cdot \frac{7}{8} = 2.35 \times 10^8$ bytes $= 235$ MB.
3. Time on NVLink: $235\text{ MB}/450\text{ GB/s} = 0.52$ ms, ratio $0.52/0.83 = 0.63$. On InfiniBand: $235\text{ MB}/50\text{ GB/s} = 4.7$ ms, ratio 5.6.

(b) A pipeline with $p = 4$, $m = 12$, $t_f = 1$, $t_b = 2$ #

• GPipe: $(12 + 3) \cdot 3 = 45$ time units, against the ideal $12 \cdot 3 = 36$. Bubble: $9/45 = 20\%$ of the total ($9/36 = 25\%$ of the ideal time). Every stage stores 12 micro-batches.
• 1F1B: also 45 units and 20%. Stage 0 stores 4 micro-batches, stage 3 only 1.
• Interleaved, $v = 3$ (12 is a multiple of 4): $36 + 9/3 = 39$ units, bubble $3/39 \approx 7.7\%$. Each micro-batch crosses $pv - 1 = 11$ boundaries in each direction instead of 3. GPU 0 stores a little more than in 1F1B: up to 15 chunk forward passes of $1/3$ stage each, the activations of 5 micro-batches.

Set the pipeline figure to $p = 4$, $m = 12$, $t_b/t_f = 2$ to see these three schedules.

(c) A layout for a 70B model on 64 H100s #

Given. Model: $\Psi = 70 \times 10^9$ parameters, $L = 80$, $h = 8192$, $a = 64$. Cluster: 8 nodes × 8 H100 80 GB; NVLink 450 GB/s and InfiniBand 50 GB/s per GPU. Global batch: 256 sequences of 4,096 tokens; micro-batch $b = 1$. For whole-step estimates we use an achieved $F = 400$ TFLOP/s.

1. Model states: $16\Psi = 1{,}120$ GB with mixed-precision Adam (lesson 6): the memory of 14 GPUs before any activation. So some model splitting is needed.
2. TP = 8 inside each node. Ratio at $F = 400$ TFLOP/s: $7 \cdot 400/(3 \cdot 8192 \cdot 0.45) = 0.25$ on NVLink. Across nodes it would be 2.3.
3. Activations per layer with TP = 8, FlashAttention and no sequence parallelism: $sbh(10 + 24/t) = 13\,sbh = 13 \cdot 4096 \cdot 8192$ bytes $= 0.44$ GB.
4. Across the 8 nodes, compare PP = 8 with ZeRO-3 DP = 8 (TP = 8 in both):

	TP 8 × PP 8, 1F1B	TP 8 × ZeRO-3 DP 8
model states per GPU	$16\Psi/64 = 17.5$ GB	$16\Psi/64 = 17.5$ GB
activations per GPU	10 layers × 8 micro-batches × 0.44 GB = 34.9 GB	80 layers × 0.44 GB = 34.9 GB
network bytes per micro-batch	$2sbh/t = 8.4$ MB per boundary and direction: 0.17 ms	$3 \cdot 2 \cdot \frac{\Psi}{8} \cdot \frac{7}{8} = 45.9$ GB: 0.92 s
compute per micro-batch	$6\Psi \cdot 4096/64/F = 67$ ms	$6\Psi \cdot 4096/8/F = 0.54$ s
verdict	bubble $7/263 = 2.7\%$; network traffic negligible	communication-bound: 4,096 tokens $< F/W = 8{,}000$

5. Decision. TP 8 × PP 8 needs about 52 GB per GPU plus buffers, and the network is almost idle. With GPipe instead of 1F1B, stage 0 would store all 256 micro-batches: about 1,117 GB. TP 8 × PP 4 × DP 2 (ZeRO-1) also works: bubble $3/131 = 2.3\%$, 21.9 GB of model states, 56.8 GB in total, and a 4.4 GB gradient all-reduce per step that hides easily (DP ratio 0.01). Both are good choices; the second leaves less memory free but uses fewer pipeline stages. TP 8 × DP 8 with ZeRO-2 (no PP) would have no bubble, but it needs about 68 GB per GPU before logits and buffers: too tight.

5. In code #

The tensor-parallel MLP with NumPy, forward and backward written by hand. Each loop iteration is one "GPU"; we check which sums are needed.

import numpy as np

C = np.sqrt(2 / np.pi)
def gelu(x):                                 # tanh form of GeLU
    return 0.5 * x * (1 + np.tanh(C * (x + 0.044715 * x**3)))
def dgelu(x):                                # its derivative
    u = np.tanh(C * (x + 0.044715 * x**3))
    return 0.5 * (1 + u) + 0.5 * x * (1 - u**2) * C * (1 + 3 * 0.044715 * x**2)

rng = np.random.default_rng(0)
T, h, t = 4, 8, 2                            # 4 tokens, width 8, two "GPUs"
X = rng.standard_normal((T, h))              # input
G = rng.standard_normal((T, h))              # upstream gradient dl/dZ
A = rng.standard_normal((h, 4 * h))
B = rng.standard_normal((4 * h, h))

U = X @ A                                    # one GPU, full weights
Z = gelu(U) @ B
dX = ((G @ B.T) * dgelu(U)) @ A.T

Z_parts, dX_parts = [], []
A_cols, B_rows = np.split(A, t, axis=1), np.split(B, t, axis=0)
for Ai, Bi in zip(A_cols, B_rows):           # GPU i: column block, row block
    Ui = X @ Ai                              # local, no communication
    Z_parts.append(gelu(Ui) @ Bi)            # a full-size partial sum
    dX_parts.append(((G @ Bi.T) * dgelu(Ui)) @ Ai.T)   # a part of dl/dX

print(np.allclose(sum(Z_parts), Z))     # True: forward needs one sum (g)
print(np.allclose(dX_parts[0], dX))     # False: one GPU has only a part
print(np.allclose(sum(dX_parts), dX))   # True: backward needs one sum (f)

F = 989e12                              # H100 dense bf16 peak, FLOP/s
for hh in (4096, 8192):                 # t = 8: (t-1)F/(3hW), NVLink and IB
    print(hh, [round(7 * F / (3 * hh * W), 2) for W in (450e9, 50e9)])
# 4096 [1.25, 11.27]
# 8192 [0.63, 5.63]

A pipeline simulator: every stage runs its operations in order; an operation starts when the stage is free and its inputs are ready. It confirms the theorem for worked example (b).

def order(kind, p, m, s):
    """Operations of stage s in order: ("F", i) or ("B", i), i = micro-batch."""
    if kind == "gpipe":
        return [("F", i) for i in range(m)] + [("B", i) for i in range(m)]
    w = min(p - s - 1, m)                    # 1F1B: warm-up forwards
    ops = [("F", i) for i in range(w)]
    for i in range(m - w):
        ops += [("F", w + i), ("B", i)]
    return ops + [("B", i) for i in range(m - w, m)]

def simulate(kind, p, m, tf, tb):
    end, free, todo = {}, [0.0] * p, [order(kind, p, m, s) for s in range(p)]
    in_flight, peak = [0] * p, [0] * p
    while any(todo):
        for s in range(p):
            if not todo[s]:
                continue
            op, i = todo[s][0]
            deps = [("F", s - 1, i)] if op == "F" and s > 0 else []
            if op == "B":
                deps = [("F", s, i)] + ([("B", s + 1, i)] if s < p - 1 else [])
            if all(d in end for d in deps):
                start = max([free[s]] + [end[d] for d in deps])
                free[s] = end[(op, s, i)] = start + (tf if op == "F" else tb)
                in_flight[s] += 1 if op == "F" else -1
                peak[s] = max(peak[s], in_flight[s])
                todo[s].pop(0)
    return max(free), peak

for kind in ("gpipe", "1f1b"):
    total, peak = simulate(kind, p=4, m=12, tf=1, tb=2)
    print(kind, total, (total - 12 * 3) / total, peak)
# gpipe 45.0 0.2 [12, 12, 12, 12]
# 1f1b 45.0 0.2 [4, 3, 2, 1]

In PyTorch, torch.distributed.tensor.parallel provides ColwiseParallel and RowwiseParallel for these splits, and torch.distributed.pipelining provides GPipe, 1F1B and interleaved schedules. The notebook builds both from NumPy instead.

6. Where it shows up #

• LLM pretraining. Llama 3 405B: TP 8, PP 16, DP 128 on 16,384 H100s at about 400 TFLOP/s per GPU (41% of the bf16 peak), with the parallel dimensions ordered [TP, CP, PP, DP] from the fastest links to the slowest. Narayanan et al. (2021) ran a 1-trillion-parameter GPT model on 3,072 A100s with TP 8 and PP 64, at 502 petaFLOP/s in total.
• Newer pipeline schedules. Zero-bubble schedules and DualPipe (see the deeper box) fill the bubble with weight-gradient work or with a second stream of micro-batches.
• Inference. TP inside a node is the standard way to serve a model that does not fit on one GPU (vLLM's tensor_parallel_size). At decode time the all-reduced tensors are small, so latency, not bandwidth, sets their cost (lesson 12).
• ML for science. ORBIT, a 113B-parameter vision transformer for Earth system prediction, combined tensor parallelism with sharded data parallelism on the Frontier supercomputer (Wang et al. 2024). On a GH200 node with 4 GPUs, $t \le 4$. For models of a few billion parameters with very long token sequences (global weather grids, genomes), sequence and context parallelism (lesson 10) usually split the activations better than TP or PP.

7. Common mistakes #

• TP across nodes. Its per-layer all-reduces need NVLink-class bandwidth: the ratio is 5.6 on InfiniBand for $h = 8192$.
• Splitting the first MLP weight by rows. It adds an all-reduce of a $T \times 4h$ tensor before the GeLU.
• "TP communication shrinks with $t$." The bytes per GPU stay near $2n$, while the compute per GPU falls as $1/t$.
• Head counts. $t$ must divide the number of attention heads; with GQA, the key/value heads must be split or copied.
• "1F1B has a smaller bubble than GPipe." It has the same bubble; 1F1B saves activation memory.
• Mixing the two bubble definitions. $(p-1)/(m+p-1)$ of the total time or $(p-1)/m$ of the ideal time: say which one you use.
• Forgetting $m = B/(bN_d)$. Adding DP ranks at a fixed global batch shrinks $m$ and grows the bubble.

8. Check yourself #

1. $h = 4096$, $t = 4$, $F = 989$ TFLOP/s. What is the forward communication-to-compute ratio on NVLink (450 GB/s) and on InfiniBand (50 GB/s)?
   Answer

   $(t-1)F/(3hW) = 3 \cdot 989 \times 10^{12}/(3 \cdot 4096 \cdot 450 \times 10^9) \approx 0.54$ on NVLink and $\approx 4.8$ on InfiniBand.
2. In the Megatron MLP, why is the backward pass of the column-parallel layer an all-reduce, and which gradients need no communication?
   Answer

   $X$ is used by every column block, so $\partial\ell/\partial X = \sum_i (\partial\ell/\partial U_i)A_i^\top$ is a sum over GPUs (operator $f$). The weight gradients $Y_i^\top G$ and $X^\top \partial\ell/\partial U_i$ use only local data, so they need no communication.
3. $p = 6$, $m = 18$, $t_f = 1$, $t_b = 2$. What is the bubble of GPipe, of 1F1B and of interleaving with $v = 2$? How many micro-batches does stage 0 store in GPipe and in 1F1B?
   Answer

   GPipe and 1F1B: $5/23 \approx 21.7\%$ (69 time units against 54). Interleaved: $5/41 \approx 12.2\%$ (61.5 units). Stage 0 stores 18 micro-batches in GPipe and 6 in 1F1B.
4. Why can no schedule with one contiguous stage per GPU and a flush beat $(p-1)/(m+p-1)$?
   Answer

   The last stage cannot start before $(p-1)t_f$, it does $m(t_f + t_b)$ of work, and after its last backward pass $p - 1$ backward passes on the other stages remain: $(p-1)t_b$. So the step takes at least $(m+p-1)(t_f+t_b)$. Interleaving escapes the bound because each GPU holds several stages.
5. Global batch 512 sequences, $b = 1$, $p = 8$. What is the bubble with $N_d = 1$ and with $N_d = 16$ data-parallel replicas?
   Answer

   $N_d = 1$: $m = 512$, bubble $7/519 \approx 1.3\%$. $N_d = 16$: $m = 32$, bubble $7/39 \approx 17.9\%$.
6. In worked example (c), why does TP 8 × PP 8 with 1F1B store as many activations per GPU as TP 8 without PP?
   Answer

   Stage 0 holds up to $p$ micro-batches, each through $L/p$ layers: $p \cdot L/p = L$ layers of one micro-batch, the same as one micro-batch through the whole model.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Why does "column split, then row split" need only one all-reduce in the forward pass?

   Show answer

2. What do Megatron's operators $f$ and $g$ do?

   Show answer

3. How many TP all-reduces does one transformer layer need, and on what size?

   Show answer

4. Write the TP communication-to-compute ratio per layer.

   Show answer

5. Compute the TP ratio for $t = 8$, $h = 8192$, $F = 989$ TFLOP/s on NVLink (450 GB/s).

   Show answer

6. What is the pipeline bubble of GPipe and 1F1B?

   Show answer

7. How many micro-batches does stage $s$ keep in memory in 1F1B, and in GPipe?

   Show answer

8. What is the bubble of the interleaved schedule with $v$ chunks per GPU?

   Show answer

9. Why does tensor parallelism stay inside a node?

   Show answer

10. In which order did Llama 3 place its parallel dimensions, and why?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Implement column-parallel and row-parallel linear layers." — Column: replicated input, sharded output, $f$ before it. Row: sharded input, partial sums, $g$ (all-reduce) after it. Check against the unsplit layer in float64, forward and backward.
• "How many all-reduces does a TP transformer layer need?" — 2 in the forward pass and 2 in the backward pass, each on $b \cdot s \cdot h$ values.
• "GPipe, 1F1B or interleaved?" — GPipe and 1F1B have the same bubble, but 1F1B keeps $p$ instead of $m$ micro-batches in memory; interleaving divides the bubble time by $v$ at $v$ times the messages.
• "64 GPUs, 8 per node, a 70B model: choose TP, PP and DP." — TP = 8 in the node; the smallest PP that fits with 1F1B; DP for the rest; check $m \ge 4p$ and $T \ge F/W$ for DP. Justify TP in the node with $(t-1)F/(3hW)$: 0.63 on NVLink against 5.6 on InfiniBand for $h = 8192$.

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. Four exercises, each with a check cell and a folded solution:

1. The tensor-parallel MLP on 4 simulated GPUs: forward and backward with explicit all-reduces, checked against the unsplit MLP (outputs and all gradients).
2. A pipeline schedule simulator for GPipe and 1F1B: check the step time $(m+p-1)(t_f+t_b)$ and the in-flight counts.
3. The interleaved schedule: simulate it and check the bubble $\frac{p-1}{vm+p-1}$ for several $v$.
4. layout_report(...): memory per GPU, bubble, TP and DP ratios for a layout; then search all layouts of 64 GPUs for a 70B model and find the best ones.

11. Watch and read #

Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM | Jared CasperJared Casper · @Scale · 24 min

Data, tensor and pipeline parallelism and how Megatron combines them, from one of its authors.

Lecture 48: The Ultra Scale PlaybookNouamane Tazi · GPU MODE · 184 min

Starts at tensor parallelism (1:32:11); pipeline parallelism starts at 2:11:20.

• Shoeybi et al. (2019), Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism: §3, the MLP and attention splits, the $f$ and $g$ operators (Figures 3 and 4).
• Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM: §2.2 (GPipe, 1F1B and interleaved bubbles), §3 (takeaways on combining TP, PP and DP), §4.1 (scatter/gather).
• Huang et al. (2019), GPipe, and Narayanan et al. (2021), Memory-Efficient Pipeline-Parallel DNN Training (PipeDream-Flush, the 1F1B schedule).
• Llama Team (2024), The Llama 3 Herd of Models: §3.3.2 and Table 4, the 4D parallelism of the 405B run.
• Qi et al. (2024), Zero Bubble Pipeline Parallelism; DeepSeek-AI (2024), DeepSeek-V3 Technical Report (DualPipe); Xu et al. (2021), An Efficient 2D Method for Training Super-Large Deep Learning Models; Bian et al. (2021), Maximizing Parallelism in Distributed Training for Huge Neural Networks.
• Wang et al. (2024), ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability.
• Hugging Face, The Ultra-Scale Playbook: the chapters on tensor and pipeline parallelism and "Finding the best training configuration"; and picotron, a small readable implementation.
• Austin et al., How to Scale Your Model, How to Parallelize a Transformer for Training: the same trade-offs in the language of TPU meshes.
• Stanford CS336 (2025), lectures 7 Parallelism 1 and 8 Parallelism 2 (on YouTube).