Offload Math

The memory wall that forces offloading

Training a transformer with mixed-precision Adam costs about 16 bytes per parameter of persistent state: fp16 weights (2) + fp16 gradients (2) + an fp32 master copy of the weights (4) + Adam’s first and second moments (4 + 4). A 10-billion-parameter model therefore needs 16 × 10e9 = 160 GB of model state alone — before a single activation is stored — which does not fit on an 80 GB GPU, let alone a 40 GB one.

The brute-force answer is more GPUs. The cheaper answer is to notice that most of that 160 GB is idle for most of the step: the optimizer moments and master weights are touched just once per step. Memory that sits idle need not sit on the most expensive silicon you own. Offloading moves it somewhere larger and slower, and pays for the move in bandwidth.

The memory-tier hierarchy

Every accelerator sits on top of a ladder of memory tiers, each roughly an order of magnitude larger and an order of magnitude slower than the one above:

The key gap is not capacity, it is the bandwidth cliff between HBM and everything below it. HBM feeds the compute units at terabytes per second; the PCIe link to CPU DRAM is 30–100× slower, and an NVMe drive slower still. So the hierarchy offers a clean deal: near-unlimited capacity, at the price of a link whose bandwidth you must spend carefully. Offloading is the art of spending that link so the reader never notices it was used.

The overlap-vs-stall condition

Offloading rests on one fact: the forward and backward passes are compute-bound. A step does on the order of 6 × P × tokens FLOPs of matrix math, keeping the tensor cores busy for tens or hundreds of milliseconds while the PCIe link sits mostly idle. That idle link is free capacity: move offloaded state across it while the GPU is busy with unrelated math and the transfer costs nothing in wall-clock time.

Formalize it. Let a stage of the step move X bytes over a link of bandwidth B, taking t_transfer = X / B. Let the compute that runs concurrently take t_compute. Two regimes:

if  t_transfer ≤ t_compute :  transfer fully hidden   (compute-bound, no stall)
if  t_transfer > t_compute :  GPU stalls for (t_transfer - t_compute)  (bandwidth-bound)

So the question ‘can I offload this?’ is really ‘is there enough concurrent compute to hide the transfer?’ The larger your batch and sequence length, the longer t_compute grows, and the more state you can hide behind it — which is why offloading and large micro-batches are natural partners. Shrink the batch and the compute window collapses, t_transfer pokes out the top, and the same offload config that flew at batch 32 crawls at batch 1.

The per-step bandwidth budget

Turn the inequality into a budget. Over one step you cannot move more than B × t_step bytes across the link without stalling, where t_step is the compute time you have to hide behind. On PCIe 4.0 ×16 (~32 GB/s) with a 250 ms compute window, that is a ceiling of 32 × 0.25 ≈ 8 GB of round-trip traffic per step, generously.

Now count a naive plan. Offloading the optimizer step to CPU sends gradients down and receives updated fp16 parameters up — roughly 4 bytes/param each step. For 10 billion params that is 40 GB, or 40 / 32 ≈ 1.25 s. Against a 250 ms window you are five times over budget: bandwidth-bound. The budget tells you, before you run anything, whether a tier is reachable or whether you must offload less, use a faster link, or grow the compute window.

What to offload, and in what order

Not all state is equally offloadable. Rank each candidate by two numbers: how big it is (offload payoff) and how often the GPU needs it (transfer cost). The best candidate is large and rarely touched.

Optimizer states (moments + master weights, 12 of the 16 bytes) win outright: they are the bulk of the state and touched exactly once per step, so their traffic overlaps the whole backward pass. Gradients are produced layer-by-layer during backward and can be streamed down as they are born. Parameters are trickier — every layer needs them on the forward and backward pass, so offloading them means fetching each layer’s weights twice per step, much heavier traffic. Activations are usually recomputed (gradient checkpointing) rather than offloaded, since recompute often beats the round trip. Hence the natural order: optimizer first, then gradients, then parameters, and only under real pressure.

Why the optimizer step is the sweet spot

The reason optimizer offload is the canonical first move is an arithmetic-intensity argument. Adam’s update is elementwise — each parameter touched a constant number of times, so O(P) FLOPs, trivial compute over a lot of data. The forward/backward is O(P × tokens): enormous compute over the same data.

That mismatch is the opening. A low-compute, high-data operation like the Adam update runs fine on the CPU, which has ample bandwidth to its own DRAM (50–200 GB/s) even though its link to the GPU is slow. So the clean split is: keep the matmuls on the GPU, ship only gradients across PCIe, do the cheap update on the CPU against the master weights that already live there, and ship back the refreshed fp16 weights — the design worked through in Optimizer State Offload Math.

Prefetch and double buffering

The overlap condition assumes the transfer is issued early enough to run alongside compute. That does not happen for free — it takes prefetching. While the GPU computes layer i, the runtime asynchronously fetches layer i+1’s offloaded state into a staging buffer, so it is resident by the time compute reaches it.

This needs two things. First, double buffering: at least two buffers, so the GPU consumes one while the next fills. Second, pinned (page-locked) host memory, because DMA transfers to and from pageable memory are far slower and cannot fully overlap. A copy engine separate from the compute stream then drives PCIe without blocking kernels. Get this plumbing wrong — a synchronous copy, unpinned memory, a single buffer — and t_transfer serializes in front of t_compute, doubling step time even when the bandwidth budget said you were fine.

The tier ladder: CPU, then NVMe

When CPU DRAM still is not enough, you drop to the next rung. Offloading to CPU DRAM is the common case: DRAM is roomy (hundreds of GB to terabytes) and the only cost is the PCIe crossing, analyzed in CPU Offload Math. When even DRAM overflows — think hundred-billion-parameter models — you push the coldest state onto NVMe SSD.

NVMe adds a second, slower link in series (SSD → DRAM → GPU), so the same overlap condition must hold on both hops, and the SSD’s few GB/s make it viable only for the most rarely touched state — the fp32 master weights and moments, as worked in NVMe Offload Math. Together the rungs form a capacity ladder: hot state in HBM, warm in DRAM, cold on NVMe, with prefetch papering over the gaps.

How the pieces assemble: ZeRO-Offload

The framework above is realized in production by ZeRO-Offload and its NVMe extension, ZeRO-Infinity. ZeRO first partitions the 16-bytes-per-param state across data-parallel GPUs so no single device holds a full copy, then offloads each partition’s optimizer state and gradients down the tier ladder to CPU or NVMe.

The two ideas compose: partitioning shrinks each GPU’s share of the state, and offloading moves that share off-device, so effective model capacity becomes the sum of every tier across every node — a single GPU can fine-tune a model whose full Adam state dwarfs its HBM. The end-to-end system is detailed in ZeRO-Offload — CPU + NVMe Offload. This overview is the ‘why and when’ that sits above all four companion pieces.

Pitfalls and when not to offload

Offloading is not free even when the bandwidth budget closes. The recurring traps:

Small batches. A tiny compute window cannot hide the transfer; offloading a small model often runs slower than just using a bigger GPU. PCIe contention. Gradient all-reduce for data parallelism competes with offload traffic on the same link; both fighting for 32 GB/s means neither is hidden. CPU-bound optimizer. If the CPU’s elementwise Adam is slower than the transfer it overlaps, the CPU becomes the bottleneck — a vectorized CPU-Adam kernel matters. Unpinned memory silently halves effective bandwidth. The honest rule: offload when the model genuinely will not fit and the batch is large enough to hide the traffic. If it already fits, HBM is always faster — offloading buys capacity, never speed.