3.6 Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism¶

In Distributed Training I we learned to replicate a model across many GPUs and split the data. That strategy — data parallelism (DP), and its memory-sharded cousins ZeRO and FSDP — has one non-negotiable requirement: a single replica’s worth of state (parameters, gradients, optimizer states, and at least one microbatch of activations) must fit on the device, or be shardable into something that fits. When that breaks, you must split the model itself.

A 70B-parameter model in bf16 is 140 GB of weights alone. Add fp32 optimizer states (Adam’s two moments plus an fp32 master copy) and you are well over half a terabyte before a single token of activation memory. No single GPU — not an 80 GB H100, not a 192 GB MI300X — holds that. You have no choice but to carve the model up and spread the pieces across devices. This chapter is about the three orthogonal axes for doing so:

• Tensor parallelism (TP) — split within a layer (each matmul is shared across GPUs).
• Pipeline parallelism (PP) — split across layers (each GPU owns a contiguous block of layers).
• Sequence / context parallelism (SP/CP) — split along the sequence dimension (each GPU owns part of the token sequence).
• Expert parallelism (EP) — split across the experts of a Mixture-of-Experts layer.

These compose with data parallelism into what practitioners call 3D, 4D, or 5D parallelism. Getting the composition right — and understanding exactly where the communication lives — is the single highest-leverage systems skill in large-scale pretraining. It is also a favorite interview topic precisely because it forces you to reason about the memory hierarchy, collective communication, and the transformer block all at once.

The Memory Wall and the Three Axes of Splitting¶

Let us first quantify why we split. Consider a dense decoder-only transformer with \(L\) layers, hidden size \(h\), and FFN expansion factor 4. The parameter count is dominated by

For training, every parameter carries a memory multiplier. With mixed-precision Adam (see Mixed Precision, bf16 & FP8 Training) and the optimizer details from Optimizers, the per-parameter footprint is roughly:

So a model needs on the order of \(16N\) bytes of static state, plus activations. For a 70B model that is \(\approx 1.1\) TB — which ZeRO/FSDP can shard across the DP group. But ZeRO does not reduce the activation memory of a single microbatch, and it does not reduce the size of the largest single tensor you must materialize. When a single layer’s activations, or a single forward pass, are too big for one device, sharding optimizer state is not enough. That is where the model-parallel axes come in.

The mental model: DP/ZeRO splits state across replicas of the whole model; model parallelism splits the computation graph itself. They are orthogonal and combine multiplicatively.

Tensor Parallelism: Splitting the Matmul¶

Tensor parallelism — introduced at scale by Megatron-LM (Shoeybi et al., 2019) — observes that the heavy lifting of a transformer is a sequence of large matrix multiplications, and a matrix multiplication can be partitioned across devices with a single collective per region. The art is choosing partitions so that consecutive matmuls compose without a collective between them.

Column and Row Parallel Linear Layers¶

