The tier below RAM: why disk becomes memory
CPU offload leans on system RAM, but RAM also runs out. A model with P parameters trained in mixed precision with Adam needs about 16 bytes per parameter of persistent state — an fp16 weight and gradient (2 + 2) plus the fp32 master weight and two Adam moments (4 + 4 + 4). For a 1 trillion-parameter model that is 1e12 × 16 = 16 TB of state.
No single node has 16 TB of DRAM; a fat training box tops out around 1 to 2 TB. But that same node can hold a dozen NVMe SSDs totaling tens of terabytes for a fraction of the cost of the equivalent RAM. The observation that justifies CPU offload, taken one step further: the optimizer state and even the parameters are idle almost all the time, touched only when their layer is active. Idle terabytes can wait on cheap, abundant flash and be paged in just before they are needed.
The bandwidths that define the problem
Every tier of the hierarchy trades capacity for speed, and the numbers span orders of magnitude. Round figures for the path an offloaded tensor travels:
GPU HBM ~1500-3000 GB/s (on-device)
PCIe 4.0 x16 ~32 GB/s (GPU <-> host)
CPU DRAM ~100-400 GB/s (host memory)
One NVMe SSD ~3.5-7 GB/s (PCIe 4.0 x4 read)
NVMe write ~2-4 GB/s (slower than read)A single NVMe drive is roughly 300-1000× slower than the GPU’s own memory and several times slower than PCIe. That is the central difficulty: the storage tier that gives you the capacity is also, by a wide margin, the narrowest pipe in the machine. Making NVMe offload work is entirely about not letting that narrow pipe stall the GPU.
ZeRO-Infinity: partition, then spill to disk
ZeRO-Infinity extends the ZeRO family of optimizations. Plain ZeRO partitions the 16 bytes of state across data-parallel GPUs so no GPU holds the whole model. ZeRO-Offload pushes the fp32 optimizer state down to CPU RAM. ZeRO-Infinity adds a third destination — NVMe — and lets each tier overflow into the next.
The result is a memory hierarchy the runtime manages explicitly: hot fp16 weights for the current layer live in GPU HBM; the fp32 optimizer state and spare parameters live in CPU DRAM if they fit; and when DRAM is exhausted, everything spills to NVMe. The optimizer step runs on the CPU, so the fp32 moments never reach the GPU — they are read from disk into host memory, updated, and written back. What the GPU sees is a just-in-time stream of the fp16 weights it needs, prefetched from disk while it works on the previous layer.
The bandwidth chain: disk to PCIe to DRAM to GPU
A byte on NVMe does not reach the GPU in one hop. It is read off the SSD, crosses PCIe into pinned host DRAM, and then crosses PCIe again into GPU HBM. Each link has its own rate, and because the hops are in series, the effective end-to-end bandwidth is bounded by the slowest one:
BW_effective ≤ min( BW_nvme, BW_pcie, BW_dram )Here the NVMe read rate is almost always the floor, so a single drive caps the path at a few GB/s however fast PCIe and DRAM are. Worse, if the disk-to-host and host-to-GPU hops share the same PCIe root complex, they contend for the same lanes and the rate drops further. The answer is to widen the floor: many drives in parallel, on separate PCIe lanes from the GPUs, with the two hops pipelined so a tensor’s disk read overlaps the previous copy.
The overlap requirement, as an inequality
Reduce one training iteration to two competing quantities. Let B be the total bytes that must move off (and onto) NVMe per iteration, and BW_agg the aggregate disk bandwidth the machine can sustain. Disk time is t_disk = B / BW_agg. Let the GPUs do F FLOPs of useful work at combined throughput R; compute time is t_compute = F / R.
If disk traffic and compute overlap, the iteration takes max(t_compute, t_disk). NVMe offload is essentially free when t_disk ≤ t_compute — the next layer’s weights finish loading before the GPU finishes the current layer, so it never stalls. Rearranged, that is a bandwidth floor: you need BW_agg ≥ B / t_compute. If the disks cannot hit that rate, the GPUs idle waiting on flash and throughput collapses. Every design choice in ZeRO-Infinity — aggregating drives, prefetching, big layers — exists to satisfy this one inequality.
Aggregate bandwidth: many drives, many nodes
The word aggregate is doing real work. A single 3.5 GB/s SSD cannot feed a GPU that wants tens of GB/s of weights, so ZeRO-Infinity never relies on one drive. It strings together the bandwidth of many: eight NVMe drives in a node, striped in parallel, deliver roughly 8 × 3.5 ≈ 28 GB/s, and across a whole cluster the aggregate reaches into the hundreds of GB/s.
This is why NVMe offload scales with node count, not just drive count. Every added node contributes both GPUs (more compute) and its own local NVMe bandwidth (more BW_agg). Because each GPU reads mostly the shard of state it owns, disk traffic is naturally distributed and aggregate bandwidth grows in step with the compute it must keep fed — the balance that lets a cluster train a model far too large for the summed HBM of all its GPUs.
A worked example at trillion scale
Take the 1 trillion-parameter model, one node of eight GPUs, and eight NVMe drives giving BW_agg ≈ 25 GB/s read. Suppose that in one iteration this node’s shard streams roughly 200 GB of weights and state off disk (its partition of the 16 TB, read once for forward and once for backward). Disk time is t_disk = 200 / 25 = 8 s.
Now the compute. A training step costs about 6 × P FLOPs per token. With a large global batch — say 2 million tokens to keep the pipeline busy — that is F = 6 × 1e12 × 2e6 = 1.2e19 FLOPs. At a sustained 2 PFLOP/s across the eight GPUs, t_compute dwarfs the 8 s of disk time, so t_disk hides completely inside it. Shrink the batch, though, and the fixed 8 s stops overlapping — the GPUs stall and the trillion-parameter run grinds. Same disks, same model; the arithmetic intensity of each step decides whether NVMe is invisible or fatal.