Optimizer State Offloading to CPU/NVMe

What Adam actually costs, per parameter

Start by counting bytes for a single parameter under standard mixed-precision training, because a total in gigabytes tells you nothing you can reason with. The GPU holds a bf16 parameter (2 bytes) and, during backward, a bf16 gradient (2 bytes). The optimizer holds three more tensors, and all three are fp32 for numerical reasons: the master weight (4 bytes), the first-moment estimate m (4 bytes), and the second-moment estimate v (4 bytes).

So Adam costs 12 bytes per parameter of persistent state on top of the 4 bytes of working copies — three quarters of your static footprint. For a 7B model that is roughly 84 GB of optimizer state against 28 GB of parameters and gradients. That ratio, not any particular model size, is why the optimizer is where you look first.

Why optimizer state is the first thing you move

Memory pressure has three sources: parameters, activations, and optimizer state. They are not equally movable, and the deciding property is access frequency, not size.

Activations are produced and consumed continuously through forward and backward; every layer touches them, often several times. Moving them across a host link means paying that link on the critical path repeatedly per step, which is why the standard answer for activations is recomputation, not offload — you spend FLOPs, which are abundant, rather than PCIe bytes, which are not. Parameters are read every forward and backward, so they are similarly hot.

Optimizer state is different in kind: it is touched exactly once per optimization step, at the very end, after the last gradient has been produced. Nothing in forward or backward reads m or v. That single-touch property is what makes it the cheapest thing to exile: you move the largest block of memory across the slowest link in the machine, but only once per step.

Two designs: CPU-side step versus streaming state back

Having decided the state lives in host DRAM, there are two ways to actually perform the update, and the difference between them is the single most important number in this article.

Stream state back to the GPU. Gradients are already in HBM, so you pull the three fp32 tensors up (12 bytes/param), run the fused CUDA Adam kernel, and write all three modified tensors back down (12 bytes/param). All three change every step, so nothing can be skipped: 24 bytes/param of PCIe traffic per step.

Run the step on the CPU. You push the bf16 gradients down (2 bytes/param), the CPU updates m, v and the master weight in place in DRAM, casts the result, and you pull the updated bf16 parameters back up (2 bytes/param). Total: 4 bytes/param. That is a 6× reduction in link traffic, and it is why every serious offload implementation computes the step host-side. (If your framework reduces gradients in fp32, the figure is 6 bytes/param rather than 4 — still four times better.)

Why running Adam on the CPU is not absurd

The obvious objection is that a CPU is orders of magnitude slower than a GPU, so moving compute onto it should be self-defeating. It is not, because the Adam step is not compute at all — it is memory traffic wearing a thin coat of arithmetic.

Count the accesses per parameter: read m, v, the master weight and the gradient; write m, v and the master weight. That is roughly 28 bytes of DRAM traffic for about a dozen flops — an arithmetic intensity well under one flop per byte. On this workload a GPU's tensor cores are irrelevant; only bandwidth to wherever the state lives matters. Running a memory-bound operation on a slower memory system costs a proportional slowdown on something that was already a rounding error in step time.

The qualification matters, though. A multi-channel server socket has host bandwidth in the hundreds of GB/s (illustratively, ~200 GB/s), which puts the 7B step around a second. A dual-channel desktop board is more like 50–100 GB/s, two to four times worse — exactly the machine where people most want offload to work.

The PCIe budget for one step

Now price the link. A PCIe Gen4 x16 slot has a theoretical ceiling around 32 GB/s per direction; achievable throughput with well-formed pinned transfers is meaningfully lower, and ~25 GB/s is a reasonable illustrative planning figure (Gen5 roughly doubles both numbers). Treat that as a budget you spend once per step.

For a 7B model with CPU-side Adam, the traffic is 4 bytes/param, or 28 GB per step split across both directions. At 25 GB/s that is roughly 1.1 seconds of link time per step. Under the streaming-state-back design the same model would move 168 GB and burn nearly seven seconds — a number that immediately explains why that design is not used.

Two adjustments before trusting the figure. This is per GPU, so if the state is sharded the per-GPU share falls accordingly. And the host link is frequently contended — other GPUs on the same switch, NIC traffic for gradient reduction, and data-loader reads all compete for it. Budget conservatively.

