Choosing a mesh for training

Chapter 5 priced each collective. This chapter adds them up for a training step, using the Scaling Book's model of a transformer layer as two matmuls, $\text{In}[B, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]$, with $B$ the global batch in tokens [SB5]. For each scheme the question is the same: does the math per layer take longer than the communication it needs? Each answer is an inequality in one number, the chip's ICI intensity:

FLOP/s per byte/s of one torus axis. The number of FLOPs a chip must do for every byte it moves over ICI to stay compute-bound.

Computed in scripts/ch06_training.py. The book's derivations use v5p at 2,550. This chapter keeps that for the book's examples and re-runs them for Ironwood.

Data parallelism and FSDP

Data parallelism keeps a full copy of the weights on every chip and splits the batch. The forward pass needs no communication. The backward pass AllReduces the weight gradients, and those AllReduces can overlap with the rest of the backward pass [SB5]. Per layer:

The batch per chip must exceed the ICI intensity, divided by the number of torus axes the data-parallel dimension spans.

This is the network version of the 240-token rule from Chapter 1. On v5p it is 2,550 tokens per chip on one axis, or 850 when the data axis spans all three torus axes [SB5]. Plain data parallelism also needs the full model plus optimizer on each chip: 10 bytes per parameter, so at most about 9B parameters on a 96 GB v5p chip.

FSDP (ZeRO-3) shards weights, gradients and optimizer state along the data axis. It AllGathers each layer's weights just before use and ReduceScatters gradients after. An AllReduce costs a ReduceScatter plus an AllGather, so FSDP moves the same bytes as data parallelism and has the same condition, 6.2, while removing the memory limit [SB5]. The book's conclusion: "we can basically always do ZeRO-3".

On Ironwood, condition 6.2 is 11,535 tokens per chip on one axis and 3,845 across three derived.

Tensor parallelism

Tensor parallelism (Megatron) shards the weights along F, AllGathers activations before $W_\text{in}$ and ReduceScatters them after $W_\text{out}$. These collectives sit on the critical path: the next layer cannot start without them [SB5]. The condition no longer involves the batch:

Tensor parallelism stays compute-bound only up to F/α ways per torus axis it spans, whatever the batch.

For LLaMA-3-70B, $F = 28{,}672$: up to about 11-way per axis on v5p, which in practice means 8-way. On Ironwood the same formula allows about 2.5-way per axis derived. Quantizing to int8 or fp8 does not help: it halves the bytes and doubles $C$, so α per byte moved is unchanged [SB5].

FSDP and tensor parallelism together

FSDP moves weights; tensor parallelism moves activations. Combining them on an $X \times Y$ mesh trades one against the other [SB5]:

More FSDP (larger X) means larger TP collectives on fewer chips; more TP means weight gathers spread over fewer chips. The best split balances the two.

Setting the two communication terms equal gives the optimal FSDP degree, and requiring math to exceed communication there gives the minimum batch per chip:

The optimal FSDP degree grows with the square root of batch and chip count. Mixed FSDP and TP needs far fewer tokens per chip than FSDP alone, because the requirement now scales with α² / F.

On v5p, with $F = 28{,}672$ and $M_X M_Y = 2$, the threshold is about 113 tokens per chip, roughly an eighth of FSDP alone across three axes.

Case study: LLaMA-3-70B on a v5p pod

The Scaling Book's Chapter 6 works this through [SB6]: a 4M-token batch (1,024 sequences of 4,096), 15T tokens, one v5p pod of 8,960 chips.

• Time: $6 \cdot 70 \times 10^9 \cdot 15 \times 10^{12} = 6.3 \times 10^{24}$ FLOPs; at 40% MFU, 44 days. The script gets 45 with 70.5B parameters.
• Memory is not the constraint: about 1.3 GB per chip spread over the pod.
• Pure FSDP over batch: at most 1,024-way, since there are only 1,024 sequences. Not enough chips.
• FSDP over batch and sequence: 468 tokens per chip, below the 850 needed. Communication-bound.
• FSDP plus TP: the threshold is 113 tokens per chip and there are 468, so it is compute-bound. $X_\text{opt} \approx 1{,}618$, so about 2,048-way FSDP × 4-way TP.

The book's summary: on a full v5p pod, a mixture of 1,024-way data parallelism, 2-way sequence parallelism and 4-way tensor parallelism avoids being communication-bound.

The same job on an Ironwood pod

Run the same numbers with Ironwood's α of 11,535 derived:

• 4,194,304 tokens over 9,216 chips is 455 tokens per chip.
• FSDP plus TP needs $11{,}535^2 / (2 \cdot 28{,}672) \approx$ 2,320 tokens per chip.
• At the best split, math is about 0.44 of communication: the step is communication-bound by about 2×.
• The largest slice on which this batch stays compute-bound is about 1,800 chips. There, the run takes about 44 days, the same as a full v5p pod.
• Raising the batch to 16M tokens, as LLaMA-3 did later in training [P-Llama3, §3.4], makes about 7,200 chips compute-bound.

The cause is the ratio from Chapter 1: Ironwood has 5× v5p's FLOPs per chip and the same ICI bandwidth per link [D-7x]. The model is simplified. It ignores overlap between collectives and other layers' compute, and gated MLPs change the constants. The direction is robust: on Ironwood, dense training wants larger batches per chip, fewer chips per job, or less ICI traffic per FLOP. Google's TPU 8t doubles ICI bandwidth relative to Ironwood [G-8dd]. With the provisional BF16 figure, that brings α back to about 2,490, v5p's level.

Pipelining

Pipeline parallelism splits the model by layers across stages and passes activations between them. Per layer it moves only $[B, D]$ activations, once, between adjacent stages, so its communication is small. The costs are the pipeline bubble and the complexity of interleaved schedules. The Scaling Book calls pipelining "a dominant strategy for GPU parallelism that is somewhat less essential on TPUs", because a TPU pod already provides thousands of chips at full ICI bandwidth [SB5]. In GSPMD, pipelining is expressed as sharding over a stage axis [P-GSPMD, §3.3], and MaxText has a stage mesh axis for it [R-maxtext].

On Ironwood the trade shifts. Pipelining's small communication is exactly what a high α favours, so expect pipeline stages to be more attractive there than on v5p derived.

Across pods: DCN

Beyond one slice, chips talk over DCN. The book's recipe is model parallelism or FSDP within a slice, and pure data parallelism across slices [SB5]. Gradients cross DCN once per step, and the condition is on the batch per slice:

Each slice's share of the batch must exceed the chip's FLOPs per byte of DCN bandwidth.

On v5p, at 6.25 GB/s of DCN per chip, that is 73,440 tokens per slice; the book rounds to "about 71,000" in places. On Ironwood, at 12.5 GB/s, it is about 185,000 tokens per slice derived. These are easy to meet: a 4M-token batch over 4 pods is 1M per pod. The Scaling Book's worked problem 1, with no published answer, scales LLaMA-3-70B to 4 v5p pods at the same batch. 1M tokens per pod clears 73,440. 117 tokens per chip just clears the FSDP-plus-TP threshold of 113. The run takes about 11 days derived.

MaxText encodes the whole chapter in two lines of its default config. FSDP takes all remaining devices inside a slice (ici_fsdp_parallelism: -1) and data parallelism takes all slices (dcn_data_parallelism: -1). Its logical_axis_rules then map named tensor dimensions such as embed, mlp and heads onto those mesh axes [R-maxtext].