Cutting Up a Model

Imagine a book too long for one person to copy out. You could hand each copyist the whole book and a different stack of blank pages - everyone does the same work on different material. You could give each copyist a different chapter, and pass the manuscript down the line. You could put four people on a single page, each writing one column. Or, for a sentence that runs on for pages, you could split that one sentence between them.

Those four arrangements are the four ways a model gets split, and a real training run uses all of them at once. What makes it interesting is that they cost very different amounts of conversation between the people doing the work - and conversation is the expensive part.

A frontier model does not fit in one GPU's 192 GBHBM3e per Blackwell GPU 180 GB usable after ECC. NVIDIA of HBM, and would take an unusable amount of time even if it did. So the work is divided along four axes at once, and they are not interchangeable.

Data parallelismData parallelismEvery GPU holds the whole model and a different slice of the batch; gradients are all-reduced each step. gives every GPU the whole model and a different slice of the batch. Communication is one gradient all-reduce per step - cheap, infrequent, and the reason this dimension scales furthest. Tensor parallelismTensor parallelismOne layer’s matrices are split across GPUs, so a single matrix multiply is executed by several GPUs together. splits individual weight matrices, so several GPUs cooperate on one matrix multiply. It communicates twice per transformer block, on the critical path. Pipeline parallelismPipeline parallelismConsecutive layers live on different GPUs, and micro-batches flow through them like a factory line. gives consecutive layers to different GPUs and streams micro-batches through them; it sends one activation tensor per stage boundary. Expert parallelismexpert parallelismSpreading an MoE model’s experts across many GPUs so each holds only a few. distributes an MoEMoEMixture of Experts - a model where each token is routed to a small subset of many parallel sub-networks. model's experts, and pays for it in all-to-all traffic.

Llama 3 405B used four dimensions written as [TP, CP, PP, DP], where the third of those is context parallelismContext parallelismThe sequence itself is split across GPUs, so one very long context can be processed by several at once. - splitting the sequence itself, which appears once context length outgrows per-GPU memory.^{Meta AI}

The distinction that matters operationally is communication volume per unit of compute, and whether that communication sits on the critical path. Tensor parallelism exchanges partial sums inside every layer and nothing can proceed until they land. Pipeline parallelism exchanges one boundary activation per micro-batch and the next micro-batch can start meanwhile. That gap - not bandwidth alone - is what assigns dimensions to tiers.

The binding constraint is memory. Weights are the smallest part of it. Training with Adam keeps, per parameter, a 32-bit master copy and two 32-bit moment estimates - twelve bytes - on top of two bytes of BF16 weights and two of gradients. Optimiser state alone is roughly six times the size of the weights it is optimising.

Tensor and pipeline parallelism divide all of it, because each GPU physically holds fewer parameters. Data parallelism does not: every replica holds the whole model. That is what sharding the optimiser stateSharded optimiser stateSplitting optimiser state, gradients and parameters across data-parallel ranks instead of replicating them. across the data-parallel group fixes, turning a per-GPU cost into a per-group one and gathering each piece when its owner needs it.

Which leaves activations, and they are the awkward term: they scale with batch size and sequence length rather than with parameters, and they are why two runs of the same model on the same hardware can have entirely different memory profiles. Activation recomputationActivation recomputationThrowing away intermediate activations during the forward pass and recomputing them during the backward pass. is the usual answer - discard them on the forward pass and recompute them on the backward - trading roughly 30% more arithmetic for a large reduction in peak memory. It is worth noticing that this makes the MFUMFUModel FLOPs Utilisation - the share of a machine’s peak arithmetic that ends up doing the model’s own maths. of the next chapter look worse while making the run possible at all.

Set the tensor-parallel degree to 72 in the planner and the memory per GPU collapses - which looks like an argument for always doing it. It is not. Tensor parallelism at degree 72 means every layer's partial sums are reduced across all 72 GPUs, twice per block, with the whole job waiting. The rack makes that survivable in a way eight GPUs behind a slower link never did, but it does not make it free.

What the rack changes is the ceiling. A tensor-parallel group has to live inside one NVLink domain, so before NVL72 the degree was capped at 8 by the server. Now it is capped at 72Blackwell GPUs per rack NVIDIA by the rack. The planner shows the group turning red the moment the degree exceeds what one domain holds, because at that point the chattiest dimension in the model is crossing the tier where bandwidth drops from 900 GB/sNVLink egress per GPU, one direction Half of the 1.8 TB/s bidirectional headline. This is the figure to compare against a scale-out line rate, because a NIC’s quoted rate is also one direction. NVIDIA to 100 GB/sConnectX-8 egress per GPU 800 Gb/s. Still 9× less than in-rack NVLink egress. NVIDIA .