Sequence, context and expert parallelism

Before this lesson: the online-softmax merge (lesson 4), activation memory $34sbh$ (lesson 6), reduce-scatter, all-gather and all-to-all (lesson 7), tensor parallelism (lesson 9). Attention with heads (deep learning 9); mixture of experts is defined again here (more in deep learning 15). Time: about 80 minutes, plus the notebook.

---

Why this matters #

Activation memory grows with the sequence length $s$. Take one Llama-2-7B-like layer ($h = 4096$), micro-batch $b = 1$, FlashAttention, and $s = 131{,}072$ tokens. The layer stores about $34sbh \approx 18.3$ GB of activations (lesson 6). 32 layers need 584 GB. Tensor parallelism with sequence parallelism over 8 GPUs still leaves 73 GB per GPU, almost all of an 80 GB H100. Long context needs one more axis of parallelism: split the sequence itself. The same problem appears in science: a global 0.25° weather grid has $721 \times 1440 \approx 1.04$ million points, so even $4 \times 4$ patches give about 65,000 tokens per field and time step.

Mixture-of-experts models have the opposite problem: many parameters, few FLOPs per token. Mixtral 8x7B has about 47B parameters but uses about 13B per token. DeepSeek-V3 has 671B parameters and activates 37B per token. The weights do not fit on one GPU, so the experts are spread over many GPUs, and every MoE layer moves every token to its experts and back.

1. Predict first #

2. The picture #

Only attention mixes tokens #

Look at one transformer layer and ask, for each operation: does the output for token $i$ need any other token?

Operation	Needs other tokens?
LayerNorm / RMSNorm, dropout, residual add	no: each token alone
Q, K, V and output projections, MLP	no: the same matrix for each token
MoE router and experts	no: each token chooses its experts alone
attention $\operatorname{softmax}(QK^\top/\sqrt{d_h})V$	yes: query $i$ needs all visible keys $j$

Simplest case: 2 GPUs, 4 tokens. GPU 0 holds tokens 0–1 and GPU 1 holds tokens 2–3. All the token-wise operations run with no communication. For attention there are three ways out, and they are the three methods of this lesson:

• Megatron sequence parallelism (SP). Keep tensor parallelism (TP) for attention and the MLP (lesson 9). Shard along the sequence only the regions that TP leaves replicated (a full copy on every GPU).
• Ring attention / context parallelism (CP). Shard everything along the sequence. Each GPU keeps its queries, and the key/value (K/V) blocks visit every GPU in turn.
• DeepSpeed-Ulysses. Shard along the sequence, but just before attention switch to sharding by heads: each GPU gets all tokens for a subset of heads.

Method	What is sharded	What moves	Collective
Megatron SP	norm/dropout regions by tokens; attention and MLP by heads/columns (TP)	activations between regions	reduce-scatter + all-gather (inside the TP group)
Ring attention (CP)	everything by tokens	K and V blocks	point-to-point sends around a ring
Ulysses	everything by tokens, attention by heads	Q, K, V and the output	two all-to-alls

Ring attention: queries stay, keys and values travel #

Cut the sequence into $N_c$ blocks of $c = s/N_c$ tokens; rank $r$ holds block $r$ of $Q$, $K$ and $V$. Attention for the queries of rank $r$ needs every K/V block. Instead of collecting them, the ranks form a ring. In each round, every rank computes attention of its queries against the K/V block it holds now, and at the same time sends that block to the next rank and receives a new one from the previous rank. After $N_c$ rounds every query has seen every key. The partial results are merged with the logsumexp formula of lesson 4, so the result is exact.

With a causal mask, query $i$ sees only keys $j \le i$. Then the order of the tokens on the ranks matters a lot, as the figure shows.

Interactive figure: loading…

3. Definitions and main results #

Megatron sequence parallelism #

With TP alone, the activation memory per layer is (Korthikanti et al. 2022, Eq. 2)

The $10\,sbh$ term is not divided by $t$. It comes from the replicated regions, which every GPU of the TP group stores in full:

Stored tensor	Bytes
inputs of the two LayerNorms	$2 \times 2sbh = 4sbh$
LayerNorm outputs = inputs of the QKV matmul and of the first MLP matmul	$2 \times 2sbh = 4sbh$
the two dropout masks (1 byte per element)	$2 \times sbh = 2sbh$
total	$10\,sbh$

These regions are token-wise, so SP shards them along $s$: GPU $i$ keeps only tokens $i\frac{s}{t}, \dots, (i+1)\frac{s}{t} - 1$. Two operators convert between the regions (Korthikanti et al., Fig. 5–6):