Take a linear layer \(Y = XA\) where \(X \in \mathbb{R}^{s \times h}\) (sequence \(\times\) hidden) and \(A \in \mathbb{R}^{h \times h'}\). There are two ways to split \(A\) across \(t\) GPUs.

Column-parallel. Split \(A\) along its output columns: \(A = [A_1, A_2, \dots, A_t]\) where each \(A_i \in \mathbb{R}^{h \times h'/t}\). Each GPU holds the full input \(X\) (replicated) and computes a slice of the output:

No communication is needed to produce \(Y\) — each GPU just holds a different chunk of the output columns. The output is sharded along the feature dimension.

Row-parallel. Split \(A\) along its input rows: \(A = [A_1; A_2; \dots; A_t]\), \(A_i \in \mathbb{R}^{h/t \times h'}\), and correspondingly split the input \(X = [X_1, \dots, X_t]\) along its columns. Each GPU computes a partial sum:

Each GPU produces a full-shaped \(Y\) but containing only its partial contribution; an all-reduce sums them to the correct result.

The magic trick: chain a column-parallel layer into a row-parallel layer and the intermediate never needs to be gathered. A column-parallel layer leaves its output sharded along features; a row-parallel layer wants its input sharded along features. They fit like puzzle pieces, and you pay exactly one all-reduce at the very end (plus one in the backward pass).

Applying It to the Transformer Block¶

Megatron maps this pattern onto both sublayers of a transformer block.

MLP block \(Y = \text{GeLU}(XA)B\):

• \(A\) (the up-projection, \(h \to 4h\)) is column-parallel → the \(4h\) intermediate is sharded across GPUs.
• GeLU is elementwise, so it acts independently on each shard — no communication, and crucially we did not have to gather before the nonlinearity (which would have been wrong to split naively).
• \(B\) (the down-projection, \(4h \to h\)) is row-parallel → one all-reduce produces the final output.

Attention block. This is even more natural because attention is already partitioned by heads. With \(a\) attention heads and \(t\) TP ranks, give each GPU \(a/t\) heads:

• The \(Q, K, V\) projections are column-parallel — each GPU produces the Q/K/V for its own heads only.
• Each GPU runs full self-attention (softmax, the \(QK^\top\), the \(\times V\)) for its heads with no cross-GPU communication. (This is why TP and Multi-Head Attention are such a good fit, and why GQA changes the K/V sharding story — see the warning below.)
• The output projection \(O\) is row-parallel → one all-reduce.

So each transformer block needs exactly two all-reduces in the forward pass (one after attention’s output proj, one after the MLP’s down proj) and two in the backward pass. In Megatron’s notation these are the operators \(f\) and \(g\): \(f\) is identity in forward / all-reduce in backward; \(g\) is all-reduce in forward / identity in backward.

import torch
import torch.distributed as dist
import torch.nn as nn

# Assume a TP process group `tp_group` of size `tp` already initialized.
# These two autograd functions place the all-reduces in exactly the
# right spots: g = forward-allreduce, f = backward-allreduce.

class _CopyToTPRegion(torch.autograd.Function):
    """f operator: identity forward, all-reduce backward."""
    @staticmethod
    def forward(ctx, x): return x
    @staticmethod
    def backward(ctx, grad):
        dist.all_reduce(grad, group=tp_group)   # sum grads from all TP ranks
        return grad

class _ReduceFromTPRegion(torch.autograd.Function):
    """g operator: all-reduce forward, identity backward."""
    @staticmethod
    def forward(ctx, x):
        dist.all_reduce(x, group=tp_group)      # sum partial outputs
        return x
    @staticmethod
    def backward(ctx, grad): return grad

copy_to_region   = _CopyToTPRegion.apply
reduce_from_region = _ReduceFromTPRegion.apply

class ColumnParallelLinear(nn.Module):
    """Y = X A, with A split along output columns across `tp` ranks.
    Output is sharded along the feature dim (gather_output=False)."""
    def __init__(self, in_f, out_f, tp, rank, bias=True):
        super().__init__()
        assert out_f % tp == 0
        self.out_local = out_f // tp
        # Each rank only allocates its slice of the weight.
        self.weight = nn.Parameter(torch.empty(self.out_local, in_f))
        self.bias   = nn.Parameter(torch.zeros(self.out_local)) if bias else None
        nn.init.normal_(self.weight, std=0.02)

    def forward(self, x):
        x = copy_to_region(x)                    # f: ensures correct bwd all-reduce
        y = torch.nn.functional.linear(x, self.weight, self.bias)
        return y                                  # shape [*, out_f/tp], sharded

class RowParallelLinear(nn.Module):
    """Y = X A, with A split along input rows; input already sharded.
    Produces the full output via an all-reduce (g operator)."""
    def __init__(self, in_f, out_f, tp, rank, bias=True):
        super().__init__()
        assert in_f % tp == 0
        self.in_local = in_f // tp
        self.weight = nn.Parameter(torch.empty(out_f, self.in_local))
        # bias is added ONCE, after the reduce, so only rank 0 should hold it
        self.bias = nn.Parameter(torch.zeros(out_f)) if (bias and rank == 0) else None
        nn.init.normal_(self.weight, std=0.02)

    def forward(self, x):                         # x sharded along features
        y = torch.nn.functional.linear(x, self.weight)   # partial sum
        y = reduce_from_region(y)                 # g: all-reduce -> full output
        if self.bias is not None:
            y = y + self.bias
        return y

# A Megatron-style MLP: column then row, nonlinearity in between, zero gathers.
class ParallelMLP(nn.Module):
    def __init__(self, h, tp, rank):
        super().__init__()
        self.fc1 = ColumnParallelLinear(h, 4 * h, tp, rank)   # h -> 4h, sharded
        self.fc2 = RowParallelLinear(4 * h, h, tp, rank)      # 4h -> h, all-reduce
    def forward(self, x):
        return self.fc2(torch.nn.functional.gelu(self.fc1(x)))

The Communication Cost of TP¶

The all-reduce in each block moves a tensor of shape \(s \times h\). With ring all-reduce over \(t\) devices, each device sends and receives \(\approx 2 \cdot \frac{t-1}{t} \cdot (s \cdot h \cdot 2\text{ bytes})\) per all-reduce. Two all-reduces forward + two backward = four all-reduces per layer per step. This is a lot of traffic, and it sits squarely on the critical path: the GPUs cannot proceed past the all-reduce until it completes. Frameworks claw some of it back by overlapping each collective with the GEMM that feeds it: Megatron-LM’s --tp-comm-overlap (built on Transformer Engine’s “userbuffers”) chunks the all-gather/reduce-scatter and pipelines the chunks against the matmul, and PyTorch’s async tensor parallel does the same on top of torch.distributed._symmetric_memory. Overlap hides exposed latency; it does not reduce the volume, so the placement rule below still binds.

This is the defining constraint of tensor parallelism: it must run over the fastest interconnect you have. On a DGX/HGX node that is NVLink/NVSwitch (hundreds of GB/s, sometimes ~900 GB/s aggregate). Cross more than the NVLink domain — e.g. over InfiniBand between nodes — and TP collapses your throughput because the per-layer all-reduces serialize behind a 10–25× slower link. Rule of thumb: keep the TP group inside one node, \(t \le 8\) (or whatever your NVLink domain is). The real rule is “TP must not leave the NVLink domain,” and 8 is merely that domain’s most common size: on rack-scale NVLink systems (GB200/GB300 NVL72, where dozens of GPUs share one NVLink fabric) the ceiling rises accordingly, which is exactly why 2026-era frontier configs can afford larger TP (and larger EP) groups than an 8-GPU box allows.

The Vocabulary: Parallel Embedding and Parallel Cross-Entropy¶

The token embedding table and the (usually weight-tied) LM head are each a \(V \times h\) matrix, with \(V\) ranging from ~32k to 256k+ — for large \(h\) this is routinely the single biggest matrix in the model, bigger than any individual attention or MLP weight. You cannot assemble an end-to-end TP transformer without sharding these along the vocabulary dimension across the \(t\) TP ranks: a third partitioning strategy, distinct from both the column-parallel and row-parallel matmul patterns above.

Vocab-parallel embedding. Rank \(r\) owns rows \([r \cdot V_{\text{loc}} : (r+1) \cdot V_{\text{loc}})\) of the table, where \(V_{\text{loc}} = V / t\). A lookup for a token id outside a rank’s local range must resolve to zero on that rank; summing across ranks then recovers the correct row, since exactly one rank owns each id.

import torch, torch.nn as nn, torch.nn.functional as F, torch.distributed as dist

class VocabParallelEmbedding(nn.Module):
    def __init__(self, V, h, tp, rank, tp_group):
        super().__init__()
        assert V % tp == 0
        self.V_loc = V // tp
        self.vocab_start = rank * self.V_loc
        self.vocab_end = self.vocab_start + self.V_loc
        self.tp_group = tp_group
        self.weight = nn.Parameter(torch.empty(self.V_loc, h))   # [V_loc, h]

    def forward(self, input_ids):                      # input_ids: [b, s], long
        mask = (input_ids < self.vocab_start) | (input_ids >= self.vocab_end)
        masked_ids = input_ids.clone() - self.vocab_start
        masked_ids[mask] = 0                            # keep in-range for the local table
        y = F.embedding(masked_ids, self.weight)         # [b, s, h]
        y = y.masked_fill(mask.unsqueeze(-1), 0.0)        # zero rows this rank doesn't own
        dist.all_reduce(y, group=self.tp_group)           # SUM: exactly one rank owns each row
        return y                                           # [b, s, h]

The backward is the identity on the local rows — gradients flow only into the rows a rank actually owns — mirroring the \(g\)/\(f\) forward-all-reduce, backward-identity (and vice versa) operator pairing used for the MLP/attention sharding above.

Vocab-parallel cross-entropy. The mirror image on the output side: a column-parallel LM head (ColumnParallelLinear with gather_output=False) leaves logits sharded along vocab as logits_local of shape [N, V_loc] (with \(N = b \cdot s\) flattened). Computing the loss without ever gathering the full [N, V] logits is the whole point — \(V\) is exactly the dimension we sharded to avoid materializing.

def vocab_parallel_cross_entropy(logits_local, target, tp_group, vocab_start, vocab_end):
    # logits_local: [N, V_loc] fp32;  target: [N] global token ids
    V_loc = logits_local.shape[-1]

    # 1) global row-max via all-reduce MAX
    m = logits_local.max(dim=-1).values                           # [N]
    dist.all_reduce(m, op=dist.ReduceOp.MAX, group=tp_group)

    # 2) global normalizer Z via all-reduce SUM of local exp-sums
    exp_local = torch.exp(logits_local - m.unsqueeze(-1))         # [N, V_loc]
    sum_exp = exp_local.sum(dim=-1)                                 # [N]
    dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group)

    # 3) the target logit: only the rank that owns `target` contributes a nonzero value
    tmask = (target >= vocab_start) & (target < vocab_end)
    local_target = (target - vocab_start).clamp(0, V_loc - 1)
    pred = logits_local.gather(1, local_target.unsqueeze(-1)).squeeze(-1)   # [N]
    pred = pred * tmask
    dist.all_reduce(pred, op=dist.ReduceOp.SUM, group=tp_group)   # now the true target logit everywhere

    # loss = logsumexp(z) - z_target = (m + log Z) - z_target
    loss = torch.log(sum_exp) + m - pred                           # [N]
    return loss.mean()

Three all-reduces of shape \([N]\) — trivial next to gathering \([N, V]\) logits. The backward is equally cheap: the gradient w.r.t. logits_local is softmax_local - onehot_local, already vocab-sharded, needing no extra collective.

Megatron-LM implements exactly these primitives as megatron.core.tensor_parallel.VocabParallelEmbedding and vocab_parallel_cross_entropy, in megatron/core/tensor_parallel/{layers.py, cross_entropy.py} — the mask-and-all-reduce embedding trick and the max-then-sum parallel softmax above are the reference formulation, not a simplification of it. The LM head is typically a ColumnParallelLinear(gather_output=False) whose sharded output feeds directly into vocab_parallel_cross_entropy, and its weight is often tied to the embedding weight.

Sequence Parallelism as TP’s Free Companion¶

Look again at the Megatron block: between the two TP regions sit the LayerNorm/RMSNorm and the dropout/residual-add, which Megatron replicates (every TP rank does the same redundant work on the full \(s \times h\) tensor). That replication wastes both compute and, more importantly, activation memory: each rank stores the full LayerNorm activations.

Sequence parallelism (SP) — in this Megatron sense (Korthikanti et al., 2022) — splits those replicated regions along the sequence dimension instead. The norm and residual are now done on \(s/t \times h\) shards. The catch: the boundaries between an SP region (sharded on sequence) and a TP region (sharded on hidden/features) require a conversion. Megatron shows that the same all-reduce of the TP region can be decomposed into a reduce-scatter + all-gather that achieves the layout conversion for the same total communication volume. So Megatron-style SP is essentially free communication-wise and meaningfully cuts activation memory — it is now standard and always-on in Megatron. (Note: this “sequence parallelism” is a memory optimization within a TP group, and is distinct from context parallelism / Ring Attention, covered later, which is a true sequence split for long context.)

One correctness detail this split forces, easy to miss when hand-rolling TP: dropout needs two different RNG regimes. Dropout applied inside a TP region acts on a distinct feature shard per rank, so each rank must draw a different mask — using the same seed everywhere would correlate the masks and effectively reduce the dropout rate. Dropout applied to a replicated tensor (e.g. on the residual stream before SP shards it) must use the same seed on all ranks, or the replicas silently diverge. Megatron keeps both states in a seed tracker (model_parallel_cuda_manual_seed / get_cuda_rng_tracker in megatron/core/tensor_parallel/random.py) and switches between them at region boundaries; PyTorch’s DTensor path handles the equivalent bookkeeping through its OffsetBasedRNGTracker. If your architecture has no dropout (as most 2026 pretraining recipes do not), this problem disappears — which is one small extra reason the modern default is dropout-free pretraining.

The Library: PyTorch-Native TP with DTensor¶

Everything above is what Megatron-LM implements by hand — and what you should implement once, to know it. In day-to-day work you would reach for PyTorch’s built-in tensor parallelism instead, which expresses exactly the same algebra declaratively. The abstraction is DTensor: a tensor carrying a placement — Shard(dim) or Replicate() — over a DeviceMesh. You write a plan saying how each submodule is sharded; the runtime inserts the \(f\)/\(g\) collectives, the SP reduce-scatter/all-gather boundary conversions, and the backward counterparts for you.

# Assumes a Llama-style `model` whose blocks expose the submodule names below,
# and parallelism degrees (dp, tp) with dp * tp == world_size.
import torch, torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, Replicate
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel,
    SequenceParallel, PrepareModuleInput, loss_parallel,
)

# One mesh, two axes: FSDP2 across the "dp" axis, TP inside the node on "tp".
mesh = init_device_mesh("cuda", (dp, tp), mesh_dim_names=("dp", "tp"))
tp_mesh = mesh["tp"]

for block in model.layers:
    parallelize_module(block, tp_mesh, {
        # Norms + residual run sequence-parallel (Korthikanti-style SP).
        "attn_norm": SequenceParallel(),
        "mlp_norm":  SequenceParallel(),
        # Boundary conversion Shard(1) [sequence] -> Replicate(): this IS the
        # all-gather half of the reduce-scatter/all-gather decomposition.
        "attn": PrepareModuleInput(input_layouts=(Shard(1),),
                                   desired_input_layouts=(Replicate(),)),
        "attn.wq": ColwiseParallel(),          # Q/K/V: column-parallel by head
        "attn.wk": ColwiseParallel(),
        "attn.wv": ColwiseParallel(),
        "attn.wo": RowwiseParallel(output_layouts=Shard(1)),   # g, then re-shard on seq
        "mlp": PrepareModuleInput(input_layouts=(Shard(1),),
                                  desired_input_layouts=(Replicate(),)),
        "mlp.w_gate": ColwiseParallel(),       # SwiGLU gate + up are column-parallel
        "mlp.w_up":   ColwiseParallel(),
        "mlp.w_down": RowwiseParallel(output_layouts=Shard(1)),  # row-parallel down-proj
    })

# Embedding sharded on vocab rows; LM head column-parallel, output left vocab-sharded.
parallelize_module(model, tp_mesh, {
    "tok_emb":   RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)),
    "final_norm": SequenceParallel(),
    "lm_head":   ColwiseParallel(output_layouts=Shard(-1), use_local_output=False),
})

logits = model(input_ids)                  # DTensor, still sharded on the vocab dim
with loss_parallel():                      # == our vocab_parallel_cross_entropy
    loss = F.cross_entropy(logits.flatten(0, 1), targets.flatten(0, 1))
    loss.backward()

Read that plan against the hand-written code above and every line maps: ColwiseParallel/RowwiseParallel are ColumnParallelLinear/RowParallelLinear, PrepareModuleInput is the SP↔TP layout conversion, and loss_parallel() is precisely the max-then-sum parallel softmax — it consumes vocab-sharded logits and never materializes the full [N, V] tensor. This DTensor path is how torchtitan, PyTorch’s reference pretraining codebase, composes TP with FSDP2, PP, and CP on a single DeviceMesh; it is the practical alternative to adopting Megatron wholesale, and the trade-off between the two stacks is weighed in Megatron-LM, DeepSpeed & Parallelism in Practice.

Pipeline Parallelism: Splitting Across Layers¶

Tensor parallelism is bounded by the NVLink domain. To go bigger we split the model depth-wise: GPU 0 holds layers \(0..k\), GPU 1 holds layers \(k+1..2k\), and so on. A microbatch flows GPU 0 → GPU 1 → … → GPU \(p-1\) in the forward pass, and the gradients flow back. The only communication is point-to-point (send/recv of the activation tensor at the stage boundary) — cheap, and tolerant of slower inter-node links. This is pipeline parallelism (PP).

The problem is the bubble. If you naively run one whole batch through the pipeline, while stage 0 computes, stages \(1..p-1\) sit idle; while stage \(p-1\) computes, the rest sit idle. Utilization is \(1/p\) — catastrophic.

GPipe: Microbatching to Fill the Pipe¶

GPipe (Huang et al., 2019) fixes this by chopping the minibatch into \(m\) microbatches and streaming them. Once the pipe is full, multiple stages work in parallel on different microbatches. The schedule is “all-forward, then all-backward”:

The bubble is the fill + drain time. With \(p\) stages and \(m\) microbatches, the bubble fraction is

Increase \(m\) and the bubble shrinks. The rule: \(m \gg p\). With \(p=8\) and \(m=8\), the bubble is \(7/15 \approx 47\%\) — terrible. With \(m = 64\), it is \(7/71 \approx 10\%\). But GPipe has a memory problem: to do all forwards before any backward, it must stash the activations of all \(m\) in-flight microbatches on each stage. Larger \(m\) means lower bubble but higher peak activation memory — a direct tension. (Activation recomputation, from Memory-Efficient Training, is the usual escape valve.)

1F1B: The Steady-State Schedule¶

PipeDream’s 1F1B (“one-forward-one-backward”) schedule, adopted by Megatron, fixes the memory blowup. Once the pipeline is full, each stage alternates: do one forward, then one backward, then one forward, and so on. The key consequence: a stage only needs to keep activations for the microbatches currently in flight through it, which is at most \(p - s\) for stage \(s\) — bounded by \(p\), independent of \(m\).

The bubble fraction is the same \(\frac{p-1}{m+p-1}\) as GPipe — 1F1B’s win is memory, not bubble. It lets you run a large \(m\) (small bubble) without storing \(m\) microbatches’ activations. This is why 1F1B is the default in essentially every production framework.

# Minimal 1F1B driver (single-stage view). In reality each rank runs this with
# send/recv to its neighbors. `num_micro` = m, `stage` in [0, p-1], `p` = #stages.
def run_1f1b(stage, p, num_micro, fwd_step, bwd_step, recv_act, send_act,
             recv_grad, send_grad):
    warmup = p - stage - 1                 # how many forwards before first backward
    warmup = min(warmup, num_micro)
    steady = num_micro - warmup
    act_queue = []                         # activations awaiting their backward

    # ---- warmup: only forwards, prime the pipe ----
    for _ in range(warmup):
        x = recv_act() if stage > 0 else next_input()
        y, act = fwd_step(x)               # act = saved tensors for backward
        send_act(y) if stage < p - 1 else None
        act_queue.append(act)

    # ---- steady state: 1 forward then 1 backward, bounded memory ----
    for i in range(steady):
        x = recv_act() if stage > 0 else next_input()
        y, act = fwd_step(x)
        send_act(y) if stage < p - 1 else None
        act_queue.append(act)
        # immediately do a backward for the OLDEST in-flight microbatch
        g = recv_grad() if stage < p - 1 else loss_grad()
        gx = bwd_step(act_queue.pop(0), g)
        send_grad(gx) if stage > 0 else None

    # ---- cooldown: drain remaining backwards ----
    for _ in range(warmup):
        g = recv_grad() if stage < p - 1 else loss_grad()
        gx = bwd_step(act_queue.pop(0), g)
        send_grad(gx) if stage > 0 else None

Interleaved 1F1B (Virtual Pipeline Stages)¶

We can shrink the bubble without more microbatches. Interleaved 1F1B (Megatron-LM, Narayanan et al., 2021) gives each physical GPU several non-contiguous chunks of layers — “virtual stages.” With \(v\) virtual stages per device, the pipeline has \(p \cdot v\) logical stages, and the bubble shrinks to

A factor \(v\) improvement in the bubble. The cost is \(v\times\) more point-to-point communication (more, smaller sends) and a more intricate schedule. With fast intra-cluster links this is usually a great trade, and interleaving is standard in large Megatron runs.

Interleaved (v=2): each GPU owns TWO chunks. GPU0 = {layers 0-1, 8-9}, etc.
The pipe is "deeper" (more stages) so fill/drain is proportionally smaller,
at the price of more boundary send/recv ops.

Zero-Bubble and the Frontier¶

The bubble is fundamentally about the forward → backward dependency. Zero-Bubble Pipeline (Qi et al., 2024) observes that the backward pass actually splits into two pieces: the gradient w.r.t. the input (needed to keep the pipeline flowing upstream) and the gradient w.r.t. the weights (needed only before the optimizer step, and not on the critical path). By scheduling the weight-gradient computation into the bubbles, the bubble can be driven to near zero. DeepSeek-V3’s “DualPipe” pushes this further by overlapping forward and backward across a bidirectional pipeline and hiding communication. These are the current frontier — but plain interleaved 1F1B remains the workhorse you should reach for first.

The Library: torch.distributed.pipelining¶

You do not have to hand-roll the driver above either. PyTorch ships the whole schedule taxonomy as classes: give each rank its own nn.Module holding that stage’s layers, wrap it in a PipelineStage, hand that to a schedule object, and call step() once per global batch.

import torch, torch.nn.functional as F
from torch.distributed.pipelining import PipelineStage, Schedule1F1B

# `stage_mod` = this rank's contiguous block of layers, already .to(dev).
# `pp_group` = the process group spanning the p pipeline ranks.
stage = PipelineStage(stage_mod, stage_index=rank, num_stages=p,
                      device=dev, group=pp_group)
sched = Schedule1F1B(stage, n_microbatches=m, loss_fn=F.cross_entropy)

losses = []
if rank == 0:
    sched.step(x)                          # first stage supplies the inputs
elif rank == p - 1:
    sched.step(target=y, losses=losses)    # last stage supplies targets, collects losses
else:
    sched.step()                           # middle stages just relay activations/grads

The microbatch splitting, the send/recv pairs, and the warmup / steady-state / cooldown bookkeeping we wrote out by hand are all internal. Swapping schedules is a one-word change: ScheduleGPipe (all-forward-then-all-backward), Schedule1F1B, ScheduleInterleaved1F1B (pass a list of stages per rank for \(v > 1\)), ScheduleInterleavedZeroBubble / ScheduleZBVZeroBubble for the split-backward zero-bubble family, and ScheduleDualPipeV for the DeepSeek-style bidirectional schedule. Megatron-Core’s equivalent knobs are --pipeline-model-parallel-size and --num-layers-per-virtual-pipeline-stage; DeepSpeed’s is PipelineModule with LayerSpec.

Sequence & Context Parallelism: Splitting the Sequence¶

For very long context (32k, 128k, 1M tokens — see Long-Context Pretraining), the bottleneck is no longer parameters but activation memory and the \(O(s^2)\) attention computation, both scaling with sequence length \(s\). Neither TP nor PP addresses this directly — TP shards heads, PP shards layers, but every device still processes all \(s\) tokens of its slice. Context parallelism (CP) shards the token sequence itself across devices: GPU 0 owns tokens \(0..s/c - 1\), GPU 1 owns the next chunk, and so on, for a CP group of size \(c\).

The MLP and the projections are trivially sequence-parallel (they act per-token). The hard part is attention, where every query must attend to all keys/values — including those living on other devices.

Ring Attention¶

Ring Attention (Liu et al., 2023) solves this by combining the FlashAttention online-softmax trick with a communication ring. Each device starts with its own block of \(Q\), \(K\), \(V\). It computes the local attention contribution, then the \(K\)/\(V\) blocks are passed around a ring (device \(i \to i+1\)) so that, over \(c\) steps, every device sees every other device’s \(K\)/\(V\) — and crucially, the \(K\)/\(V\) for the next step is being sent while the current step is being computed, so communication hides under computation.

The online-softmax running statistics (the running max \(m\) and running sum \(\ell\), from FlashAttention) let each device incrementally fold in each incoming K/V block without ever materializing the full \(s \times s\) attention matrix:

# Ring Attention: each of c devices owns a contiguous query/key/value block.
# We rotate K,V around the ring; online softmax accumulates the result.
# `send_recv_ring(t)` sends t to rank+1 and returns the tensor from rank-1.
import torch, math

def ring_attention(q, k, v, cp_group_size, head_dim):
    # q,k,v: local blocks, shape [b, heads, s_local, d]
    scale = 1.0 / math.sqrt(head_dim)
    # running online-softmax state, FlashAttention-style
    out   = torch.zeros_like(q)                      # accumulated output
    l_run = torch.zeros(*q.shape[:-1], 1, device=q.device)   # running sum of exp
    m_run = torch.full((*q.shape[:-1], 1), -1e30, device=q.device)  # running max
    k_cur, v_cur = k, v
    for step in range(cp_group_size):
        # Begin sending current K,V to the next rank; overlaps with the matmul.
        k_next, v_next = send_recv_ring(k_cur), send_recv_ring(v_cur)
        s = torch.matmul(q, k_cur.transpose(-1, -2)) * scale   # [b,h,s_loc,s_loc]
        # (apply causal mask here if needed, accounting for block offsets)
        m_new = torch.maximum(m_run, s.max(dim=-1, keepdim=True).values)
        p = torch.exp(s - m_new)                     # rescaled exp
        corr = torch.exp(m_run - m_new)              # correction for old terms
        l_run = corr * l_run + p.sum(dim=-1, keepdim=True)
        out   = corr * out + torch.matmul(p, v_cur)  # fold this block's V
        m_run = m_new
        k_cur, v_cur = k_next, v_next                # rotate to next block
    return out / l_run                               # normalize at the very end

The communication per step is one \(K\) and one \(V\) block (\(\approx 2 \cdot \frac{s}{c} \cdot h \cdot 2\) bytes), and there are \(c\) steps — so total volume per device is \(O(s \cdot h)\), independent of \(c\), and it overlaps with the \(O(s^2/c)\) local compute. As \(s\) grows, compute dominates communication and Ring Attention scales to arbitrarily long sequences as long as you add devices.

In production you rarely write that loop yourself. PyTorch exposes it as torch.distributed.tensor.experimental.context_parallel, a context manager that shards the Q/K/V buffers over a CP mesh dimension and swaps in a ring-communicating SDPA implementation for the duration of the block; Megatron-Core exposes --context-parallel-size (fused with its FlashAttention kernels and the balanced token reordering below); and the community ring-flash-attention package wraps the real FlashAttention CUDA kernels in the same ring so you keep kernel-level IO-awareness and the sequence split.

Expert Parallelism: Scaling Mixture-of-Experts¶

A Mixture-of-Experts (MoE) layer (see Mixture-of-Experts Architectures) replaces the single FFN with \(E\) expert FFNs and a router that sends each token to its top-\(k\) experts. The whole point is that the active compute per token stays constant (only \(k\) of \(E\) experts fire) while the parameter count grows with \(E\). But that means the parameters are enormous — far too big to replicate. Expert parallelism (EP) places different experts on different devices: with an EP group of size \(e\), each device holds \(E/e\) experts.

Because the router sends tokens to experts that live on other devices, EP’s signature communication is a pair of all-to-all collectives:

1. Dispatch all-to-all: after routing, each device sends each token to the device holding its chosen expert(s).
2. (Each device runs its local experts on the tokens it received.)
3. Combine all-to-all: send the expert outputs back to the device that owns each token, where they are weighted by the router scores and summed.

# Sketch of one expert-parallel MoE layer. `ep_group` spans `e` ranks;
# this rank owns experts [rank*E_local : (rank+1)*E_local].
import torch, torch.distributed as dist, torch.nn.functional as F

def bucket_tokens_by_rank(x, dest_rank, e):
    # x: [T, h] (for k>1, each token is expanded to k rows before this call);
    # dest_rank: [T] long, the target EP rank per row.
    order = torch.argsort(dest_rank)
    send_buf = x[order]                                  # [T, h], grouped by dest rank
    counts = torch.bincount(dest_rank, minlength=e)       # [e], tokens destined for each rank
    return send_buf, counts, order                        # keep `order` to unsort after combine

def run_local_experts(recv_buf, experts_local, recv_counts, E_local):
    # recv_buf: [T_recv, h]. Production code re-buckets received tokens by which of the
    # E_local local experts they hit (each row carries its expert id); this sketch keeps
    # it simple and applies the single local expert when E_local == 1.
    assert E_local == 1
    return experts_local[0](recv_buf)                     # [T_recv, h]

def moe_forward(x, gate, experts_local, E, e, ep_group, k=1):
    # x: [tokens, h];  gate: Linear(h, E);  experts_local: list of FFNs on this rank
    logits = gate(x)                                   # [tokens, E]
    topk = logits.topk(k, dim=-1)                      # choose k experts/token
    probs = F.softmax(topk.values, dim=-1)             # routing weights, [tokens, k]
    expert_ids = topk.indices                          # [tokens, k]

    # Build send buffers: bucket tokens by the DEVICE that owns their expert.
    # (k=1 shown; for k>1, expand each token to k rows first, one per chosen expert.)
    E_local = E // e
    dest_rank = expert_ids[:, 0] // E_local             # [tokens], device owning this token (k=1)
    send_buf, counts, order = bucket_tokens_by_rank(x, dest_rank, e)   # counts: [e]

    # 1) EXCHANGE COUNTS first: each rank must learn how many tokens it will receive
    #    from every peer, since ranks receive different numbers of tokens (load imbalance)
    #    — this is the part a fixed-size torch.empty_like recv buffer silently gets wrong.
    recv_counts = torch.empty_like(counts)
    dist.all_to_all_single(recv_counts, counts, group=ep_group)        # [e]
    input_splits = counts.tolist()
    output_splits = recv_counts.tolist()

    # 2) DISPATCH: all-to-all with EXPLICIT variable split sizes, sized from the exchange.
    recv_buf = send_buf.new_empty((int(recv_counts.sum()), x.shape[-1]))   # [T_recv, h]
    dist.all_to_all_single(recv_buf, send_buf, output_splits, input_splits, group=ep_group)

    # 3) Run the local experts on received tokens.
    local_out = run_local_experts(recv_buf, experts_local, recv_counts, E_local)  # [T_recv, h]

    # 4) COMBINE: all-to-all back to each token's original owner — note the split sizes
    #    are SWAPPED relative to dispatch (we're now sending back what we received).
    combine_buf = local_out.new_empty((int(counts.sum()), x.shape[-1]))    # [T, h]
    dist.all_to_all_single(combine_buf, local_out, input_splits, output_splits, group=ep_group)

    # 5) Unsort back to original token order, then weight by router probs.
    out = torch.empty_like(combine_buf)
    out[order] = combine_buf                            # [T, h]
    out = out * probs.squeeze(-1).unsqueeze(-1)          # k=1: scale by the single routing weight
    # for k>1: repeat dispatch/combine per expert slot and torch.scatter_add_ the k
    # weighted contributions back onto the T original token rows.
    return out

The exchanged-counts-then-variable-split-size handling above is the actual hard part of expert parallelism — it is exactly what Megatron’s token dispatcher (megatron.core.transformer.moe) and DeepSpeed/Tutel implement, since a naive fixed-size recv buffer is only correct when every rank happens to receive the same number of tokens, which routing essentially never guarantees. Two open-source libraries are worth knowing at this layer: DeepEP (DeepSeek’s expert-parallel communication library) supplies NVLink- and RDMA-tuned dispatch/combine kernels, including a low-latency path for MoE inference; and MegaBlocks reformulates the expert FFNs as block-sparse GEMMs, giving dropless MoE — no capacity factor, no dropped tokens — at the cost of variable-shaped kernels.

The deciding factor for EP performance is load balance. If the router sends most tokens to a few popular experts, those devices become stragglers while others idle, and the all-to-all is dominated by the heaviest bucket. This is why MoE training relies on an auxiliary load-balancing loss (or DeepSeek-style auxiliary-loss-free bias correction) to spread tokens evenly, and on expert capacity limits that drop or reroute overflow tokens. EP is almost always combined with TP and DP, and the all-to-all collectives are extremely bandwidth-sensitive — keep the EP group on fast links, and overlap dispatch/combine with the attention compute of the next layer where possible.

Combining Everything: 3D, 4D & 5D Parallelism¶

No single axis suffices at frontier scale. Real systems compose them. The total number of GPUs is the product:

Each GPU belongs to one group per axis. The orchestration trick is mapping these groups onto the physical network topology so that the chattiest collectives ride the fastest links. The canonical ordering, fastest-comm axis innermost:

This is why you see configs like “TP=8 (within a node), PP=12 (across nodes), DP=16 (across racks).” TP is locked inside the NVLink domain; PP and DP span the InfiniBand/Ethernet fabric where their lighter, less frequent communication is tolerable.

A useful way to think about it: TP and PP both reduce per-device memory and let you fit a bigger model; DP buys throughput; CP buys context length; EP buys parameter count. You first pick TP/PP/CP/EP large enough that one model replica fits and trains efficiently, then set DP to consume the remaining GPUs for throughput.

torch.testing.assert_close(loss_tp2, loss_ref, rtol=1e-3, atol=1e-3)
torch.testing.assert_close(logits_tp2, logits_ref, rtol=1e-3, atol=1e-3)

Where the Communication Lives: A Summary Table¶

The single most important thing to internalize is what collective each axis pays, how often, and over which tensor. This is what lets you reason about whether a config will be fast on your hardware.

Notice the frequency column. TP and EP pay per layer, so they demand the fastest links and the tightest placement. DP pays per step (after gradient accumulation over all microbatches), so it tolerates the slow outer fabric. PP’s communication is cheap point-to-point but introduces the bubble, a compute-utilization tax rather than a bandwidth one. CP’s communication overlaps with compute and so is nearly hidden when sequences are long. This table — internalized — is most of what you need to design or debug a large training run, and it connects directly to the collective communication primitives and the practical framework details in Megatron-LM, DeepSpeed & Parallelism in Practice.

Key Takeaways¶

Further reading¶

• Shoeybi, Patwary, Puri, et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019) — the column/row tensor-parallel formulation.
• Narayanan, Shoeybi, Casper, et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021) — interleaved 1F1B and the 3D-parallelism analysis.
• Korthikanti, Casper, Lym, et al., Reducing Activation Recomputation in Large Transformer Models (2022) — Megatron sequence parallelism and selective recomputation.
• Huang, Cheng, Bapna, et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (2019).
• Narayanan, Harlap, Phanishayee, et al., PipeDream: Generalized Pipeline Parallelism for DNN Training (2019) — the 1F1B schedule.
• Qi, Wan, Huang, et al., Zero Bubble Pipeline Parallelism (2024).
• Liu, Zaharia, Abbeel, Ring Attention with Blockwise Transformers for Near-Infinite Context (2023).
• Lepikhin, Lee, Xu, et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020) — expert parallelism and all-to-all dispatch/combine.
• Rajbhandari, Rasley, Ruwase, He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2020) — the DP-side counterpart to this chapter.
• The Megatron-LM and DeepSpeed open-source repositories — the reference implementations of every schedule discussed here.

