Tensor parallelism splits a matmul, pipeline parallelism splits a network, and only one of them scales past a single node

FSDP shards a model’s persistent state across GPUs and hides most of the resulting communication behind compute, but sharding state doesn’t shrink a single matrix multiplication, and it doesn’t shrink a model’s depth. Two separate walls remain after ZeRO-3 has done everything it can: one weight matrix inside one layer can still be too large for one GPU’s memory or too slow for one GPU’s compute to be worth running alone, and a sufficiently deep model can exceed what fits even after every optimizer trick has been applied. Tensor parallelism answers the first wall by splitting the algebra inside a layer. Pipeline parallelism answers the second by splitting the model’s depth across groups of GPUs. They solve different problems, they communicate in structurally different ways, and, worth stating up front because it shapes every real deployment: only one of the two tolerates being spread across nodes.

Column-parallel and row-parallel are the same layer, split two different ways#

For a linear layer $Y = XW$, splitting $W$ by columns puts a different slice of the output features on each rank:

$W = [\,W_0 \mid W_1\,], \qquad Y_i = X W_i$

Each rank computes its own slice of $Y$ independently, no communication required to produce it, the result is simply sharded along the feature dimension across ranks.

Splitting $W$ by rows instead puts a different slice of the input features on each rank, which forces $X$ to already be split the same way:

$W = \begin{bmatrix}W_0 \\ W_1\end{bmatrix}, \quad X = [\,X_0 \mid X_1\,], \qquad Y_i = X_i W_i$

Now each rank only computes a partial contribution to the full output, and the true result is the sum across ranks, $Y = \sum_i Y_i$, which does require communication, an AllReduce, to actually materialize.

Worked example, $2$ ranks, $X$ a single $1{\times}4$ row, $W$ a $4{\times}2$ matrix, $X = [1, 2, 3, 4]$, $W = \begin{bmatrix}1&0\\0&1\\1&0\\0&1\end{bmatrix}$.

import numpy as np
X = np.array([[1, 2, 3, 4]])
W = np.array([[1, 0], [0, 1], [1, 0], [0, 1]])
full = X @ W
print(full)   # [[4 6]]  -- the reference result

# Row split: rank 0 owns rows 0-1 of W and columns 0-1 of X, rank 1 owns rows 2-3 and columns 2-3
X0, X1 = X[:, :2], X[:, 2:]
W0, W1 = W[:2, :], W[2:, :]
Y0, Y1 = X0 @ W0, X1 @ W1
print(Y0 + Y1)   # [[4 6]] -- matches only after the AllReduce-equivalent sum

Y0 and Y1 individually are $[1,2]$ and $[3,4]$, neither one is the answer, they’re both wrong until summed, and that sum is exactly what an AllReduce over the row-parallel output computes in a real, multi-GPU run.

Why Megatron’s actual layer only needs one AllReduce, not two#

A transformer MLP block is two linear layers with a nonlinearity between them, and the placement of column-split versus row-split across those two layers isn’t arbitrary, it’s the entire point of the design. Make the first linear layer column-parallel: its output stays sharded, no communication yet, and the nonlinearity applies element-wise, so it runs correctly on each rank’s shard with zero communication either. Make the second linear layer row-parallel, taking that already-sharded activation as its already-split input: its output needs exactly one AllReduce to reduce into the correct final result. Two matrix multiplications, one nonlinearity, and the communication cost of the entire block is one AllReduce, not one after every matmul. This is the specific insight worth taking from Megatron-LM’s design: which operation gets column- versus row-split determines whether intermediate activations stay usefully sharded or have to be reconstructed at every step, and getting that placement right is most of what makes tensor parallelism affordable at all.

PyTorch’s native tensor-parallel API is a direct, literal expression of that placement decision, not a separate abstraction on top of it:

from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel

parallelize_module(
    mlp_block,
    device_mesh=tp_mesh,
    parallelize_plan={
        "gate_proj": ColwiseParallel(),    # first linear: column-split, output stays sharded
        "up_proj":   ColwiseParallel(),
        "down_proj": RowwiseParallel(),    # second linear: row-split, triggers the one AllReduce
    },
)

The plan dictionary is the design decision from the paragraph above, spelled out per-submodule: gate_proj and up_proj column-parallel because their sharded output can feed directly into a row-parallel down_proj without anything being reconstructed in between, and the single AllReduce happens exactly where RowwiseParallel is applied, nowhere else in the block.

Attention partitions along a dimension the architecture already provides for free: each head’s computation is independent until the final concatenation, so heads split cleanly across TP ranks the same way a column-parallel layer does, sharded through the entire attention computation, one communication step at the boundary.

Vocabulary-parallel embeddings and loss extend the same idea to the largest single tensor in many models, the input embedding and output projection, which scale with vocabulary size. Sharding the vocabulary dimension across TP ranks means no single GPU ever materializes the full embedding table, but it means computing cross-entropy loss correctly requires care: each rank only holds logits for its own slice of the vocabulary, so the softmax normalizer has to be computed as a local partial reduction on each rank and then combined across ranks, rather than naively computing softmax as if the full vocabulary were locally present.

Sequence parallelism, paired with TP, is not context parallelism, a distinction worth being precise about since the two names invite confusion and get covered separately for exactly that reason. Operations like LayerNorm and dropout aren’t naturally tensor-parallelizable the way a matmul is, so in a pure TP scheme they end up redundantly replicated in full on every TP rank, wasting activation memory. Megatron’s sequence parallelism shards exactly these operations along the sequence dimension instead, reclaiming that redundant activation memory, but it’s a memory optimization riding alongside tensor parallelism, operating on the same TP process group. It does not shard attention’s actual computation across a longer context the way context parallelism does.

TP=8 in a real production config is not a round number, it’s a bandwidth constraint#

Tensor parallelism issues an AllReduce (or equivalent) at every layer boundary, for every microbatch, which makes it by far the most communication-frequent form of parallelism in this series. That frequency is exactly why TP is kept inside a single node in practice: an NVLink-connected GPU pair moves data at a different order of magnitude than an inter-node network link, and paying inter-node latency at every single layer boundary, many times per second, is a cost tensor parallelism specifically cannot absorb the way pipeline parallelism’s much less frequent point-to-point sends can.

Llama 3 405B’s published training configuration makes this concrete rather than a rule of thumb: TP is fixed at $8$ across every stage of its pretraining, matching exactly the 8-GPU NVLink domain of one node. TP is capped there deliberately, “optimal given the batch size and hierarchical network bandwidth constraints,” setting TP no higher than the node size specifically to guarantee it only ever uses intra-node NVLink, never the slower inter-node fabric.

Pipeline parallelism splits depth, and pays for it in bubble instead of bandwidth#

GPU group 0        GPU group 1        GPU group 2
Layers 0-7   ───▶  Layers 8-15  ───▶ Layers 16-23

Run one full batch through this naively and every stage but the first sits idle waiting for work, then idle again waiting for the backward pass to arrive back. The standard fix is splitting the batch into $M$ microbatches and pipelining them through the stages:

Time ─────────────────────────────────────────▶
Stage 0:  F0  F1  F2  F3  B0  B1  B2  B3
Stage 1:      F0  F1  F2  F3  B0  B1  B2  B3
Stage 2:          F0  F1  F2  F3  B0  B1  B2  B3

The empty leading and trailing corners are the bubble, compute time paid for nothing. GPipe’s bubble fraction, for $P$ stages and $M$ microbatches:

$\text{bubble} \approx \frac{P-1}{M+P-1}$