• $g$: all-gather along $s$ in the forward pass, before the column-parallel matmul (which needs all tokens); reduce-scatter in the backward pass;
• $\bar g$: reduce-scatter along $s$ in the forward pass, after the row-parallel matmul; all-gather in the backward pass.

Interactive figure: loading…

Claim: SP costs no extra bandwidth.

1. TP needs one all-reduce after each row-parallel matmul: two per layer in the forward pass.
2. A ring all-reduce of $n$ bytes over $t$ GPUs is a reduce-scatter followed by an all-gather (lesson 7). Each half sends $\frac{t-1}{t}n$ bytes per GPU.
3. SP does the reduce-scatter half right after the matmul. GPU $i$ now holds the summed output for its $s/t$ tokens only.
4. Dropout, the residual add and LayerNorm run on these $s/t$ tokens. They are token-wise, so the result equals the replicated result restricted to these tokens.
5. The all-gather half runs just before the next column-parallel matmul.
6. So per layer and pass, SP runs 2 reduce-scatters and 2 all-gathers instead of 2 all-reduces: the same $2 \cdot 2\frac{t-1}{t}n$ bytes. $\blacksquare$

Memory. Now every activation is divided by $t$ (Korthikanti et al., Eq. 4):

One detail makes this exact: SP stores only the local shard of each LayerNorm output and repeats the all-gather in the backward pass, overlapped with the gradient computation. With FlashAttention the $5as/h$ term is not stored (lesson 4), and the memory is about $34\,sbh/t$.

Limits. SP lives inside the TP group, so it scales only as far as TP: usually one NVLink node of 8 GPUs. And every GPU still runs attention over all $s$ tokens.

Ring attention #

Each of the $N_c$ ranks holds a block of $c = s/N_c$ queries and the matching K/V block. In round $k = 0, \dots, N_c - 1$:

1. rank $r$ computes attention of its queries against the K/V block of rank $(r - k) \bmod N_c$;
2. it merges this partial result into its running result;
3. at the same time it sends the K/V block to rank $r + 1$ and receives the next one from rank $r - 1$ (no send in the last round).

After $N_c - 1$ shifts every query has seen every key.

Why the merge is exact. For one query, let $o_A$ and $\mathrm{lse}_A = \log \sum_{j \in A} e^{x_j}$ be the output and the logsumexp over a key set $A$ (scores $x_j$). For a disjoint key set $B$:

(lesson 4 proves it). The result is the exact softmax over $A \cup B$, and the merge is associative and commutative. So the ring may visit the blocks in any order.

Theorem (minimum block size). For multi-head attention in bf16, ring attention hides the K/V transfer behind the attention compute when

Proof.

1. In one round, a rank computes $QK^\top$ for $c$ queries and $c$ keys over all heads: $2c^2 h$ FLOPs (each of $a$ heads does $2c^2 d_h$, and $a d_h = h$). Then $PV$: another $2c^2 h$. Total $4c^2 h$ FLOPs, time $4c^2 h / F$.
2. In the same round it sends one K block and one V block: $2 \cdot c \cdot h$ elements of 2 bytes, so $4ch$ bytes, time $4ch / W$.
3. Sending and computing run at the same time. The transfer is hidden if $4c^2 h / F \ge 4ch / W$.
4. Divide both sides by $4ch$: $c \ge F/W$. $\blacksquare$

The bound does not depend on $h$ or on the number of ranks. Three refinements matter in practice:

• GQA. K and V have only $n_{kv} d_h = h/g$ columns, so the bytes shrink to $4ch/g$ while the FLOPs stay $4c^2 h$: $c \ge F/(gW)$.
• Causal mask. With a balanced layout (below), each round does only half a block of work, $2c^2 h$ FLOPs, so $c \ge 2F/W$.
• Backward pass. dK and dV travel with K and V (about twice the bytes), and the attention backward does about 2.5 times the forward FLOPs (five matmuls instead of two), so the condition is about the same. For other formats, multiply the bound by $e/2$.

Memory per rank: the Ring Attention paper counts six blocks per host (the query block, the current K and V blocks, the K and V blocks being received, and the output). This is linear in $c$ and does not depend on $s$.

Interactive figure: loading…

Causal load balancing #

With a causal mask and contiguous blocks, rank $r$ needs only K/V blocks $0, \dots, r$. Count work in units of one full $c \times c$ block, where a diagonal block costs $\tfrac12$. Rank 0 does $\tfrac12$ block; rank $N_c - 1$ does $N_c - \tfrac12$ blocks. All ranks move in lockstep, so the ring takes $N_c - \tfrac12$ block-times, while the average load is $N_c/2$: about 50% efficiency.