Does it hide behind backward? Do the seconds

Offload is free if and only if the transfer overlaps with compute you were going to do anyway. The overlap you get is real: backward produces gradients last layer first, so the moment a layer's gradient is complete it can start its journey to the host while the remaining layers are still computing. The question is whether backward lasts long enough to cover the whole 1.1 seconds.

Backward is about 4 × N × T FLOPs for N parameters and T tokens processed on that GPU per step. Take the 7B model at an illustrative 400 TFLOP/s of achieved throughput:

This is the whole decision, and it reframes offload as a batch-size question. Small per-GPU batches leave the link exposed; large ones bury it. If your step is short, offload will hurt, and the cheapest fix is often to increase gradient accumulation rather than to abandon the plan.

Pinned memory, or the whole plan collapses

Everything above assumes the transfer overlaps with compute. That assumption is only valid if the host buffer is page-locked, and this is the detail that most often turns a sound design into a mysterious slowdown.

The DMA engine can only move data to or from physical pages guaranteed not to be swapped or relocated. Hand the driver a normal pageable allocation and it cannot DMA from it directly; it stages the data through an internal pinned bounce buffer, which means an extra CPU-side memcpy and, critically, a copy that cannot be asynchronous. You lose roughly half your effective bandwidth to the extra copy, then lose the overlap entirely, because the transfer now serialises against the stream instead of riding alongside backward. A 1.1-second hidden transfer becomes a two-second exposed one.

So allocate the host-side state as pinned memory up front and keep it for the run's lifetime — pinning is itself expensive and synchronising. The counterweight: pinned pages are unswappable, so pinning tens of gigabytes on a host with modest RAM starves the OS and the data loader.

The penalty you should expect, and how it scales

Expect a real but bounded cost. With pinned buffers, a CPU-side fused Adam and a per-GPU batch large enough to cover the link, a well-tuned offload typically lands in the region of a modest tens-of-percent step-time increase; with any of those three missing it can be a multiple.

The exposed portion has a specific shape. Because gradients arrive last-layer-first, the CPU pipelines updates against the tail of backward — but the first layers' gradients only exist once backward has finished, so their transfer, update and return trip have nothing left to hide behind. That leading-layer tail is the residual penalty, and it is why overhead never quite reaches zero however large the batch.

The scaling is more forgiving than intuition suggests. Both the transfer (4 bytes × N) and the backward pass (4 × N × T) are linear in parameter count, so the ratio is independent of model size. Doubling the model does not worsen the percentage penalty — it worsens your host RAM problem instead.

NVMe as the next tier, and where it stops

Host DRAM is finite. At 12 bytes/param, a node with 1.5 TB of RAM tops out around 125B parameters of optimizer state, and that is before the data loader, the pinned staging buffers and the OS get their share. The next tier down is NVMe, and the DeepSpeed lineage extends the same idea to it.

The economics change sharply. A single Gen4 NVMe drive delivers on the order of 7 GB/s sequential, roughly a quarter of the host link, so matching PCIe-class throughput needs several drives striped in parallel — and the aggregate still shares the bus with everything else. You also inherit an I/O stack: without O_DIRECT and an asynchronous submission path, the page cache adds a second full copy of every byte.

Durability is the underrated limit. Rewriting the full state every step writes terabytes per hour; consumer drives with modest endurance ratings can be consumed in weeks. NVMe offload buys you a model that fits at all — not one that trains at a competitive rate.

When offload is the wrong answer

Offload is a way to trade time for capacity on a GPU count you cannot change. When you can change it, sharding is almost always better: interconnects between GPUs run an order of magnitude faster than the host link, and splitting state across devices reduces per-GPU footprint without ever crossing PCIe on the critical path. If adding GPUs is possible, add GPUs.

Before either, ask the question that dissolves the problem most often: how many parameters are actually trainable? The 12 bytes/param term scales with trainable parameters, not total ones. A LoRA adapter on a 7B model takes 84 GB of optimizer state down to well under a gigabyte, and the reason to offload simply evaporates. For the consumer-hardware fine-tuning case this is the first move, not the last resort.

Two more disqualifiers. If activations dominate your footprint, offloading the optimizer solves the wrong problem — recompute instead. And inference has no optimizer state at all, so none of this applies to a serving bottleneck.