Splitting a Model Across Devices

Data, tensor and pipeline parallelism compared by the one quantity that decides between them — bytes communicated per FLOP of useful work.

AI SystemsGraduate⏱ 3 hours9 min readdata paralleltensor parallelpipeline parallelcommunicationbubble

AssumesThe Cost of Collective Communication · The Memory Arithmetic of Training

Intuition

There are only three things you can cut when a model will not fit on one device: the batch, the individual matrices, or the stack of layers. Cutting the batch means every device holds the whole model and they must agree on the gradient. Cutting the matrices means a single layer's arithmetic is spread across devices, so they must exchange activations inside every layer. Cutting the layer stack means each device owns a contiguous slice of depth, and activations flow between neighbours.

Each cut produces a different communication pattern, and — this is the useful part — each is amortised by a different quantity. Data parallelism gets cheaper with more tokens per device. Tensor parallelism gets cheaper with a wider model. Pipeline parallelism gets cheaper with more microbatches. That is why large training runs use all three at once rather than picking a favourite.

Counting rules

Throughout, let $N$ be the parameter count, $d$ the model width, $L$ the layer count, $F$ the achievable FLOP rate per device, $B$ the achievable interconnect bandwidth per device, and $b$ the bytes per element communicated. For a decoder-only transformer we use the standard counts established in Inference Arithmetic and the Roofline:

parameters per layer $\approx 12d^2$ (attention $4d^2$ , MLP $8d^2$ ), so $N \approx 12Ld^2$ ;
training cost $\approx 6N$ FLOPs per token (forward $2N$ , backward $4N$ ).

Data parallelism

Every device holds a full replica and processes a different slice of the batch. After the backward pass the gradients must be summed across replicas — one all-reduce of $N$ elements per step.

PropositionData-parallel communication does not depend on model size

Let each device process $T$ tokens per optimiser step. Using a ring all-reduce, whose per-device volume is at most $2Nb$ bytes,

\frac{T_{\text{comm}}}{T_{\text{compute}}} \;=\; \frac{2Nb/B}{6NT/F} \;=\; \frac{b\,F}{3\,B\,T} .

The parameter count cancels. The ratio depends only on the hardware ratio $F/B$ , the gradient precision $b$ , and the tokens per device per step.

Assumptions

The all-reduce is bandwidth-bound, so its latency term $2(p-1)\alpha$ is negligible. For a multi-gigabyte gradient this holds; for small models it does not, and the latency term restores a dependence on $p$ .
Communication is not overlapped with computation. Overlapping changes the conclusion from "this is the cost" to "this is the amount that must be hidden", which is the same number.

Key idea

Data parallelism is amortised by tokens per device per step. Doubling the per-device microbatch halves the relative cost of gradient synchronisation. This is why data-parallel training degrades at small per-device batch — and why gradient accumulation, which raises $T$ without raising memory for activations proportionally, is the standard fix.

Tensor parallelism

Each weight matrix is sharded across $t$ devices, so a single layer's matmuls are computed jointly. The Megatron decomposition requires one all-reduce of the activation tensor after the attention block and one after the MLP block, in the forward pass, and the same in the backward pass.

PropositionTensor-parallel communication does not depend on batch size

With $T$ tokens in flight, a hidden size $d$ , and $4$ all-reduces per layer per step (two forward, two backward) each of $Td$ elements,

\frac{T_{\text{comm}}}{T_{\text{compute}}} \;=\; \frac{4L\cdot 2Tdb/B}{6\,(12Ld^2)\,T/F} \;=\; \frac{8LTdb/B}{72Ld^2T/F} \;=\; \frac{b\,F}{9\,d\,B} .

Both $L$ and $T$ cancel. The ratio falls like $1/d$ : wider models tolerate tensor parallelism better.

Pipeline parallelism

The layer stack is cut into $p$ contiguous stages, one per device. A batch is split into $m$ microbatches which flow through the stages; only activations at stage boundaries are communicated, which is a small point-to-point transfer rather than a collective. The cost is not bandwidth but idleness: at the start, stages downstream have nothing to do, and at the end, upstream stages are finished.

PropositionThe pipeline bubble fraction

With $p$ stages, $m$ microbatches, and equal per-stage time $\tau$ per microbatch, the schedule completes in $(m+p-1)\tau$ against an ideal $m\tau$ . The fraction of wall-clock time lost to idle stages is

\text{bubble} \;=\; \frac{p-1}{m+p-1},

and the overhead relative to the ideal is $(p-1)/m$ .

With $p = 4$ and $m = 8$ the bubble is $3/11 \approx 27\%$ ; raising $m$ to $32$ brings it to $3/35 \approx 8.6\%$ .

Key idea

Pipeline parallelism is amortised by microbatch count. But $m$ cannot grow freely: every in-flight microbatch holds its activations for the backward pass, so memory grows with $m$ . The bubble and the activation memory pull in opposite directions, and that trade is what interleaved and 1F1B schedules exist to improve.

Why all three

CorollaryWhy the three are combined rather than chosen between

The three overheads are amortised by three different quantities:

\text{data} \;\propto\; \frac{1}{T}, \qquad \text{tensor} \;\propto\; \frac{1}{d}, \qquad \text{pipeline} \;\propto\; \frac{p-1}{m} .

Since $d$ is fixed by the model and $T$ and $m$ are bounded by memory, no single axis can be pushed to absorb an arbitrary device count; minimising total overhead subject to a fixed device budget $p_{\text{total}} = p_{\text{data}} \cdot p_{\text{tensor}} \cdot p_{\text{pipe}}$ therefore has an interior optimum in general.

What this does not say

This page counts bytes and FLOPs. It says nothing about convergence: a configuration that minimises time per step can still be worse per unit of progress, because the global batch size it implies may sit past the point where larger batches stop helping. Step time and steps-to-target are separate quantities, and a parallelism plan that optimises only the first is incomplete.

References

PaperM. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, B. Catanzaro (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model ParallelismSource of the tensor-parallel decomposition whose communication is counted here.
PaperY. Huang et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline ParallelismThe microbatched pipeline schedule whose bubble fraction is derived here.
PaperD. Narayanan et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMCombines the three strategies and reports the interaction effects discussed at the end.
PaperS. Rajbhandari, J. Rasley, O. Ruwase, Y. He (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter ModelsShards optimiser state, gradients and parameters across data-parallel ranks.

These works are cited as sources to read, not reproduced. All exposition, examples and problems on this site are written independently; where an argument is standard, the citation points to a canonical treatment.

Results stated on this page

PropositionData-parallel communication does not depend on model sizeproof includedThe ratio of gradient-synchronisation time to compute time is independent of the parameter count and falls inversely with tokens processed per device per step.
PropositionTensor-parallel communication does not depend on batch sizeproof includedThe ratio of activation-exchange time to compute time is independent of layer count and sequence length, and falls inversely with model width.
PropositionThe pipeline bubble fractionproof includedWith p stages and m microbatches, the fraction of wall-clock time in which some stage is idle is (p−1)/(m+p−1).
CorollaryWhy the three are combined rather than chosen betweenproof includedEach strategy is amortised by a different quantity, so the cheapest configuration mixes them rather than maximising any one.