Zig-zag layout. Cut the sequence into $2N_c$ chunks of $c/2$ tokens. Rank $r$ gets chunk $r$ (early) and chunk $2N_c - 1 - r$ (late).

Proposition. With the zig-zag layout, every rank does the same work in every round, so the ring takes $N_c/2$ block-times: the optimum.

Proof. Measure work in chunk pairs. Rank $r$ holds query chunks $r$ and $2N_c - 1 - r$. In some round it holds the K/V chunks $j$ and $2N_c - 1 - j$ of rank $j \ne r$. Check the four chunk pairs:

1. queries $r$, keys $j$: fully visible if $j < r$, fully masked if $j > r$;
2. queries $r$, keys $2N_c - 1 - j$: always masked, because $2N_c - 1 - j \ge N_c > r$;
3. queries $2N_c - 1 - r$, keys $j$: always visible, because $j < N_c \le 2N_c - 1 - r$;
4. queries $2N_c - 1 - r$, keys $2N_c - 1 - j$: visible if $j > r$, masked if $j < r$.

Exactly one of pairs 1 and 4 is visible, and pair 3 is visible. So every rank does 2 chunk pairs ($= \tfrac12$ block) in every round with $j \ne r$. In round 0 ($j = r$) the rank does the two diagonal chunk pairs and pair 3: $\tfrac12 + \tfrac12 + 1 = 2$ chunk pairs as well. All ranks are equal in all rounds, so no rank waits. $\blacksquare$

Llama 3 uses exactly this split: "the i-th CP rank received both the i-th and the (2×CP−1−i)-th chunks" (Llama 3 paper, §3.3.2).

DeepSpeed-Ulysses #

Activations are sharded along $s$. Before attention, one all-to-all on Q, K and V switches to head sharding: rank $r$ receives all $s$ tokens for $a/N_c$ heads. Attention then runs locally with any kernel (FlashAttention, causal or sparse masks), because each head needs nothing from the other heads. A second all-to-all on the output switches back to sequence sharding.

An all-to-all is a block transpose (lesson 7): think of the tensor as an $N_c \times N_c$ grid of blocks (token blocks × head groups). Before, rank $r$ holds row $r$ of the grid; after, it holds column $r$. Each rank keeps 1 block and sends $N_c - 1$ blocks.

Interactive figure: loading…

Bytes per GPU and layer (forward, $b = 1$, bf16, multi-head attention):

Method	Bytes sent per GPU	Overlapped with compute?
Megatron SP (degree $t$)	$4 \cdot \frac{t-1}{t} \cdot 2sh \approx 8sh$	partly
Ring attention (degree $N_c$)	$(N_c - 1) \cdot 4ch = 4sh\,\frac{N_c-1}{N_c} \approx 4sh$	yes, if $c \ge F/W$
Ulysses (degree $N_c$)	$4 \cdot \frac{N_c-1}{N_c} \cdot \frac{2sh}{N_c} \approx 8sh/N_c$	no

Where the Ulysses line comes from: each rank holds $s/N_c$ tokens of Q, K, V and the output, $2sh/N_c$ bytes each; an all-to-all sends $\frac{N_c - 1}{N_c}$ of it; there are four such tensors. With GQA, ring attention sends $g$ times fewer bytes, and Ulysses' K and V terms shrink by $g$.

For multi-head attention, Ulysses sends $N_c/2$ times fewer bytes than ring attention, and its bytes per GPU stay constant when $s$ and $N_c$ grow together (the term $s/N_c$ does not change). Its limits:

• Head-count limit. Each rank needs whole heads, so $N_c$ cannot exceed the number of heads. With GQA, K and V have only $n_{kv}$ heads, so $N_c \le n_{kv}$ unless K/V heads are copied.
• Exposed communication. The all-to-all sits on the critical path; attention cannot start before it ends. Across nodes, with slow links, this hurts.

Ring attention has no head limit (it needs only $c \ge F/W$ to hide its transfers) and works across nodes. The two combine: Ulysses inside a node, ring attention across nodes.

Mixture of experts and expert parallelism #

Parameters vs active parameters. Each token uses only $k$ of the $E$ experts. So the FLOPs per token follow the active parameters (the $k$ chosen experts plus attention), and the memory follows all parameters. Mixtral 8x7B: about 47B parameters, about 13B active per token.

Expert parallelism (EP). Put $E/N_e$ experts on each of $N_e$ GPUs. The attention part stays data-parallel: each GPU holds a different set of $T$ tokens. An MoE layer then runs in four steps:

1. route: each GPU runs the router on its own $T$ tokens;
2. dispatch (all-to-all): each token vector is sent to the GPUs that hold its $k$ experts;
3. expert compute: each GPU runs its experts on the tokens it received;
4. combine (all-to-all): the expert outputs go back to the GPU that holds the token, which forms the weighted sum $y$.

