Model, Tensor & Pipeline Parallelism
When the model itself does not fit on one device, you cut the model rather than the data — across layers, inside matrix multiplies, or across the optimizer state. Every cut moves activations or weights over the wire, and the wire becomes the bottleneck.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
The model's weights, gradients and optimizer state exceed one device's memory — which way do you cut it, and what does each cut cost in communication?
A team fine-tuning a large language model for their support domain cannot fit it on any single GPU they can rent: the weights alone fill the card before the optimizer state and activations are counted. They have eight cards on one node and are trying to work out what "split the model" means in practice.
Put the first half of the layers on GPU 0 and the second half on GPU 1. Activations flow forward across the boundary, gradients flow back. Two cards, twice the memory, problem solved.
GPU 1 idles while GPU 0 runs the forward pass of the first half, then GPU 0 idles while GPU 1 finishes and starts the backward pass. With two stages half of the hardware is idle at any time; the run fits, and it is nearly as slow as one card.
- GPU 1 idles while GPU 0 runs the forward pass of the first half, then GPU 0 idles while GPU 1 finishes and starts the backward pass. With two stages half of the hardware is idle at any time; the run fits, and it is nearly as slow as one card.
- Memory did not halve. Each card still holds its layers' weights, gradients and optimizer state, and the optimizer state was the largest term. The split moved the problem without shrinking it.
- Attempting the same split across two nodes moves every layer-boundary activation over the inter-node network on every step in both directions. Communication time exceeds compute time and the run is slower than a single card would have been if it had fit.
- The learning rate and batch that worked for the fitted small variant of the model are reused; the run diverges, and the team spends a week discovering that the parallelism was fine and the schedule was not.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The surrounding task is next-token prediction on domain text, fine-tuned from a pretrained checkpoint (Fine-Tuning). The target of this lesson is to make that training run *exist* on hardware where the naive layout does not fit.
- Fitting is a memory budget: parameters, gradients, optimizer state (two extra copies per parameter for Adam) and activations, all per device. The split is chosen to bring each device's share under its memory.
- One example is a sequence of a few thousand tokens. Activations per layer scale with sequence length times hidden width times batch, and for a deep model at long sequence they can exceed the weights.
- The pretrained checkpoint is tens of gigabytes of FP16 weights on object storage. In training, each weight needs a gradient and two optimizer moments, and the moments are usually kept in FP32 — roughly six times the FP16 weight footprint before activations.
- The node has eight GPUs on a fast intra-node interconnect; a second node would be reachable only over a much slower network.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- Pipeline parallelism places consecutive layers on consecutive devices. Only activations at the stage boundaries cross the wire — small relative to weights — but the stages run sequentially, so the naive version leaves every device but one idle. Splitting the batch into micro-batches and streaming them through lets stages overlap; the idle time that remains at the start and end of each batch is the *pipeline bubble*, and its fraction shrinks as micro-batch count grows relative to stage count.
- Tensor parallelism splits a single layer's matrix multiply across devices — each holds a column or row slice of the weight and computes a partial product, and the partials are combined with a collective operation. It cuts per-device weight memory for the largest layers and needs an exchange *inside* every layer, so it only works over a very fast interconnect, in practice within one node.
- Sharded optimizer state (the ZeRO-style family) keeps data parallelism's shape — every device runs the whole model on its shard of data — but partitions the optimizer moments, and optionally the gradients and weights, across devices, gathering each piece just in time. It is the first thing to try because it changes the memory budget without changing the compute pattern, at the cost of extra all-gather traffic.
- All three move bytes over the interconnect every step, and the bytes moved scale differently: pipeline with activation size at boundaries, tensor with activation size at every layer, sharded state with parameter count. On an interconnect with a fixed bandwidth, the strategy that moves the fewest bytes per unit of compute wins, and the answer changes with the model shape and the hardware (Memory Bandwidth & VRAM).
The memory budget that decides the cut
A parameter in training is not one number. In mixed precision it is an FP16 weight, an FP16 gradient, an FP32 master copy and two FP32 optimizer moments — sixteen bytes per parameter before any activation is stored. Activations add a term that scales with batch, sequence length and depth, and for long sequences it is the largest of all.
Whichever term dominates decides the first cut. Optimizer state is removed from each device by sharding it; activations by checkpointing and recomputing them; weights, last, by tensor or pipeline splits. Doing the arithmetic first is what stops a team from pipelining across nodes to solve a problem that sharding the optimizer would have solved on one.
| Term | Scales with | Bytes per parameter (mixed precision) | First remedy |
|---|---|---|---|
| Weights (FP16) | parameter count | 2 | tensor / pipeline split |
| Gradients (FP16) | parameter count | 2 | shard across data-parallel workers |
| Master weights (FP32) | parameter count | 4 | shard across data-parallel workers |
| Optimizer moments (FP32 × 2) | parameter count | 8 | shard across data-parallel workers |
| Activations | batch × sequence × width × depth | not per parameter | activation checkpointing, smaller micro-batch |
Three cuts, three traffic patterns
Pipeline parallelism moves the smallest thing — activations at a handful of stage boundaries — but leaves devices idle in the bubble. Tensor parallelism moves activations inside every layer, which is only tolerable on the fastest link in the system. Sharded state moves parameters and moments as they are gathered, on the data-parallel pattern, and keeps every device computing the full model.
The combination that works is dictated by the hardware topology: tensor groups within a node where the link is fast, pipeline stages across nodes where it is slow, data parallelism with sharded state across the replicas. The layout is a map of the interconnect as much as of the model.
- 1Split the batch
Divide the batch into M micro-batches so stages can overlap on different micro-batches.
fails by M too small relative to the stage count: the bubble — start-up and drain when only some stages have work — dominates the step.
- 2Forward through stages
Stage k runs its layers on micro-batch j while stage k−1 runs micro-batch j+1; boundary activations are sent onward.
fails by Unbalanced stages: the stage with the heaviest layers paces every other one.
- 3Backward through stages
Gradients flow back across the same boundaries in reverse; each stage accumulates gradients over all micro-batches.
fails by Stored activations for all in-flight micro-batches exceed memory; needs activation checkpointing.
- 4Update
After all micro-batches, every stage applies its accumulated gradient — one optimizer step for the whole batch.
fails by Effective batch is the full batch, not the micro-batch; the schedule must be tuned for it.
The bubble fraction is roughly (stages − 1) / (micro-batches + stages − 1). Eight stages with eight micro-batches idle nearly half the hardware; with sixty-four micro-batches, under a tenth. The micro-batch count is bounded by activation memory, which is why the two problems are solved together.
Every layer's partial products are combined over the slow network on every forward and backward pass; the step is dominated by waiting.
The per-layer exchanges stay on the intra-node interconnect; only boundary activations cross between nodes, a few times per micro-batch.
The strategies differ in bytes moved per layer, and the interconnects differ by an order of magnitude in bandwidth. Matching the chattiest split to the fastest link is the entire design.
What must remain true about the layout
A split model is a promise that the pieces compute the same function the whole would have. Nothing enforces that promise: a wrong slice boundary, a partial sum combined in the wrong order, or a mixed-precision cast at a stage boundary all produce a model that trains — to something — and the loss curve does not say which.
The layout is also a promise about the hardware. A tensor-parallel group placed across nodes by a scheduler that did not know the difference runs, slowly, and the slowness looks like a large model being large.
The split forward and backward pass computes the same function and gradients as the unsplit model, on hardware whose topology matches the split's design.
holds when A small configuration of the architecture has been checked split-against-unsplit to floating-point tolerance; placement constraints pin tensor-parallel groups to one node; the mixed-precision policy is identical at and across boundaries.
breaks when A code change moves a slice boundary; the scheduler places a group across nodes; a new layer type is added without a tensor-parallel implementation and silently falls back to replication.
respond Do not tune the model. Fix the layout, re-run the equivalence test, then profile again.
How to build it
Most important first.
- Do the memory arithmetic before choosing anything: parameters × (weight bytes + gradient bytes + optimizer bytes) + activations at the intended batch and sequence length. The term that dominates decides the first cut.
- Reduce before splitting. Mixed precision halves weights and activations; activation checkpointing trades recomputation for activation memory; a parameter-efficient method may make the trainable state tiny (Parameter-Efficient Fine-Tuning). Many "does not fit" problems fit after this step.
- Then shard the optimizer state across the data-parallel workers. It keeps the training loop's shape and removes the largest term.
- If the weights still do not fit, tensor-parallelise the largest layers within a node, and pipeline across nodes — the fast interconnect carries the intra-layer traffic, the slow one carries only boundary activations.
- Choose micro-batch count to shrink the bubble, and re-tune the schedule for the resulting effective batch, exactly as in Data Parallelism.
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Peak memory per device against capacity, per term: weights, gradients, optimizer state, activations. The breakdown decides the next cut; a single "out of memory" does not.
- Device idle fraction per step — the pipeline bubble plus time blocked on collectives. This is the number that says whether the split is using the hardware or merely fitting on it.
- Bytes moved per step per device, against the interconnect's measured bandwidth. When it approaches the compute time per step, the split has reached its limit on this hardware.
- Do not use aggregate GPU utilisation. A device blocked in a collective reports busy.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The per-device memory budget — weights, gradients, optimizer state and activations at the chosen batch and sequence length — stays within capacity, including the transient buffers the collectives allocate.
- The interconnect topology the split was designed for is the one the run gets: tensor-parallel groups stay within a node on the fast link, and the scheduler does not silently place them across nodes.
- The split is semantically invisible: the combined forward and backward pass computes the same function as the unsplit model, to numerical tolerance, on every step.
- Offline: on a small configuration of the same architecture that fits on one device, run the split and the unsplit model on a fixed batch and compare loss and gradients to floating-point tolerance. This is the only test that proves the layout computes the right thing.
- Before the long run: profile one step and confirm the idle fraction and bytes moved are what the design predicted; a bubble much larger than expected means the stages are unbalanced.
- During: track peak memory and step time per device; a growing step time on one device is the straggler forming.
What can go wrong
- Tensor parallelism is stretched across nodes because the in-node cards ran out, and the per-layer exchanges over the slow link make every step communication-bound.
- Pipeline stages are cut by layer count rather than by compute, so one stage with the heavy layers is the straggler and every other stage waits on it.
- Sharded optimizer state works until one device dies; the checkpoint format has to reassemble a consistent full state from shards, and it was never tested (Checkpointing).
- Numerical behaviour changes with the split: partial sums combined in a different order, mixed-precision accumulation at the boundaries. A loss that differs from the small-model reference is blamed on the model when it is the layout.
- Every cut moves data every step; the strategies differ in what they move and how often, and all of them turn the interconnect into the critical resource.
- Pipeline parallelism wastes hardware in the bubble; tensor parallelism demands a fast link; sharded state adds gather traffic. Combining them to get the best of each is the hardest engineering in the domain.
- The layout is coupled to the hardware. A split designed for eight cards on one node does not transfer to sixteen cards on two nodes without redesign, and the checkpoint format may not either.
- "Two GPUs means double the memory for the model." Each still holds its share of weights, gradients and optimizer state, and the optimizer state dominated. Sharding the state is what actually halves the largest term.
- "Pipeline parallelism is just putting layers on different cards." The naive version idles all but one card. Micro-batching and the bubble are the whole design.
- "Use tensor parallelism across the cluster; it is the most memory-efficient." It is the most communication-intensive, with an exchange inside every layer. Across a slow link it is slower than not training at all.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- SCALE-SPECIFICThis lesson applies only to models whose training state exceeds one device — in practice large transformers; for any model that fits after mixed precision and activation checkpointing, data parallelism with sharded optimizer state is the whole answer and the layer and tensor splits are unnecessary complexity.
- SIMPLIFIEDThe three strategies are described at the level of what crosses the wire and when; real systems combine all three with expert parallelism, sequence parallelism, overlapped communication and hardware-specific collectives, and the arithmetic in the section below is for the shape of the argument, not a sizing tool.
- CONTESTEDOne serious position holds that a team that is not pretraining should never operate these splits: a parameter-efficient method or a smaller model fits on one device, the operational burden of a multi-dimensional parallel layout is a full-time job, and managed training services hide it for a price. The counter is that full fine-tuning of a large model on domain data sometimes measurably beats the cheap alternatives and there is then no way around the layout.
Where the depth lives
This domain teaches the model and hands the rest off by name.