The memory the GPU cannot afford
Start with why offloading is worth the trouble. In mixed-precision Adam training, the memory that scales with the parameter count Ψ is not dominated by the model itself — it is dominated by the optimizer. For every parameter you pay for an fp16 weight, an fp16 gradient, and then a full fp32 triple: a master copy of the weight, a first-moment (momentum) estimate, and a second-moment (variance) estimate. That fp32 triple is the expensive part, and it lives on the GPU purely so the optimizer can touch it once per step.
The tragedy is that this memory is almost idle. The forward and backward passes — the compute-hungry work — never read momentum or variance. They only need the fp16 weights and produce fp16 gradients. So the largest block of GPU memory is reserved for a step that runs once per iteration and is arithmetically light. ZeRO-Offload’s whole thesis is that this idle-but-huge state is exactly what you should relocate to cheaper, more plentiful CPU RAM.
The 16-bytes-per-parameter budget
Make it concrete. For each parameter, standard mixed-precision Adam holds:
| State | Precision | Bytes / param | Where it must live |
|---|---|---|---|
| Weight | fp16 | 2 | GPU (used in fwd/bwd) |
| Gradient | fp16 | 2 | GPU (produced in bwd) |
| Master weight | fp32 | 4 | optimizer only |
| Momentum (m) | fp32 | 4 | optimizer only |
| Variance (v) | fp32 | 4 | optimizer only |
That sums to 16Ψ bytes, and the split is stark: 4Ψ for the fp16 weight-and-gradient pair that the GPU genuinely needs during compute, and 12Ψ — three quarters of the total — for the fp32 optimizer states that only the update step reads. A 10-billion-parameter model therefore demands 16 × 10×10^9 = 160 GB of state before a single activation is stored. No single commodity GPU has that. But 120 GB of it is optimizer state you can send elsewhere.
The split: what stays, what moves
ZeRO-Offload draws the line exactly along that 4Ψ / 12Ψ seam. The GPU keeps the fp16 weights (it needs them to run the forward pass) and computes gradients in the backward pass. Everything the optimizer owns — the fp32 master weights, momentum m, and variance v — is pinned in CPU memory, and the Adam update runs on the CPU.
The per-step data flow follows naturally. As the backward pass finishes each gradient, that fp16 gradient is streamed off the GPU to the CPU. The CPU then runs Adam entirely in fp32, mutating m, v, and the master weights in place. It casts the updated master weights back to fp16 and copies them onto the GPU, overwriting the stale weights just in time for the next forward pass. The GPU never holds the fp32 triple at all: its model-state footprint drops from 16Ψ to roughly 2Ψ (the fp16 weights), an 8× reduction — leaving VRAM free for activations and a bigger batch.
The CPU-side Adam step
Why is it acceptable to run the update on a CPU that is far slower than the GPU at dense linear algebra? Because the Adam update is not linear algebra — it is a cheap, fully element-wise recipe applied once per parameter:
m ← β1·m + (1-β1)·g
v ← β2·v + (1-β2)·g^2
m̂ = m / (1-β1^t) v̂ = v / (1-β2^t)
w ← w - lr · m̂ / (√v̂ + ε)That is a handful of multiplies, adds, a square, and a square root per parameter — on the order of a dozen FLOPs. Crucially there are no matrix products, so the total optimizer compute is O(Ψ): it scales with the parameter count and, unlike the forward and backward passes, is completely independent of the batch size. DeepSpeed ships a hand-tuned SIMD “CPU-Adam” kernel for exactly this loop. The bottleneck on the CPU is not arithmetic but memory bandwidth — the step must stream the 12Ψ bytes of fp32 state through the cores once — which is precisely what CPU DRAM is good at.
Why the update is cheap enough for the CPU
The trade only works because of an asymmetry in how the two halves scale. The GPU’s forward and backward passes cost about 6Ψ FLOPs per token, so a step that processes a batch of B tokens burns roughly 6Ψ·B FLOPs. The CPU’s Adam update costs a fixed ~12Ψ FLOPs regardless of B.
The ratio of GPU work to CPU work is therefore about 6B / 12 = B/2 — it grows linearly with the batch. Process a handful of tokens per step and the puny CPU update dominates wall-clock time; process tens of thousands of tokens per step and the GPU compute towers over it, so the optimizer step becomes a rounding error you can hide behind the next backward pass. This is the design’s central insight: the one part of training whose cost does not grow with batch size is the one part you offload, because a large batch amortizes it into invisibility. Offload the forward pass instead and you would be shipping O(B) work to the wrong processor.
The PCIe transfer cost
The second cost is communication, and it obeys the same O(Ψ) law. Each step moves two things across the PCIe bus: the gradients out to the CPU (2Ψ bytes in fp16) and the freshly updated weights back to the GPU (2Ψ bytes in fp16). That is 4Ψ bytes of PCIe traffic per step — and ZeRO-Offload is built so this is provably the minimum: you cannot update on the CPU without sending gradients down and results back up.
Like the CPU compute, this volume is fixed per step and independent of batch size. Bandwidth is the constraint: PCIe Gen3 ×16 delivers roughly 12 GB/s per direction and Gen4 about 25 GB/s. Divide 4Ψ bytes by that rate to get the transfer time. Because it is a fixed cost per step, the same lever that hides the CPU update — a large batch, which stretches the GPU compute per step — is what buys you the time to overlap these copies with computation rather than stalling on them.
A worked example: a 10B model on one GPU
Take Ψ = 10×10^9 parameters. Full Adam state is 16Ψ = 160 GB — impossible on a 32 GB or even 80 GB GPU. ZeRO-Offload leaves 2Ψ = 20 GB of fp16 weights on the GPU and relocates 12Ψ = 120 GB of optimizer state (plus the streamed gradients) into CPU RAM, where 120 GB is ordinary. The GPU footprint for model state collapses from 160 GB to ~20 GB, and the freed VRAM absorbs activations — a 10B model now trains on a single 32 GB card.
Now the per-step overhead. PCIe traffic is 4Ψ = 40 GB; at Gen4’s 25 GB/s that is about 1.6 s of copying. The CPU-Adam step streams 12Ψ = 120 GB of state at, say, 80 GB/s DRAM bandwidth — roughly 1.5 s. If the GPU forward and backward take 6 s for a large batch, both the ~1.5 s update and the ~1.6 s of copies overlap behind it and cost almost nothing. Shrink the batch until the GPU pass is under 3 s, and that same fixed overhead becomes the thing you wait on.