When a giant language model is trained across many GPUs in a pipeline, the standard schedule forces the first stage to remember far more intermediate data than the last — so to fit in memory, every stage recomputes everything it dropped, a one-size-fits-all tax that overworks the stages with memory to spare; AdaPipe (Tsinghua) instead lets each stage recompute only as much as its own memory pressure demands, and then, because recomputing turns a memory imbalance into a compute imbalance, moves whole layers off the busy early stages onto the idle late ones — two dynamic-programming searches that together speed up GPT-3 and Llama 2 training up to 1.32×.
The argument in one glance
-
Problem
A pipeline that remembers unevenly
A model too big for one GPU is split into a PIPELINE of stages across GPUs — stage 0 runs the first layers, stage 1 the next, and so on, with data flowing forward and gradients flowing back. To keep every GPU busy, many small "micro-batches" are kept in flight at once. But computing a gradient needs the ACTIVATIONS (intermediate results) from that micro-batch's forward pass, held in memory until its backward pass finishes. The standard memory-efficient schedule (1F1B) means an early stage must hold the activations of MANY in-flight micro-batches while a late stage holds almost none — so memory use is savagely imbalanced across the pipeline, and with longer sequences it only gets worse.
-
What people did
Recompute to save memory
The universal fix is recomputation (activation checkpointing): in the forward pass, throw most activations away; in the backward pass, redo the forward computation to regenerate them just in time. This slashes memory at the cost of extra compute. The typical "full recomputation" keeps only each layer's input and regenerates everything else — enough to make even the memory-heavy first stage fit.
-
Why it fell short
One recomputation policy for stages with wildly different memory
The trouble is that recomputation is applied UNIFORMLY — every stage recomputes everything, even the last stage that had tens of gigabytes of memory sitting idle. That is pure wasted compute: the memory was there to keep those activations, but the one-size-fits-all policy threw them away and paid to regenerate them anyway. Prior attempts to be smarter recompute at too coarse a granularity (treating a whole attention or feed-forward layer as one indivisible block, though it mixes memory-hungry and compute-hungry operators) or hand-pick operators uniformly — never adapting to the fact that stage 0 and stage 7 face completely different memory pressure. Profiling GPT-3 shows the first stage using 2.33× the memory of the last, with 30+ GB wasted per stage under full recomputation.
-
This paper's idea
Recompute per stage, then rebalance the layers
AdaPipe (Tsinghua University) makes two coupled moves. ADAPTIVE RECOMPUTATION: split each layer into fine-grained "computation units," and let every stage independently choose which units to save and which to recompute — a stage with spare memory saves more and recomputes less, a memory-starved stage recomputes more. A dynamic-programming search (a knapsack: maximize saved recompute-cost under the memory budget) finds each stage's optimal choice. But this converts a MEMORY imbalance into a COMPUTE imbalance — early stages now do more recompute work, becoming the pipeline's bottleneck. So ADAPTIVE PARTITIONING: a second DP search moves whole layers off the busy early stages onto the idle late ones, rebalancing compute time. Together they cut wasted work: up to 1.32× faster GPT-3 and 1.22× faster Llama 2 training, the search finishing in seconds.
Why this problem is worth solving
Training a frontier language model is a memory problem before
it is a compute problem: the model, its gradients, its
optimizer state, and above all the ACTIVATIONS of every
in-flight sample must fit in GPU memory, and a single GPU
holds a small fraction of what a 175-billion-parameter model
needs. So the model is spread across GPUs with
splitting a model’s layers
into sequential STAGES on different GPUs — stage 0 holds the
first layers, stage 1 the next, and so on. A training sample
flows forward through all stages, then its gradients flow
backward. Because one sample cannot be in two stages at once,
the pipeline is kept busy by streaming many small
micro-batches through it simultaneously, like cars on an
assembly line. It is the parallelism of choice BETWEEN nodes,
because it communicates little — only the activations at
stage boundaries — and tolerates the slower links between
machines
That discipline,
one-forward-one-backward: instead
of running all forward passes then all backward passes
(which would force every stage to remember every
micro-batch’s activations at once), 1F1B interleaves them so
each stage does a forward pass, then a backward pass, in
alternation once the pipeline is full. This caps how many
micro-batches’ activations any stage must hold — but the cap
is UNEVEN: stage s (out of p stages) must retain the
activations of p−s micro-batches during the steady phase. So
the first stage holds many, the last holds one
What people did before, and where it broke
- Uniform full recomputation. Keep only each layer’s input, regenerate the rest, in every stage. It makes the worst stage fit — and taxes every other stage with recompute it did not need, leaving tens of gigabytes unused. The default, and the thing to beat.
- Coarse-grained recomputation (chain-structure optimizers). Find a good save/recompute split, but treat a whole attention or feed-forward layer as one atomic unit — even though each layer mixes cheap-to-recompute, memory-heavy operators with expensive-to-recompute, small ones. Too blunt to make the right trade.
- Manual selective recomputation. Hand-pick specific operators (softmax, dropout, the attention matmuls) to recompute uniformly. Better, but applied identically to every stage, and largely superseded by Flash Attention, which fuses those very operators and removes their activations outright.
- Finer layer-splitting (vPipe). Divides attention and feed-forward layers along residual connections — but still cannot handle operators with multiple activation inputs (like the attention-score computation), so its granularity stops short of where the real trade-offs live.
- The gap. Nobody let each stage choose its OWN recomputation at operator granularity, and nobody then rebalanced the layer assignment to absorb the compute imbalance that per-stage recomputation creates. AdaPipe couples exactly those two.
How AdaPipe works
Profile, then two nested searches, then rebalance
-
Split layers into computation units new
Rather than deciding save-or-recompute per operator, AdaPipe groups operators into computation UNITS: a unit is a minimal cluster whose internal tensors are not saved even in a normal backward pass (transposes, adds, divisions ride along with a GEMM), bounded by the operators whose outputs DO get saved. This is the right granularity — each unit is a single, meaningful save-or-recompute decision, and the division fits any transformer (GPT-3, Llama 2, BERT).
-
Profile time and memory per unit new
A short preliminary run (5–10 iterations) records each unit's forward time, backward time, and the memory its saved tensors occupy — including tensors that fused ops like Flash Attention save internally. Because transformer layers are homogeneous, these per-unit costs generalize across the whole model.
-
Adaptive recomputation: a knapsack per stage new
For each stage, choosing which units to SAVE (versus recompute) to minimize backward-pass time under a memory budget is exactly a knapsack problem — each saved unit "costs" (p−s)× its memory (it is held for p−s micro-batches) and "pays back" the forward time no longer recomputed. A dynamic-programming search solves it optimally per stage: memory-rich late stages save many units, starved early stages save few.
-
The imbalance flips from memory to compute new
Now every stage fits comfortably (~70 GB, balanced) — but early stages, forced to recompute more, take longer per micro-batch. Profiling shows the slowest stage running 1.17× the fastest. In the pipeline's long STEADY phase, whose pace is set by the slowest stage, that imbalance is the new bottleneck. Fixing memory created a compute problem.
-
Adaptive partitioning: a second DP over layer boundaries new
Because transformer layers are identical in shape, a stage boundary can be moved without changing communication. A second dynamic-programming search — built ON TOP of the recomputation DP so the two co-optimize rather than settling into a local optimum — reassigns whole layers: take layers off the recompute-heavy early stages, give them to the idle late stages, until per-stage compute time is even. Early stages end up with fewer layers, later stages with more.
-
Fast search via structure
Naively the two-level search is costly, but transformer homogeneity makes many layer sub-sequences isomorphic (any k-layer run starting with attention is interchangeable), collapsing redundant sub-searches, and activation sizes being powers of two lets the knapsack's memory axis be divided by a common factor. Net: the whole plan for GPT-3 or Llama 2 is found in seconds.
The idea
Match the policy to the profile instead of averaging over it. The whole insight is that 1F1B’s memory imbalance is not noise to be smoothed but STRUCTURE to be exploited: stage s deterministically holds p−s micro-batches’ activations, so the memory pressure at every stage is known in advance and strictly decreasing down the pipeline. A uniform recomputation policy discards that knowledge, treating the memory-rich last stage exactly like the starved first stage. AdaPipe instead reads the profile and shapes the policy to it: recompute a lot where memory is tight, recompute little where it is plentiful. It is the same move this site keeps seeing — when a resource is unevenly distributed and the distribution is predictable, adapt per-member rather than provision for the worst case — here spending exactly enough recompute at each stage rather than the worst stage’s dose everywhere.
Fine granularity is what makes the trade-off tradeable. The reason coarse methods fail is that a transformer layer is not one thing: inside it sit operators that are cheap to recompute but hold huge activations (ideal to drop) alongside operators that are expensive to recompute but hold little (ideal to keep). Bundling them into one save-or-recompute decision forces a bad average. AdaPipe’s computation units cut at exactly the seams where activations are saved, so each decision is a clean, independent knapsack item with a real memory-cost and time-payoff — and only at that granularity does the per-stage optimization have meaningful choices to make. The unit abstraction is the quiet enabling contribution.
Solving one imbalance reveals another — so co-optimize both. The paper’s cleanest lesson is that per-stage recomputation does not eliminate imbalance; it TRANSFORMS memory imbalance into compute imbalance. Early stages now grind through more recompute work, and since the pipeline’s steady phase runs at the speed of its slowest stage, the win would be squandered without a second fix. Adaptive partitioning is that fix — shift layers from the busy early stages to the idle late ones — but crucially it is not run after recomputation; the two dynamic programs are NESTED, the partitioning DP calling the recomputation DP for each candidate layer range, so they reach a joint optimum instead of each locally optimizing and fighting. The two moves also help in complementary regimes: adaptive recomputation shines when micro-batches are few (the warmup and ending phases dominate), adaptive partitioning when they are many (the steady phase dominates).
GPT-3 across eight pipeline stages
Following the memory-to-compute-to-balance chain
GPT-3 (175B), sequence length 16384, an 8-stage pipeline (tensor-parallel 8). Numbers from the paper's per-stage profiling on A100 GPUs.
Step 1 of 5
The 1F1B baseline: lopsided memory
With no recomputation, stage 0 must hold the activations of 8 in-flight micro-batches while stage 7 holds just 1 — so stage 0 needs over 80 GB and simply does not fit, while stage 7 has tens of gigabytes free. The imbalance is 2.33× first-to-last. This is why uniform full recomputation gets switched on: it is the only way to make stage 0 fit.
stage s holds p−s micro-batches' activations — the imbalance is structural, not incidental full recomputation makes stage 0 fit but leaves 30+ GB idle on every other stage
Why it works
- 1F1B’s memory pressure is deterministic and strictly decreasing down the pipeline → each stage’s optimal recomputation can be computed in advance from its known budget → memory-rich stages stop paying recompute they never needed.
- Computation units cut at the seams where activations are actually saved → each save-or-recompute choice is a clean knapsack item with real memory-cost and time-payoff → the per-stage optimization has meaningful, independent decisions to make.
- Per-stage recomputation converts memory imbalance into compute imbalance → a second search reassigns whole layers from busy early stages to idle late ones → the steady phase, gated by the slowest stage, speeds up.
- The partitioning DP calls the recomputation DP for each candidate layer range → the two co-optimize instead of locally optimizing and conflicting → the plan reaches a joint near-optimum, not a compromise.
- Transformer layers are homogeneous (isomorphic sub-sequences, power-of-two activation sizes) → redundant sub-searches collapse and the knapsack’s memory axis shrinks → the whole plan is found in seconds, cheap enough to run per configuration.
What it costs
- A modest, bounded speedup. The gains are 1.18–1.32×, not order-of-magnitude — this is a scheduling optimization that recovers wasted recompute, capped by how much waste existed. Biggest when sequences are long and memory tight; smaller when full recomputation was barely needed.
- A simplifying memory restriction. AdaPipe forces each attention and feed-forward layer’s OUTPUT tensor to be saved (never recomputed), which keeps the memory model tractable but means it can use MORE memory than full recomputation in the extreme — so at tensor-parallel size 1 (huge output tensors) it can run out of memory where the baseline fits.
- Depends on the profile holding. The per-unit costs come from a short profiling run and the homogeneity assumption; heterogeneous or non-transformer architectures, or shifting runtime conditions, would weaken the model.
- Scope is recomputation + partitioning. It does not search tensor/data-parallel sizes per stage, does not do offloading, and layers the two DPs on 1F1B specifically — it optimizes within a fixed 3D-parallel strategy rather than replacing the strategy search.
- Restricted to save-the-boundary layouts. The unit and buffer model assumes buffer size never exceeds one decoder layer’s intermediates; unusual layer structures outside the attention/feed-forward pattern are out of the modeled space.
The evaluation
GPT-3 (175B) and Llama 2 (70B) trained on NVIDIA A100 GPUs (up to 64) and Huawei Ascend 910 NPUs (up to 2,048), against DAPPLE and Chimera pipelines with full and no recomputation, with an “Even Partitioning” variant isolating adaptive recomputation alone.
Recomputation and partitioning policies for a pipeline
| Property | Uniform recomputation (DAPPLE / Chimera, full or none) | Adaptive recomputation only (Even Partitioning) | AdaPipe (recomputation + partitioning) |
|---|---|---|---|
| Recomputation policy | |||
| Memory across stages | |||
| Compute balance across stages | Full: even (all recompute equally); Non: even but infeasible | ||
| Handling long sequences (16384) | No-recompute OOMs everywhere; full-recompute leaves memory idle | Fits and uses memory well; 1.21× over baseline | |
| End-to-end speedup | Reference (DAPPLE) |
Check yourself
Check your understanding
1. Under 1F1B pipeline scheduling, the first stage runs out of memory while the last stage has tens of gigabytes free. Predict why this imbalance is structural — and why uniform full recomputation is a poor response to it.
2. AdaPipe splits transformer layers into fine-grained "computation units" rather than deciding recomputation per whole layer or per individual operator. Why is this granularity the enabling choice?
3. After adaptive recomputation balances memory across stages, AdaPipe adds a SECOND optimization (adaptive partitioning). Predict what problem the first optimization created that the second must solve, and why the two are nested rather than run in sequence.
4. AdaPipe delivers 1.18–1.32× speedups — real but bounded, and largest when sequences are long and memory is tight. Predict what sets this ceiling, and why the regime dependence follows from it.