Worked at two microbatch counts, same $4$-stage pipeline. At $M=8$: $\frac{4-1}{8+4-1} = \frac{3}{11} \approx 27\%$, over a quarter of every step spent idle. At $M=32$: $\frac{3}{32+3} = \frac{3}{35} \approx 8.6\%$. Same model, same stage count, same hardware, quadrupling the microbatch count alone drops idle time by roughly $19$ percentage points, which is exactly why “how many microbatches” is a first-order tuning question for pipeline parallelism and not a minor knob.

def gipe_bubble_fraction(stages, microbatches):
    return (stages - 1) / (microbatches + stages - 1)

for m in (4, 8, 16, 32, 64):
    print(m, round(gipe_bubble_fraction(4, m), 3))
# 4 0.429   8 0.273   16 0.158   32 0.086   64 0.045

GPipe, 1F1B, and interleaved are three different answers to the same trade. GPipe, as diagrammed above, runs all forward passes before any backward pass, which keeps $M$ microbatches’ worth of activations alive simultaneously, a real activation-memory cost. 1F1B (one-forward-one-backward) reaches a steady state where each stage alternates a forward and a backward pass, releasing activation memory for a completed microbatch far sooner, same bubble fraction, meaningfully lower peak activation memory. Interleaved 1F1B goes further, assigning each device more than one, non-contiguous, pipeline stage, which shrinks the bubble below the plain formula above at the cost of more, smaller point-to-point communication rounds between stages, trading bubble size for communication count, the same fundamental trade tensor parallelism already made in the other direction.

Splitting a model and running one of these schedules against it is two calls, not a hand-built microbatch loop:

from torch.distributed.pipelining import pipeline, SplitPoint, ScheduleGPipe

pipe = pipeline(
    model,
    mb_args=(example_microbatch,),
    split_spec={"layers.8": SplitPoint.BEGINNING, "layers.16": SplitPoint.BEGINNING},
)   # cuts a 24-layer model into the 3 stages diagrammed above, at layers 8 and 16

stage = pipe.get_stage(stage_index=rank, device=device)
schedule = ScheduleGPipe(stage, n_microbatches=32)   # the M=32 configuration worked out above

if rank == 0:
    losses = schedule.step(full_batch)   # rank 0 feeds the batch in; every other rank calls step() with no argument

split_spec is where the stage boundaries from the ASCII diagram earlier become an actual, executable cut, and swapping ScheduleGPipe for Schedule1F1B or an interleaved variant changes only which schedule class runs against the same split model, the stage-partitioning code above it doesn’t change at all.

Table 4, and a citation worth checking against its own arithmetic#

Llama 3 405B’s published parallelism configuration gives two rows worth comparing directly, and the comparison is a real illustration of how these dimensions trade off in production, not a hypothetical:

Stage	GPUs	TP	CP	PP	DP	Seq. len	TFLOPs/GPU	BF16 MFU
Short-context	8,192	8	1	16	64	8,192	430	43%
Long-context	8,192	8	16	16	4	131,072	380	38%

Worth verifying rather than repeating: some secondary summaries of this table state the long-context row’s GPU count as $16{,}384$, presumably borrowing Meta’s separate, headline figure for the full training fleet. The row’s own numbers don’t support that: $8 \times 16 \times 16 \times 4 = 8{,}192$, identical to the short-context row’s GPU count, not double it. Both configurations run on the same 8,192-GPU slice; what actually changes between them is the mesh. TP and PP hold fixed at $8$ and $16$, TP pinned to the node size as already derived, PP presumably balanced against model depth and available stages. What moves is context parallelism, $1 \to 16$, absorbing the much larger per-token activation memory that a $16\times$ longer sequence demands, funded by shrinking the data-parallel width, $64 \to 4$, to keep total GPU count constant. And the cost of that trade shows up directly in the last column: MFU drops from $43\%$ to $38\%$, the real price of context parallelism’s added communication, paid to make a $131$K-token context trainable at all rather than not fitting in memory.

When PP earns its bubble, and when it doesn’t#

