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.