CPU and NVMe Offloading

Every other technique on this branch reduces what must be resident. Offloading accepts the memory requirement and moves part of it somewhere larger and slower.

The hierarchy is steep — three orders of magnitude from HBM to NVMe — so the question is not whether it costs performance but whether the transfer can be hidden behind compute.

What to offload

The right candidates are tensors that are large, and used rarely.

Optimiser state is the ideal case: two thirds of the memory of a mixed-precision model, and touched exactly once per step. ZeRO-Offload puts it on the host along with the fp32 master weights, and runs the optimiser update on the CPU — the update is elementwise and bandwidth-bound, so a CPU is not badly matched to it, and the gradients have to cross the bus anyway.

Parameters are harder: used every layer, every step. ZeRO-Infinity offloads them and prefetches layer $\ell+1$ while layer $\ell$ computes, which works precisely when a layer’s compute time exceeds the next layer’s transfer time.

Rearranged, this is a statement about arithmetic intensity, and it is why offloading works far better during training than during decoding. Training a layer at batch 8 and sequence 4096 does thousands of FLOPs per parameter byte; decoding one token does two. Same weights, same bus, and the ratio decides everything.

The tiers

ZeRO-Infinity spans all three and reports trillion-parameter training on a single node, at a throughput that is a fraction of what the same model would achieve resident. The honest framing is the one its authors use: this is about accessibility, not speed. A model that trains slowly on hardware you have beats one that does not train at all on hardware you do not.

import torch
from torch import nn


class OffloadedLayers:
    """Keep one layer resident; prefetch the next on a side stream."""

    def __init__(self, layers: list[nn.Module], device="cuda"):
        # pin_memory: required for async H2D. Pageable memory silently
        # serialises through a staging buffer and the overlap is lost.
        self.cpu = [{k: v.pin_memory() for k, v in l.state_dict().items()} for l in layers]
        self.layers, self.device = layers, device
        self.stream = torch.cuda.Stream()

    def _prefetch(self, i: int) -> None:
        if i >= len(self.layers):
            return
        with torch.cuda.stream(self.stream):
            for k, v in self.cpu[i].items():
                self.layers[i].state_dict()[k].copy_(v, non_blocking=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self._prefetch(0)
        for i, layer in enumerate(self.layers):
            # Wait for this layer's weights, then start the next transfer
            # before computing — so the copy overlaps the matmul.
            torch.cuda.current_stream().wait_stream(self.stream)
            self._prefetch(i + 1)
            x = layer(x)
        return x

The ordering in that loop is the whole technique. Prefetching after the compute call rather than before serialises the two, and the code still runs and still produces correct results — just at half the speed, with no error to explain why.

The same arithmetic, with the inequality now failing badly: decoding has arithmetic intensity of about 2, so streaming weights over PCIe is straightforwardly bandwidth-bound and slow.

LLM in a Flash is the interesting response. It combines offloading with activation sparsity — if only 10% of the FFN neurons fire, only their rows need to be read from SSD, which reduces the transfer by the sparsity factor rather than trying to hide it. Against a baseline of the model not running at all, a 4–5× speed-up over naïve streaming is the difference between usable and not.