Proposition (all-to-all bytes). With uniform routing, one copy of the token vector per (token, expert) pair, and $e$ bytes per element, each GPU sends per MoE layer in the forward pass

Proof.

1. A GPU holds $T$ tokens, and each token has $k$ chosen experts: $kT$ (token, expert) pairs.
2. Each pair sends one token vector of $h$ elements, $h e$ bytes, to the GPU of its expert.
3. With uniform routing, the expert lies on any of the $N_e$ GPUs with probability $1/N_e$. With probability $1/N_e$ it is on the same GPU, and nothing is sent. So the expected bytes of the dispatch are $kT \cdot he \cdot \frac{N_e - 1}{N_e}$.
4. The combine returns one output vector of $h$ elements per pair, on the same paths in reverse: the same number of bytes.
5. Add the two. $\blacksquare$

The notebook checks this with a simulation. Compare it with the expert compute. Each pair costs $2 \cdot 2h\,d_{ff}$ FLOPs in a two-matrix MLP with inner width $d_{ff}$ ($3 \cdot 2h\,d_{ff}$ with SwiGLU, which has three matrices). For the two-matrix MLP in bf16, the time ratio is

So the all-to-all is cheap only when $d_{ff}$ is large compared with $F/W$. Over NVLink ($F/W \approx 2{,}200$ at peak) this is often true; over InfiniBand ($F/W \approx 20{,}000$) it is not. So EP is cheapest inside a node. When it must cross nodes, systems limit how many nodes a token may visit and overlap the all-to-all with other work (see the folded box below).

The busiest GPU sets the time. Routing is learned, so it is not uniform. The all-to-all and the next layer wait for the slowest GPU. If one expert gets twice the mean load, its GPU computes twice as long, and every other GPU waits. Two tools fight this:

• Capacity factor. Each expert accepts at most $C = \mathrm{CF} \cdot \frac{kN}{E}$ tokens, where $N$ is the number of tokens routed together and $\frac{kN}{E}$ is the mean load (Switch Transformer writes it for $k = 1$). A token over the capacity is dropped for that expert: its update skips the expert and the token passes on through the residual connection. $\mathrm{CF} > 1$ leaves room for imbalance but costs memory and padding.
• Load-balancing loss (Switch Transformer). Add $\mathcal{L}_\text{aux} = \alpha\, E \sum_{i=1}^{E} f_i P_i$ to the loss, where $f_i$ is the fraction of tokens sent to expert $i$ and $P_i$ the mean router probability of expert $i$. It equals $\alpha$ for uniform routing and grows when a few experts get many tokens and high probability; the paper uses $\alpha = 10^{-2}$.

Interactive figure: loading…

Beyond the basic recipe

