ZeRO-Infinity is what you get when you stop treating offload as a destination and start treating it as a hierarchy. Stage 3 already divides every byte of training state by the data-parallel degree; CPU and NVMe offload each already move one bucket of state down one rung. Combine all three and something new appears — a runtime whose effective capacity is the sum of HBM plus DRAM plus flash across every node, and whose viability is decided not by any single bandwidth number but by a ratio: how much arithmetic you do per byte you drag up the ladder. This article works that ratio out, then follows it into the three design choices it forces: partitioning tensors for bandwidth rather than ownership, running the fetch pipeline a layer ahead, and tiling operators that no longer fit on one device.
What ZeRO-Infinity adds on top of stage 3
The companion pieces each own one rung. Stages 1–3 partition optimizer state, gradients and parameters across dp ranks, cutting the 16 bytes-per-parameter budget to 16Ψ/dp. CPU offload moves the fp32 optimizer bucket to host DRAM; NVMe offload pushes it onto flash. Each is a single move.
ZeRO-Infinity is the composition, and composing changes the objective. Once every tier is available simultaneously, the question is no longer ‘does this bucket fit?’ but ‘what is the cheapest placement of every bucket such that no tier stalls the GPU?’ Capacity becomes an aggregate — C_total = Σ_nodes (HBM + DRAM + NVMe) — and the binding constraint moves from bytes to bandwidth. Three mechanisms exist purely to make that constraint satisfiable, and none appear in a single-tier design: bandwidth-centric partitioning, a deep prefetch pipeline, and memory-centric tiling.
The FLOPs-per-byte test every tier must pass
Overlap is usually written t_transfer ≤ t_compute. Divide through and it becomes scale-free and far more useful. Let a GPU sustain R FLOP/s and a tier deliver BW bytes/s to it; that tier is invisible exactly when the work’s arithmetic intensity clears the machine’s balance point:
ait = FLOPs done / bytes moved
tier is hidden <=> ait >= R / BW (the machine balance)With activation checkpointing a training step costs forward (2P) + recomputed forward (2P) + backward (4P) = 8P FLOPs per token, so T tokens cost 8·P·T. The fp16 parameters must climb the ladder twice per step (once for forward, once for the recompute-plus-backward), moving 4P bytes. So ait_params = 8PT / 4P = 2T. Parameter streaming is hidden when 2T ≥ R/BW — a condition on batch size, not on the model. This single inequality decides every placement below.
Applying the test: which state can live where
Take R = 125 TFLOP/s sustained per GPU and read off what each tier demands:
| Tier | BW to one GPU | Required ait | Tokens needed per link |
|---|---|---|---|
| HBM | 2 TB/s | 63 FLOPs/byte | ~32 |
| DRAM over PCIe 4.0 x16 | 32 GB/s | 3,900 | ~2,000 |
| NVMe (one drive’s share) | 3 GB/s | 41,700 | ~20,900 |
Read that last row carefully: it is a per-link bar, the cost when one rank must pull everything it needs off its own drive. On those terms batch 16 × seq 2048 = 32,768 clears it and a lone 2048-token sequence misses by 10×. That is the problem statement, not the verdict — the next section divides the bar by dp. The Adam update, by contrast, fails the test at every tier and every batch size: it is elementwise, roughly 12 FLOPs against 24 bytes read and written, so ait ≈ 0.5. It can never be hidden behind GPU compute, which is why it does not run on the GPU at all — it runs on the CPU, beside the DRAM and flash it touches, where its absolute time is small next to a long step.
Bandwidth-centric partitioning: allgather, not broadcast
Classic stage 3 assigns each parameter tensor an owner. To materialize a layer, the owner reads the whole tensor from its own slow tier and broadcasts it. The aggregate bytes read per rank are the same either way — but the concurrency is not. Under owner-broadcast, the layer you need right now is gated by one rank’s link:
owner-broadcast: t_fetch = layer_bytes / bw_one
bandwidth-centric: t_fetch = layer_bytes / (dp * bw_one)ZeRO-Infinity instead stripes every tensor across all dp ranks and materializes it with an allgather, so all dp NVMe links and PCIe lanes read their slice at the same instant. For a 1T-parameter model with 128 layers, a layer is 2e12/128 = 15.7 GB of fp16 weights; at 3 GB/s one owner needs 5.2 s, while 64 ranks striped need 0.08 s. The flash tier’s required arithmetic intensity falls by a factor of dp — roughly 21,000 tokens per link becomes a few hundred — which is what makes NVMe-resident parameters practical at batch sizes a single-link design could never hide.
Prefetch depth: how far ahead the pipeline must run
Striping buys bandwidth, not timeliness: a fetch issued when compute reaches the layer is a stall however fast it runs. Let t_layer be one layer’s compute time and t_fetch the time to materialize the next. The lookahead depth d must satisfy d ≥ t_fetch / t_layer, and each unit of depth costs a resident buffer:
d >= t_fetch / t_layer buffer_bytes = d * layer_bytesWork it for the 1T model at a modest 4,096 tokens per GPU: step compute is 8 × 1e12 × 4096 / 1.25e14 ≈ 262 s, so t_layer ≈ 2.05 s. Under owner-broadcast, d ≥ 5.2/2.05 = 2.5 → d = 3, i.e. 47 GB of prefetch buffers — more than an 80 GB card can spare once activations are resident. Bandwidth-centric partitioning drops t_fetch to 0.08 s, so d = 1 suffices and the buffer is one layer. The two mechanisms are the same mechanism: striping is what makes shallow, affordable prefetch possible.
Memory-centric tiling: operators bigger than the device
Everything so far assumes a materialized layer fits. At extreme width it does not: for the 1T model, d_model = 25600 and the FFN’s first matrix alone is 25600 × 102400 = 2.6e9 parameters, 5.2 GB in fp16 — and there are two such matrices plus attention projections per layer. The historical fix was tensor parallelism, which drags in extra collectives and a fixed device group.
Memory-centric tiling avoids that. Split the operator along its output dimension into k tiles and execute them sequentially, gathering and releasing each tile in turn. Working memory becomes layer_bytes / k (doubled if you double-buffer), so with k = 8 the 5.2 GB matrix needs 655 MB resident, or 1.3 GB with a tile in flight. Because a column-split matmul produces a slice of the output, no cross-tile communication is needed. The cost is matmul shape: tile too finely and the kernels get small and utilization drops.
The ceiling that is advertised, and the one that binds
The headline number is a division. If persistent state is 16 bytes per parameter, the aggregate flash capacity sets a hard limit:
P_max = C_nvme_total / 16 bytesOne node with 8 × 8 TB drives holds 64 TB, so P_max = 6.4e13/16 = 4e12 — four trillion parameters on a single 8-GPU box. Sixty-four such nodes give 4 PB and a nominal 256 T parameters. Two other ceilings bind first. Width: a tile of the widest operator plus its activations must fit in HBM, capping d_model independently of capacity. And time — a 256T model wants roughly 20 tokens per parameter, so 5.1e15 tokens, so 8 × 2.56e14 × 5.1e15 ≈ 1e31 FLOPs — about 33,000 years at a sustained 10 EFLOP/s. The capacity ceiling exceeds the compute ceiling by four orders of magnitude, so it was never the real limit.
The honest cost: efficiency versus simply buying GPUs
ZeRO-Infinity does not make offload free; it makes the tax predictable. Per-step efficiency against an all-HBM baseline is roughly min(1, ait / (R/BW)) for the binding tier, and a fully NVMe-resident run lands well under what the same GPUs achieve with state in memory.
Compare honestly. A 1T model needs 16 TB of state; at perhaps 50 GB of an 80 GB card usable for state after activations and buffers, plain stage 3 needs on the order of 320–500 GPUs before it can run at all. ZeRO-Infinity runs it on 64, at lower per-GPU efficiency. Total throughput still favours the big cluster by a wide margin — scale is not free. What you get is access: a hard out-of-memory wall converted into a soft, quantifiable bandwidth tax. The rule follows directly. If aggregate HBM comfortably exceeds 16Ψ plus activations, use stage 3 and keep the efficiency. Reach for Infinity when it does not, or when GPU count is set by budget rather than by the model.
Pitfalls the synthesis creates
Composing tiers creates failure modes none of the single moves have. Margin collapse: your headroom over the flash tier is dp · 2T against R/BW, so it erodes from both ends — halving the micro-batch and halving the data-parallel degree cost exactly the same. Small-dp, small-batch fine-tuning is where the design quietly stops working, and against the DRAM tier, whose bar is ~2,000 tokens per link, a single short sequence on a handful of ranks is already borderline. Link contention: the allgather that makes striping work and the NVMe reads that feed it may share one PCIe root complex, in which case the dp × bw_one figure is fiction; place storage and interconnect on separate lanes.
Prefetch mistuned both ways: too shallow and fetches serialize in front of compute; too deep and the buffers eat the HBM the activations needed. Page-cache shadowing: buffered file I/O re-copies flash data into the same DRAM you were trying to relieve, so direct I/O into pinned buffers is mandatory, not an optimization. And the whole edifice assumes throughput training: autoregressive decoding does about 2P FLOPs against 2P bytes of weights, ait ≈ 1, four orders short of the flash tier’s requirement.
R/BW; with checkpointing, parameter streaming has ait = 2T, which puts the flash tier’s bar at ~21,000 tokens per pass per link — a bar striping then cuts by dp to a few hundred. Three mechanisms satisfy that inequality: bandwidth-centric partitioning, which stripes every tensor so dp links read at once instead of one owner broadcasting; a prefetch pipeline deep enough that d ≥ t_fetch/t_layer, which striping makes shallow and therefore affordable; and memory-centric tiling, which shrinks a too-wide operator’s working set to layer_bytes/k without tensor parallelism. Capacity then becomes P_max = C_nvme/16 — trillions of parameters, far past what any compute budget could actually train. The honest reading: Infinity does not beat more GPUs on throughput, it converts an out-of-memory wall into a bandwidth tax you can compute in advance.