Exercises¶

1. In the Megatron MLP block \(Y = \text{GeLU}(XA)B\), the up-projection \(A\) is column-parallel and the down-projection \(B\) is row-parallel. Explain why the order matters: what goes wrong if you instead make \(A\) row-parallel and \(B\) column-parallel? In particular, what extra collective would the GeLU force, and why?

2. A pipeline has \(p = 16\) stages. (a) Using GPipe/1F1B, how many microbatches \(m\) do you need to keep the bubble fraction at or below 10%? (b) If you switch to interleaved 1F1B with \(v = 4\) virtual stages per device and keep that same \(m\), what is the new bubble fraction? © The chapter says 1F1B and GPipe have the same bubble but 1F1B is preferred — what exactly does 1F1B buy you, and how does the peak activation memory scale with \(m\) in each?

3. Consider the chapter’s 70B model (\(L=80\), \(h=8192\), \(a=64\)) but now on a cluster of 256 H100s (80 GB), arranged as 32 nodes x 8 GPUs. Static training state is \(\approx 16 \times 70\text{B} = 1120\) GB. (a) With TP=8 and PP=4, how much static state sits on each GPU, and how many transformer layers does each pipeline stage own? (b) How many data-parallel replicas \(d\) does this leave, and what is the DP group size? © With \(m = 32\) microbatches, what is the plain-1F1B bubble fraction?