Reach for pipeline parallelism once FSDP and tensor parallelism together still can’t fit the model, and only once the batch supports enough microbatches to keep the bubble formula’s denominator meaningfully larger than $P{-}1$. A pipeline with too few microbatches relative to its stage count pays a bubble tax that can rival or exceed whatever memory problem it was meant to solve. Stage balance matters just as much as microbatch count: one disproportionately heavy stage means every other stage waits on it every single microbatch, turning a theoretically well-amortized bubble into a real, measured stall that shows up on exactly one stage’s profiler trace and nowhere else.

What this takes to be frontier-job-ready#

The technical axis: being able to say, for a specific model and specific hardware topology, which wall TP solves and which wall PP solves, and deriving the actual bubble fraction and communication frequency rather than reaching for “more parallelism” as an undifferentiated response to an OOM.

The operational axis is stated close to verbatim in real postings for this class of work: Google DeepMind’s Research Engineer, Pretraining role for Gemini lists, as a hard requirement rather than a nice-to-have, the ability to “understand common Accelerated Linear Algebra (XLA) primitives and how JAX code runs on TPUs in practice” and to “work across the LLM preparation stack (pretraining, fine tuning, serving),” with distributed training on TPU pods named alongside it as its own requirement. Worth being precise about what that implies: it’s a JAX/XLA/TPU stack, not PyTorch/GPU, and the TP and PP mental models built here transfer conceptually but the actual primitives still have to be relearned for a role built on Google’s own hardware.

The autonomy axis: reading Table 4 correctly, as this post did, means treating a stated GPU count as a claim to check against the row’s own numbers rather than accepting it because a paper published it, the identical discipline this series applies to secondhand hiring-language quotes.

Common mistakes#

Treating TP as a knob that scales like DP does: TP’s AllReduce frequency, once per layer boundary per microbatch, makes crossing node boundaries with it a bandwidth mistake, not a scaling choice, which is why real configurations pin it to the NVLink domain size.

Conflating Megatron’s sequence parallelism with context parallelism because both shard something along the sequence axis: one reclaims activation memory for LayerNorm and dropout inside an existing TP group, the other shards attention’s actual computation across a long context, and the fix for one problem doesn’t touch the other.

Picking a microbatch count without checking the bubble formula: the difference between $M=8$ and $M=32$ on a 4-stage pipeline is roughly 19 percentage points of pure idle time, not a rounding error.

Repeating a GPU count from a secondary summary without checking it against the row’s own TP×CP×PP×DP product: exactly the arithmetic mismatch caught above.

Try it yourself#

Beginner. Using the row-parallel worked example, split $X = [2, 0, 1, 3]$ and $W = \begin{bmatrix}1&1\\0&1\\1&0\\2&1\end{bmatrix}$ across 2 ranks the same way, compute each rank’s partial $Y_i$, sum them, and confirm the result matches the unsharded $XW$.

Intermediate. Using gipe_bubble_fraction, find the smallest microbatch count $M$ for an 8-stage pipeline that gets the bubble fraction under 10%, and explain why an 8-stage pipeline needs a larger $M$ than the 4-stage example above to hit the same target.

Advanced. Given Llama 3’s two published rows, explain in your own derivation why context parallelism was funded specifically by shrinking DP rather than by shrinking PP or TP, tying your answer to which of the three dimensions each additional GPU actually helps with when the bottleneck is per-token attention memory at long context.

---

The one-sentence version: tensor parallelism buys memory and compute headroom inside a layer at the cost of an AllReduce at every layer boundary, which is exactly why real configurations never let it leave the NVLink domain, pipeline parallelism buys depth headroom at the cost of a bubble that shrinks as microbatch count grows, and Llama 3’s own published table is a clean, verifiable record of trading data-parallel width for context-parallel width, at a real, measured MFU cost, to make long-context training fit at all. Attention’s own memory wall at long context is what the next dimension exists to solve, context parallelism shards the sequence itself, a fundamentally different communication pattern than anything covered here.