NVMe offload is the bottom tier of the memory hierarchy pressed into service as model memory: when a model’s state is too large not just for GPU memory but for CPU RAM as well, you park it on solid-state disk and stream it up through the machine only when a layer needs it. This is the trick behind ZeRO-Infinity, which pushes parameters and optimizer state all the way down to NVMe and thereby fits trillion-parameter models on a handful of GPUs. The catch is that NVMe is the slowest link by far, and every offloaded byte must climb a chain of buses — disk to PCIe to DRAM to GPU. Whether that works reduces to one question this article answers with arithmetic: can the aggregate disk bandwidth keep the GPUs fed?

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.

Advertisement

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.

Advertisement

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.

Writes and the endurance tax

Offload traffic is not symmetric: the runtime reads parameters and optimizer state off NVMe, then writes the changed fp32 state back. NVMe writes run at perhaps half the read rate, so the write leg can dominate t_disk even though it moves fewer bytes. Writes also wear the flash — rewriting terabytes of optimizer state every step accumulates petabytes over a long run, so enterprise drives rated for high drive-writes-per-day matter, and a consumer SSD can be worn to failure. Size the drives for total bytes written, not just peak bandwidth.

When NVMe offload is viable versus a bottleneck

The inequality BW_agg ≥ B / t_compute sorts the cases cleanly. NVMe offload is a genuine win when compute per step is large relative to the bytes streamed: big models with big layers, large global batches, long sequences, and plenty of drives per node. Those are exactly the extreme-scale training regimes ZeRO-Infinity was built for, where the alternative is not a slower run but no run at all.

It is a bottleneck wherever compute per byte is thin. Small batches starve the overlap. Autoregressive inference, spending only about 2 × P FLOPs per token, offers almost nothing to hide a full weight-load behind — decoding one token at a time off NVMe is disk-bandwidth-bound. And too few drives, or drives sharing PCIe lanes with the GPUs, cap BW_agg below what the inequality demands. The rule of thumb: NVMe offload rewards high-arithmetic-intensity throughput work and punishes latency-oriented, small-step usage.

Pitfalls the bandwidth math predicts

The failure modes fall straight out of the chain and the inequality. Too few drives: one SSD cannot feed a GPU, so a single-drive setup guarantees t_disk >> t_compute; aggregate several. Shared PCIe lanes: if disk reads and host-to-GPU copies contend for the same root complex, the effective bandwidth is far below the paper sum — isolate the storage lanes. Pageable host buffers: stage transfers through pinned memory or the copy cannot overlap the disk read.

No prefetch: if the runtime loads a layer’s weights only when compute reaches it, disk time and compute time add instead of overlapping; a good scheduler prefetches the next layers while the current ones run. Before enabling NVMe offload, estimate the bytes moved per iteration, divide by your real aggregate bandwidth, and compare with compute time — the verdict is usually clear on paper before a single step runs.

NVMe offload buys capacity with the slowest bandwidth in the machine: when a model’s state overflows both GPU HBM and CPU DRAM, ZeRO-Infinity spills parameters and optimizer state to solid-state disk and streams them up the chain — NVMe to PCIe to DRAM to GPU — just in time for each layer. Because the hops are in series, the effective rate is bounded by the slowest link, which is almost always the disk, so no single drive suffices; you aggregate many drives across many nodes to raise BW_agg. Whether it works is one inequality: disk time t_disk = B / BW_agg against compute time t_compute = F / R, and overlap hides the disk only when BW_agg ≥ B / t_compute. That holds for big models, big batches, and long sequences — high arithmetic intensity per step — and fails for small batches and single-token decoding, where the GPUs stall on flash. Budget for the slower write leg and flash endurance too, and NVMe offload turns a model too large for any node’s memory into one that simply trains.