4. Estimate the tensor-parallel communication volume for one transformer layer forward pass. Use \(s = 8192\), \(h = 8192\), activations in bf16 (2 bytes), and TP degree \(t = 8\) with ring all-reduce. (a) How many bytes does each GPU send+receive per all-reduce, using the chapter’s \(2\cdot\frac{t-1}{t}\cdot(s\,h\cdot 2)\) estimate? (b) A layer forward does two all-reduces; what is the per-GPU per-layer forward volume? © Why does this quantity, multiplied out over 80 layers, force the TP group onto NVLink rather than InfiniBand?

5. The chapter’s VocabParallelEmbedding masks out-of-range token ids, zeros the corresponding output rows, and then does an all_reduce(SUM). (a) Explain precisely why a SUM all-reduce yields the correct embedding, and why exactly one rank contributes a nonzero row per token. (b) Implement a gather_output option for ColumnParallelLinear (from the chapter’s code) that, when True, returns the full unsharded output by all-gathering the feature-sharded slices along the last dimension. Give runnable code consistent with the chapter’s style.

class ColumnParallelLinear(nn.Module):
    """Y = X A, A split along output columns across `tp` ranks.
    If gather_output=True, all-gather the shards to return the full [*, out_f]."""
    def __init__(self, in_f, out_f, tp, rank, tp_group, bias=True,
                 gather_output=False):
        super().__init__()
        assert out_f % tp == 0
        self.tp = tp
        self.tp_group = tp_group
        self.gather_output = gather_output
        self.out_local = out_f // tp
        self.weight = nn.Parameter(torch.empty(self.out_local, in_f))
        self.bias = nn.Parameter(torch.zeros(self.out_local)) if bias else None
        nn.init.normal_(self.weight, std=0.02)

    def forward(self, x):
        x = copy_to_region(x)                              # f: correct bwd all-reduce
        y = torch.nn.functional.linear(x, self.weight, self.bias)  # [*, out_f/tp]
        if not self.gather_output:
            return y                                       # sharded (feeds a RowParallel)
        # all-gather the per-rank feature shards, then concat along last dim
        shards = [torch.empty_like(y) for _ in range(self.tp)]
        dist.all_gather(shards, y.contiguous(), group=self.tp_group)
        return torch.cat(shards, dim=-1)                   # [*, out_f], full output