• Dropless MoE. MegaBlocks (Gale et al. 2022) computes all experts as one block-sparse matmul, so no token is dropped and no padding to the capacity is needed.
• DeepSeek-V3 (2024) trains with 64-way EP spread over 8 nodes. Each token is sent to at most 4 nodes, and custom all-to-all kernels use InfiniBand and NVLink together. It balances the load without a large auxiliary loss: a bias per expert is added to the routing scores only, and after each step it is decreased for overloaded experts and increased for underloaded ones.
• With $d_{ff} = 2048$ (DeepSeek-V3's experts) and InfiniBand, the ratio above is far larger than 1. That is why such systems limit how many nodes a token may visit and overlap the all-to-all with the computation of other micro-batches.

Putting the axes together #

• Order from inner to outer. Llama 3 405B orders its axes as [TP, CP, PP, DP]: TP (with SP) inside the node, where the bandwidth is highest, then CP, then PP, then DP. For the 131,072-token stage it used TP 8, CP 16, PP 16 and DP 8 on 16,384 GPUs (Llama 3 paper, Table 4).
• Gradients. CP ranks hold the same weights but different tokens, so their weight gradients must be summed: Megatron reduces gradients over the DP × CP group. Under SP, the LayerNorm weight gradients are partial (each GPU saw $s/t$ tokens) and are all-reduced over the TP group.
• EP and DP. Non-expert weights are replicated over the data-parallel GPUs; expert weights are split over the EP group. So expert gradients are reduced only over the copies of each expert.

Answers to section 1. (1) It stays the same: an all-reduce is a reduce-scatter plus an all-gather. (2) About 20,000: $c \ge F/W = 989 \cdot 10^{12} / 50 \cdot 10^{9} \approx 19{,}780$ tokens for multi-head attention, and about 4,900 with GQA at $g = 4$. (3) 8 times the parameters, but only about 2 times the FLOPs per token (2 experts run, plus a tiny router). Memory follows all parameters; compute follows the active ones.

4. Worked example #

(a) SP on a toy #

$t = 2$, 4 tokens, $h = 1$. After the row-parallel matmul, GPU 0 holds the partial sums $(1, 2, 3, 4)$ and GPU 1 holds $(10, 20, 30, 40)$.

• TP: all-reduce. Both GPUs get $(11, 22, 33, 44)$, and both store the dropout and LayerNorm activations for 4 tokens. Sent per GPU: $2 \cdot \frac{1}{2} \cdot 4 = 4$ elements.
• SP: reduce-scatter. GPU 0 sends $(3, 4)$ and gets $(11, 22)$; GPU 1 sends $(10, 20)$ and gets $(33, 44)$. Each runs dropout and LayerNorm on 2 tokens and stores 2 tokens. The all-gather before the next matmul sends 2 more elements each. Total: 4 elements, as before.

(b) SP memory for a 7B-like model #

$h = 4096$, $s = 8192$, $b = 1$, $t = 8$, FlashAttention, so $sbh = 33{,}554{,}432$:

• no parallelism: $34\,sbh = 1.14$ GB per layer;
• TP: $(10 + 24/8)\,sbh = 13\,sbh = 436$ MB per layer, 14.0 GB for 32 layers;
• TP + SP: $34\,sbh/8 = 4.25\,sbh = 143$ MB per layer, 4.6 GB for 32 layers.

At $t = 8$ the replicated part is $10/13 = 77\%$ of the TP memory, so SP saves a factor of 3.1.

(c) Causal ring attention with zig-zag, every block #

$N_c = 4$ ranks, $s = 8$ tokens, causal mask. We count work in (query, key) pairs with key ≤ query. The whole sequence has $8 \cdot 9 / 2 = 36$ such pairs, so the best possible is $36/4 = 9$ per rank.

Zig-zag. $2N_c = 8$ chunks of one token. Rank $r$ holds tokens $r$ and $7 - r$: rank 0 has $\{0, 7\}$, rank 1 $\{1, 6\}$, rank 2 $\{2, 5\}$, rank 3 $\{3, 4\}$. In round $k$, rank $r$ holds the K/V tokens of rank $(r - k) \bmod 4$.

Round	rank 0: Q $\{0,7\}$	rank 1: Q $\{1,6\}$	rank 2: Q $\{2,5\}$	rank 3: Q $\{3,4\}$	round time
0	K/V $\{0,7\}$: (0,0) (7,0) (7,7) = 3	$\{1,6\}$: (1,1) (6,1) (6,6) = 3	$\{2,5\}$: (2,2) (5,2) (5,5) = 3	$\{3,4\}$: (3,3) (4,3) (4,4) = 3	3
1	$\{3,4\}$: (7,3) (7,4) = 2	$\{0,7\}$: (1,0) (6,0) = 2	$\{1,6\}$: (2,1) (5,1) = 2	$\{2,5\}$: (3,2) (4,2) = 2	2
2	$\{2,5\}$: (7,2) (7,5) = 2	$\{3,4\}$: (6,3) (6,4) = 2	$\{0,7\}$: (2,0) (5,0) = 2	$\{1,6\}$: (3,1) (4,1) = 2	2
3	$\{1,6\}$: (7,1) (7,6) = 2	$\{2,5\}$: (6,2) (6,5) = 2	$\{3,4\}$: (5,3) (5,4) = 2	$\{0,7\}$: (3,0) (4,0) = 2	2
total	9	9	9	9	9

Each entry lists the visible pairs (query, key). Every rank does 9 pairs, and the ring takes $3 + 2 + 2 + 2 = 9$ pair-times: the optimum.

Contiguous. Rank $r$ holds tokens $2r$ and $2r + 1$. Against its own block a rank does 3 pairs, against an earlier block 4 pairs, against a later block 0. The work per rank is $3, 7, 11, 15$. Rank 3 holds an earlier block in rounds 1, 2 and 3, so the ring takes $3 + 4 + 4 + 4 = 15$ pair-times. Zig-zag is $15/9 \approx 1.67$ times faster here. For long blocks (where the diagonal costs half a block) the ratio is $(N_c - \tfrac12)/(N_c/2) = 3.5/2 = 1.75$ for $N_c = 4$.

The merge for one query. Rank 0's query 7 meets keys $\{0, 7\}$, then $\{3, 4\}$, then $\{2, 5\}$, then $\{1, 6\}$. After each round it merges the new $(o, \mathrm{lse})$ into its running pair. The order differs from $0, 1, \dots, 7$, but the merge is exact in any order, so the result equals full causal attention (the notebook checks this to $10^{-12}$).

(d) Ring attention at 128K tokens #

$s = 131{,}072$ on $N_c = 8$ H100s, so $c = 16{,}384$; $h = 4096$, multi-head attention, at peak.

• Compute per round: $4c^2h = 4.40 \times 10^{12}$ FLOPs, 4.45 ms at 989 TFLOP/s.
• Transfer per round: $4ch = 268$ MB. Over NVLink (450 GB/s): 0.60 ms, hidden. Over a 50 GB/s network: 5.37 ms, not hidden, as $c = 16{,}384 < F/W = 19{,}780$ predicts.
• With GQA at $g = 4$: 67 MB, 1.34 ms, hidden again.
• Per layer, ring attention sends $7 \times 268 = 1.88$ GB per GPU. Ulysses would send $4 \cdot \frac{7}{8} \cdot \frac{2 \cdot 131{,}072 \cdot 4096}{8} = 470$ MB, $N_c/2 = 4$ times less, but on the critical path.

(e) One MoE layer #

A Mixtral-like layer: $h = 4096$, $E = 8$ experts on $N_e = 8$ GPUs, top-2, SwiGLU experts with $d_{ff} = 14{,}336$, $T = 8192$ tokens per GPU, bf16, uniform routing, at peak.

• All-to-all bytes per GPU: $2 \cdot 2 \cdot 8192 \cdot 4096 \cdot 2 \cdot \frac78 = 235$ MB. Over NVLink: 0.52 ms. Over InfiniBand: 4.70 ms.
• Expert FLOPs per GPU: $kT = 16{,}384$ copies, each $6h\,d_{ff} = 352$ MFLOP: $5.77 \times 10^{12}$ FLOPs, 5.84 ms.
• Ratio: 0.09 over NVLink, 0.80 over InfiniBand. Check with the formula (SwiGLU has 1.5 times the FLOPs): $\frac{19{,}780}{1.5 \cdot 14{,}336} \cdot \frac78 = 0.80$. Across nodes, the all-to-all costs almost as much as the experts.

5. In code #

NumPy checks of sections 3 and 4: sequence parallelism, the causal work per round, the minimum block size, and the MoE bytes.

import numpy as np

rng = np.random.default_rng(0)

# (a) Sequence parallelism: reduce-scatter + token-wise op + all-gather
t, s, h = 4, 8, 16                          # TP degree, tokens, hidden size
partials = rng.standard_normal((t, s, h))   # row-parallel matmul output on each GPU

def layer_norm(x):                          # token-wise: each row is normalised alone
    mu, var = x.mean(-1, keepdims=True), x.var(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + 1e-5)

tp = layer_norm(partials.sum(0))            # TP: all-reduce, then every GPU runs LN on all s tokens
shards = [partials[:, i*s//t:(i+1)*s//t].sum(0) for i in range(t)]   # reduce-scatter
sp = np.concatenate([layer_norm(x) for x in shards])                  # LN on s/t tokens, all-gather
print(np.abs(sp - tp).max())                # 0.0

n = s * h * 2                               # bytes of the activation in bf16
print(2 * (t-1)/t * n, (t-1)/t * n + (t-1)/t * n)   # all-reduce vs RS + AG: 384.0 384.0

# (b) Causal work per round, token level (s = 8 tokens, N_c = 4 ranks)
Nc, s = 4, 8
layouts = {"contiguous": [[2*r, 2*r+1] for r in range(Nc)],
           "zig-zag":    [[r, 2*Nc-1-r] for r in range(Nc)]}
for name, tok in layouts.items():
    # round k: rank r holds the K/V tokens of rank (r - k) mod N_c; count pairs with key <= query
    w = np.array([[sum(kk <= q for q in tok[r] for kk in tok[(r-k) % Nc]) for r in range(Nc)]
                  for k in range(Nc)])
    print(f"{name:10s} per rank {w.sum(0)}  per round max {w.max(1)}  ring time {w.max(1).sum()}")
# contiguous per rank [ 3  7 11 15]  per round max [3 4 4 4]  ring time 15
# zig-zag    per rank [9 9 9 9]  per round max [3 2 2 2]  ring time 9

# (c) Minimum block size c >= F / W (bf16, multi-head attention)
F = 989e12                                  # H100 dense bf16 peak, FLOP/s
for name, W in [("NVLink", 450e9), ("InfiniBand", 50e9)]:
    print(name, round(F / W))               # NVLink 2198, InfiniBand 19780

# (d) MoE dispatch: bytes per GPU, simulated vs formula
Ne, E, k, T, h, e = 8, 64, 2, 4096, 1024, 2
experts_per_gpu = E // Ne
sent = []
for g in range(Ne):                          # tokens of GPU g pick k distinct experts uniformly
    choice = np.argsort(rng.random((T, E)), axis=1)[:, :k]
    sent.append((choice // experts_per_gpu != g).sum() * h * e)
print(2 * np.mean(sent) / 1e6, 2 * k * T * h * e * (Ne-1) / Ne / 1e6)   # MB: dispatch + combine
# 29.425152 29.360128

The simulated bytes differ from the formula by 0.2%: the formula is the expected value, and the simulation draws random experts.

6. Where it shows up #

• Long-context LLM training. Llama 3 trains up to 128K tokens with CP (Llama 3 paper, §3.3.2). It uses the zig-zag split, but it all-gathers K and V instead of passing them around a ring: with GQA, K and V are much smaller than Q, and an all-gather makes masks such as document masks easy. The exposed all-gather is $O(s)$, while attention is $O(s^2)$.
• Mixture-of-experts models. Switch Transformer, GShard, Mixtral and DeepSeek-V3 all use expert parallelism with dispatch and combine all-to-alls. Load balancing (a loss, a capacity factor, or a routing bias) is part of every recipe.
• ML for science. Global weather and climate grids, long DNA sequences, long seismic records and video all give $10^4$–$10^6$ tokens. Ring attention is domain decomposition for attention: a stencil code exchanges only halos with its neighbours, but global attention needs every block, so the blocks travel around the ring.
• Inference. The same ideas split the prefill of very long prompts and the attention of video diffusion transformers. MoE serving also spreads the experts over GPUs with EP.

7. Common mistakes #

• "SP adds communication." It replaces each all-reduce by its two halves. The bytes are the same.
• Forgetting the gradient sums. Under SP each GPU sees only $s/t$ tokens, so the LayerNorm weight gradients are partial and must be all-reduced over the TP group. Under CP, weight gradients must be summed over the CP ranks too.
• Contiguous causal sharding. It wastes about half of the compute. Use zig-zag (or a striped layout).
• Using $c \ge F/W$ blindly. With a causal mask the work per round halves ($c \ge 2F/W$); with GQA the bytes shrink by $g$; and $F$ is the attention FLOP/s you really achieve, not the peak.
• Ulysses with more GPUs than K/V heads. The degree is limited by the head count; beyond it, combine with ring attention.
• "An MoE model with 8 experts costs 8 times the FLOPs." FLOPs follow the $k$ active experts; memory follows all $E$.
• Ignoring load imbalance. The all-to-all waits for the busiest GPU. Measure the max load, not the mean.

8. Check yourself #

1. Under TP + SP, how many collectives run per layer in one forward and one backward pass?
   Answer

   Four all-gathers and four reduce-scatters (TP alone: four all-reduces), with the same bytes. Megatron adds one extra all-gather per stored LayerNorm-output shard in the backward pass, overlapped with the gradient computation.
2. $h = 5120$, $s = 4096$, $b = 1$, $t = 4$, FlashAttention. Give the activation memory per layer with TP, and with TP + SP.
   Answer

   $sbh = 20{,}971{,}520$. TP: $(10 + 24/4)\,sbh = 16\,sbh \approx 336$ MB. TP + SP: $34\,sbh/4 = 8.5\,sbh \approx 178$ MB.
3. Why does the minimum block size double with a causal mask and a zig-zag layout?
   Answer

   Each round a rank computes only half a $c \times c$ block ($2c^2h$ FLOPs), but it still sends full K and V blocks ($4ch$ bytes): $2c^2h/F \ge 4ch/W \iff c \ge 2F/W$.
4. Why can the ring visit the K/V blocks in any order?
   Answer

   The logsumexp merge gives the exact softmax over the union of disjoint key sets, and it is associative and commutative.
5. A model has 32 query heads and 8 K/V heads. What is the largest Ulysses degree without copying K/V heads, and what can you use beyond it?
   Answer

   8: each rank needs whole K/V heads. Beyond it, combine Ulysses with ring attention (ring across groups of 8), or copy K/V heads at a higher communication cost.
6. EP over $N_e = 4$ GPUs, $T = 1024$ tokens per GPU, top-1, $h = 2048$, bf16, uniform routing. How many bytes does each GPU send per MoE layer in the forward pass?
   Answer

   $2 \cdot 1 \cdot 1024 \cdot 2048 \cdot 2 \cdot \frac34 = 6{,}291{,}456$ bytes $\approx 6.3$ MB: dispatch and combine, each $\frac34$ of $1024 \cdot 4096$ bytes.

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. Which regions does Megatron sequence parallelism shard, and along which axis?

   Show answer

2. Why does sequence parallelism send the same bytes as plain TP?

   Show answer

3. Activation memory per layer with TP alone, and with TP + SP?

   Show answer

4. State the minimum block size for ring attention and its derivation in one line.

   Show answer

5. Ring attention on H100s: $F = 989$ TFLOP/s, $W = 50$ GB/s, multi-head attention. Minimum tokens per rank?

   Show answer

6. In the zig-zag layout with $N_c$ ranks, which chunks does rank $r$ hold?

   Show answer

7. What do the two all-to-alls of DeepSpeed-Ulysses do?

   Show answer

8. What limits the Ulysses degree $N_c$?

   Show answer

9. Bytes per GPU of the MoE all-to-alls in one forward pass, with uniform routing?

   Show answer

10. What does the capacity factor control in an MoE layer?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How would you train with a 1M-token context?" FlashAttention on each GPU, TP + SP inside the node, ring attention (CP) across nodes with a zig-zag layout; check $c \ge F/W$ (or $2F/(gW)$ with a causal mask and GQA). Activation memory per GPU falls as $1/N_c$.
• "SP vs CP?" SP shards only the norm/dropout regions inside the TP group and keeps the bytes of TP. CP shards all activations, attention included, and works across nodes.
• "Ring attention vs Ulysses?" Ring: point-to-point, overlapped, no head limit, needs $s/N_c \ge F/W$, more bytes. Ulysses: two all-to-alls, $N_c/2$ times fewer bytes (multi-head attention), exposed, $N_c \le$ heads. Combine: Ulysses in the node, ring across nodes.
• "Why is MoE training hard to scale?" Two all-to-alls per MoE layer, with $2kTh\,e\,(N_e - 1)/N_e$ bytes per GPU; slow across nodes unless $d_{ff} \gg F/W$; the busiest expert sets the time, so you need capacity factors, a balance loss or a routing bias, and you trade dropped tokens against padding.
• "What is the difference between total and active parameters?" Total sets the memory (all experts are stored); active sets the FLOPs per token (only $k$ experts run).

10. Practice #

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

1. SP check. Reduce-scatter, a token-wise op and all-gather on a simulated TP group give the same result as all-reduce and the op, with the same bytes.
2. Ring attention. Four simulated ranks pass K/V blocks around a ring and merge partial results with the logsumexp formula; check against full (causal) attention.
3. Load balance. Work per rank and per round for the contiguous and zig-zag layouts, and the ring time, for several $N_c$.
4. Ulysses. An all-to-all reshard on NumPy arrays (sequence → heads → sequence), attention per head group, checked against full attention; count the bytes.
5. MoE dispatch. A top-$k$ routing simulator with a capacity factor: dropped-token rate, all-to-all bytes and the max/mean load as the router skew grows.

11. Watch and read #

Lecture 13: Ring AttentionGPU MODE · 76 min

How ring attention passes K/V blocks between devices and merges the partial softmax results, with code.

Stanford CS25: V1 I Mixture of Experts (MoE) paradigm and the Switch TransformerStanford Online · CS25 · 66 min

The authors of Switch Transformer explain top-1 routing, the capacity factor and the load-balancing loss.

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient SparsityYannic Kilcher · 34 min

A paper walk-through; watch for how experts are placed on devices and what happens to dropped tokens.

• Korthikanti et al. (2022), Reducing Activation Recomputation in Large Transformer Models: §4.1 (the $34sbh$ count), §4.2.1 (Eq. 2), §4.2.2 (Fig. 5–6, Eqs. 3–4, sequence parallelism).
• Liu, Zaharia and Abbeel (2023), Ring Attention with Blockwise Transformers for Near-Infinite Context: §3 (arithmetic intensity between hosts, memory requirement).
• Jacobs et al. (2023), DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models: §3.1 (design), §3.2 (communication analysis).
• Llama Team (2024), The Llama 3 Herd of Models: §3.3.2 (context parallelism, zig-zag chunks, all-gather of K/V) and Table 4.
• Fedus, Zoph and Shazeer (2021), Switch Transformers: §2.1–2.2 (routing, capacity factor, load-balancing loss).
• Lepikhin et al. (2020), GShard: top-2 gating and expert placement across devices.
• Jiang et al. (2024), Mixtral of Experts.
• DeepSeek-AI (2024), DeepSeek-V3 Technical Report: §2.1.2 (auxiliary-loss-free load balancing), §3.2 (expert parallelism and all-to-all kernels).
• Gale et al. (2022), MegaBlocks: Efficient Sparse Training with Mixture-of-Experts.
• Hugging Face, The Ultra-Scale Playbook: sections on sequence parallelism, context parallelism (ring and zig-zag) and expert parallelism.