NOTE-006: GPipe

Contribution

Training a model past one accelerator’s memory previously meant architecture-specific infrastructure that did not transfer between tasks. GPipe makes model parallelism generic for anything expressible as a sequence of layers: partition the sequence into cells, place each on an accelerator, and — the actual contribution — split the mini-batch into micro-batches so the partitions pipeline instead of idling. Combined with re-materialization it gives almost linear speedup, demonstrated on a 557M-parameter AmoebaNet (84.4% top-1 on ImageNet-2012) and a 6B-parameter multilingual translation model.

Key insight

Naive model parallelism is a queue, not a pipeline: partition k waits for k−1 on every example, so K accelerators do one accelerator’s work. The fix is that the unit flowing through the pipeline need not be the unit the optimizer steps on. Split the mini-batch, pipeline the pieces, accumulate, then step — synchronous mini-batch SGD is preserved exactly, so the optimization is unchanged and only the schedule differs.

The second insight is that the two costs trade against each other through the same knob. More micro-batches means less bubble and less peak memory, but smaller per-device work; the paper gives both terms in closed form so the trade is arithmetic rather than guesswork.

Assumptions

• The model is a sequence of layers. That is the scope condition, stated as such. Branching or non-sequential topologies are out.
• Synchronous mini-batch gradient descent. Gradients are accumulated across micro-batches and applied once, so this is not an asynchronous or stale-gradient method and inherits none of their convergence questions.
• Partitions are evenly balanced — Figure 2(c)‘s bubble analysis assumes it, and the paper says so explicitly before discussing what happens when they are not.
• Communication is only activation tensors at partition boundaries, so the method does not assume a high-speed interconnect.

Key results

• Bubble overhead is O((K−1)/(M + K−1)), amortized over M micro-steps, with K partitions.
• M ≥ 4×K makes it negligible — the paper’s own empirical threshold, and the rule the record was missing. Helped by re-computation in the backward pass being schedulable early, without waiting for gradients from earlier layers.
• Peak activation memory with re-materialization and partitioning: O(N + (L/K) × (N/M)), where N/M is the micro-batch size and L/K the layers per partition. Without either it is O(N × L).
• Low communication: only activation tensors cross partition boundaries, so efficient scaling holds even without high-speed interconnects.
• Almost linear speedup when a model is partitioned across accelerators.
• Demonstrated across two unrelated architectures and tasks — 557M AmoebaNet on ImageNet, 6B-parameter multilingual NMT — which is the transferability claim the paper opens with.

Claims

id	claim	strength	support
C1	Micro-batch splitting converts naive model parallelism into a pipeline with almost linear speedup	strong	the bubble formula plus measured scaling
C2	Bubble overhead is negligible at M ≥ 4K	moderate	an empirical threshold from their experiments, not derived
C3	Re-materialization plus partitioning reduces peak activation memory to O(N + (L/K)(N/M))	strong	derived, and the basis of the memory claims
C4	The method is architecture-independent for sequential models	strong	two very different architectures and tasks
C5	It works without high-speed interconnects	moderate	follows from only activations crossing boundaries; argued rather than measured across fabrics
C6	Optimization semantics are unchanged	strong	synchronous mini-batch SGD by construction

Method

Algorithm: GPipe — batch-splitting pipeline parallelism.

Express the model as a sequence of L layers. Partition into K cells, one per accelerator. Split each mini-batch into M micro-batches. Pipeline the micro-batches forward through the cells, then backward, re-materializing activations rather than caching them. Accumulate gradients across all M and apply one synchronous update.

Key components

• Micro-batch splitting — the pipelining mechanism
• Re-materialization on every accelerator — the memory mechanism
• Partition balancing — the assumption the bubble analysis rests on
• Early scheduling of backward re-computation, which is part of why M ≥ 4K suffices

Concepts

• Bubble overhead — accelerator idle time introduced by partitioning, before the pipeline fills and after it drains. O((K−1)/(M + K−1)).
• Micro-batch — a slice of the mini-batch that flows through the pipeline independently. Distinct from the optimizer’s batch, which is still the whole mini-batch.
• Cell — a consecutive group of layers placed on one accelerator.
• Re-materialization — recomputing forward activations during the backward pass instead of storing them. Also called gradient checkpointing.

Connections

Builds on re-materialization / gradient checkpointing as prior work rather than inventing it; the contribution is combining it with batch-split pipelining. It stands beside data parallelism and tensor parallelism as the third axis, and the record’s later distributed practices assume all three.

Recommendations

• R1 — Split the mini-batch into micro-batches when using pipeline parallelism. Topic: distributed training. Status: standard. Strength: strong. Applies when: always, under pipeline parallelism — without it the pipeline is a queue.
• R2 — Choose M ≥ 4×K. Topic: distributed training. Status: standard. Strength: moderate. Applies when: the per-device micro-batch stays large enough to keep the accelerator efficient, which is the competing constraint.
• R3 — Balance partitions. Topic: distributed training. Status: standard. Strength: strong. Applies when: always — the whole bubble analysis assumes it, and the slowest cell sets the pipeline’s rate.
• R4 — Size cells from O(N + (L/K)(N/M)). Topic: distributed training. Status: standard. Strength: strong. Applies when: the memory/compute trade needs deciding; the formula makes it arithmetic.

Bearing on the record

All three practices sourced to this note are confirmed — the fourth confirming cluster in #114.

practice	disposition
SOTA-017 use micro-batch splitting	confirmed — R1, the paper’s central mechanism
SOTA-018 balance pipeline stages to minimize bubble overhead	confirmed — R3, and the bubble formula makes “minimize” quantitative
SOTA-019 choose pipeline chunks by memory vs. compute	confirmed — R4, and the memory expression is the trade in closed form

All three gain the same thing: numbers where they had adjectives. SOTA-017 gets M ≥ 4×K; SOTA-018 gets O((K−1)/(M + K−1)) as what “bubble overhead” actually is; SOTA-019 gets O(N + (L/K)(N/M)) against O(N × L) as the trade it names.

That is the pattern across every confirming cluster so far — the practices were right and stated at a level that could not be acted on, which is what a note nobody read against its paper produces even when nothing is wrong.

Limitations

• M ≥ 4K is empirical and from their hardware and models. It is a good default and not a derived bound.
• Even partitioning is assumed for the analysis; the paper discusses imbalance but the clean formula does not survive it.
• Sequential models only.
• 2018 accelerators. The bubble/memory trade is unchanged in form, but where it lands has moved with memory capacity and interconnect speed.

Open questions

• Where does M ≥ 4K land on current hardware, and does the early-scheduling argument still carry it?
• Pipeline parallelism now composes with tensor and data parallelism and with sharding (SOTA-116’s line). Nothing in this record states how the bubble formula interacts with those.
• Imbalanced partitions are the common case in practice. What is the bubble overhead as a function of the imbalance?