Offload to CPU and NVMe
A GPU is not the only memory in the machine. A node has system memory, usually far more of it than the GPUs have, and it has solid-state storage beyond that. State that is needed rarely can be kept in those slower places and brought in when required.
The optimizer state is the natural candidate, because it is twelve of the sixteen bytes and it is touched exactly once per step. Holding it in system memory removes it from the GPU entirely.
Two versions of the same idea
The simpler one parks the state in system memory and copies it to the GPU for the update. The more effective one performs the update on the CPU as well, so the state never moves and the GPU is freed of that work altogether.
Copying that back and forth over PCIe every step would cost 12 seconds, which is longer than the whole computation. Doing the update on the CPU avoids the round trip entirely, which is why the second version is the one that gets used.
What it costs
The cost is the link between the CPU and the GPU, which is PCIe. It carries roughly 60 GB per second, which is far below NVLink and comparable to a network cable. Offload therefore converts a memory problem into a bandwidth problem. It earns its place in a run whose binding constraint is GPU memory.
| Link | GB/s | Carries |
|---|---|---|
| NVLink, inside a node | 450 | gradients between GPUs |
| PCIe, CPU to GPU | 60 | offloaded state and gradients |
| InfiniBand, between nodes | 50 | gradients across the cluster |
When it is the wrong tool
If the run is already limited by the wire rather than by memory, offload makes it worse: it adds traffic to solve a problem the run did not have. The check is the one from the cost page — if communication is already a large fraction of step time, spend the effort on moving less, not on moving more to a slower place.
NVMe extends the same idea one step further out, and the same reasoning applies with worse numbers. It makes runs possible that would otherwise not fit at all, at a speed that only makes sense when the alternative is not running.