Memory Offload Across the Hierarchy

"They moved me to host memory to make room. I am still part of the model; I just take the scenic route to every step now."

A Shard That Believes It Is the Whole Model

Sharding the model across devices, the subject of Chapter 16, divides parameters, gradients, and optimizer state evenly across the GPUs you have, so that no single device holds the whole model. That works until you run out of GPUs. A practitioner with one accelerator, or a small budget-limited pool, cannot shard their way to a model that needs a hundred gigabytes of state when the cluster only offers sixteen. The state has to live somewhere, and if it cannot all live in GPU memory, the only direction left is down: into the host CPU's much larger DRAM, and below that into the still larger NVMe SSD attached to the machine. Training proceeds by hauling each piece of state back up to the GPU exactly when the step needs it, then sending it back down to make room. This is memory offload, and the technique that made it practical at scale is the ZeRO-Offload and ZeRO-Infinity line of work built into DeepSpeed.

1. Why the Hierarchy Exists, and Why It Is So Uneven Beginner

Every machine that trains models has a memory hierarchy, and the defining fact about it is that the tiers differ by orders of magnitude in both directions. GPU HBM is measured in tens of gigabytes and delivers something like a terabyte or more per second; host DRAM is measured in hundreds of gigabytes but is reached only across the PCIe link at a few tens of gigabytes per second; an NVMe SSD is measured in terabytes but sustains only single-digit gigabytes per second. Capacity grows as you descend, and bandwidth collapses. This is the locality principle from Chapter 2 applied inside a single node: data you touch constantly belongs in the fast, scarce tier, and data you touch rarely can be exiled to a slow, abundant one without much harm.

Mixed-precision Adam training makes the tiering natural because not all of a model's state is touched with the same frequency. Per parameter, a typical setup keeps a half-precision parameter and gradient (two bytes each) that the forward and backward passes touch constantly, plus a full-precision master copy and the two Adam moments (twelve bytes total) that are read and written only once per optimizer step. That twelve-byte optimizer trio is three quarters of the memory and the least frequently touched, which makes it the obvious thing to push downhill. ZeRO-Offload's central insight is exactly this: keep the hot two-byte params and grads on the GPU, move the cold optimizer state to the CPU, and run the optimizer update on the CPU so the heavy state never has to climb back up.

2. The Bandwidth-Capacity Trade, Made Quantitative Intermediate

To decide whether offload is worth it, you need the cost in numbers, and the roofline reasoning of Chapter 3 gives the shape of the answer. A training step does a fixed amount of GPU compute and, under offload, a fixed amount of data movement across a slow link. Let a model have $P$ parameters, let $b_{\text{off}}$ be the bytes of state offloaded per parameter, and let $B_{\text{link}}$ be the bandwidth of the slow link the offloaded state crosses. Each step must move that state down and back up, so the transfer time is

The GPU compute for a step depends on the model size and the number of tokens $T$ processed, at roughly six floating-point operations per parameter per token (the standard forward-plus-backward count), so with sustained throughput $F$ in FLOP/s,

The two can overlap: while the GPU computes one slice of the step, the system prefetches the next slice's state and writes back the previous slice's. If a fraction $\rho$ of the transfer hides behind compute, the step takes $t_{\text{compute}} + (1 - \rho)\,t_{\text{transfer}}$. Two readings fall out immediately. First, the slowdown shrinks as $T$ grows, because more compute means more cover for the transfer; this is why offload pairs well with large micro-batches. Second, the slowdown grows sharply as $B_{\text{link}}$ falls, which is why moving from DRAM (PCIe, tens of GB/s) to NVMe (single-digit GB/s) hurts far more than moving from HBM to DRAM. The capacity, meanwhile, is simply how much state the combined tiers can hold: the largest model that fits is $(C_{\text{HBM}} + C_{\text{DRAM}} + C_{\text{NVMe}}) / b_{\text{total}}$, growing with every tier you enlist. The code below puts both sides of the trade on one screen.

import sys

# Mixed-precision Adam memory accounting, bytes per parameter.
BYTES_PARAM_GRAD = 4   # fp16 param + fp16 grad: touched every step, stays on GPU
BYTES_OPT        = 12  # fp32 master + momentum + variance: cold, offloadable
BYTES_TOTAL      = BYTES_PARAM_GRAD + BYTES_OPT   # 16 B/param

# One modest node and its three tiers (capacity GB, bandwidth GB/s).
GPU_HBM_GB, BW_HBM   = 16.0,  1500.0   # on-package HBM
HOST_DRAM_GB, BW_PCIE = 256.0,   25.0   # host DRAM over PCIe 4.0 x16
NVME_GB, BW_NVME      = 4000.0,    5.0   # NVMe SSD sustained

TOKENS_PER_STEP = 2 ** 13          # 8192 tokens of per-GPU work to hide behind
FLOPS_PER_PARAM_PER_TOKEN = 6.0    # fwd + bwd, the standard 6N rule
GPU_TFLOPS = 120.0                 # sustained fp16 throughput, TFLOP/s
OVERLAP = 0.85                     # fraction of transfer prefetching can hide

def largest_fit(cap_gb):           # max billions of params whose 16 B/param fits
    return cap_gb * 1e9 / BYTES_TOTAL / 1e9

def step_time(n_params_b, link_bw):
    n = n_params_b * 1e9
    compute  = n * TOKENS_PER_STEP * FLOPS_PER_PARAM_PER_TOKEN / (GPU_TFLOPS * 1e12)
    transfer = 0.0 if link_bw is None else 2 * n * BYTES_OPT / (link_bw * 1e9)
    return compute, transfer, compute + transfer * (1 - OVERLAP)

print("CAPACITY: largest model that fits as state moves down the hierarchy")
cap = {"HBM only": largest_fit(GPU_HBM_GB),
       "HBM + CPU DRAM": largest_fit(GPU_HBM_GB + HOST_DRAM_GB),
       "HBM + CPU + NVMe": largest_fit(GPU_HBM_GB + HOST_DRAM_GB + NVME_GB)}
for name, b in cap.items():
    print(f"  {name:<18}{b:>7.2f} B params   ({b/cap['HBM only']:>5.1f}x)")

print("\nTHROUGHPUT COST: per-step time for a 3 B-param model under each tier")
base = None
for name, bw in [("HBM only", None), ("HBM + CPU", BW_PCIE), ("HBM + CPU + NVMe", BW_NVME)]:
    comp, xfer, step = step_time(3.0, bw)
    base = step if base is None else base
    print(f"  {name:<18}compute {comp:6.3f}s  transfer {xfer:7.3f}s  "
          f"step {step:6.3f}s  ({step/base:.2f}x)")

CAPACITY: largest model that fits as state moves down the hierarchy
  HBM only             1.00 B params   (  1.0x)
  HBM + CPU DRAM      17.00 B params   ( 17.0x)
  HBM + CPU + NVMe   267.00 B params   (267.0x)

THROUGHPUT COST: per-step time for a 3 B-param model under each tier
  HBM only          compute  1.229s  transfer   0.000s  step  1.229s  (1.00x)
  HBM + CPU         compute  1.229s  transfer   2.880s  step  1.661s  (1.35x)
  HBM + CPU + NVMe  compute  1.229s  transfer  14.400s  step  3.389s  (2.76x)

The output states the trade in one screen. The capacity column climbs by 17x and then by 267x as each tier joins; the throughput column shows the price, a 1.35x slowdown for the DRAM hop and 2.76x once NVMe is in the loop. Note how the transfer time for NVMe (14.4 seconds) dwarfs the compute (1.23 seconds), yet the exposed step time is only 3.39 seconds: prefetching hides 85 percent of the movement behind the GPU's work, which is the single engineering trick that keeps NVMe offload from being unusable. The numbers also make the regime boundaries visible, which is the subject of the next section.

3. When Offload Is the Right Tool, and When It Is Not Intermediate

Offload is a capacity technique, not a speed technique, and reading the numbers backward tells you exactly when to reach for it. It is the right tool when you are capacity-constrained or budget-constrained: when the model simply will not fit on the GPUs you can afford, and a slower run is strictly better than no run at all. A researcher fine-tuning a 13-billion-parameter model on a single workstation GPU, or a startup training on a handful of consumer cards rather than a rented cluster, is exactly the case offload was built for. The slowdown is real but bounded, and the alternative is not a faster run; it is an out-of-memory error.

It is the wrong tool when throughput is what you are optimizing and you have the hardware to fit the model without it. A frontier training run that already shards across thousands of GPUs and is measured in cost-per-token does not want to multiply its step time to save memory it can buy. In that regime the PCIe and NVMe links become the bottleneck the whole system waits on, and the right answer is more GPUs and better sharding, not slower memory. The honest framing is that offload converts a hard capacity wall into a soft throughput penalty, which is a wonderful deal when you have hit the wall and a bad deal when you have not.

4. From the Model to the Library Beginner

The arithmetic in Code 18.7.1 captures the trade, but the real machinery (prefetching the next slice while computing the current one, partitioning state across DRAM and NVMe, running the optimizer on the CPU, and overlapping every transfer with compute) is intricate enough that you should never build it by hand. DeepSpeed packages all of it behind a configuration file, as part of the same ZeRO family introduced for sharding in Chapter 16. ZeRO-Offload adds the CPU tier and ZeRO-Infinity adds the NVMe tier, both as a few keys in a JSON config rather than a change to your training loop.

# deepspeed_config.json: ZeRO stage 3 with offload to CPU and NVMe.
zero_config = {
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu",  "pin_memory": True},
    "offload_param":     {"device": "nvme", "nvme_path": "/local_nvme"},
    "overlap_comm": True               # hide transfers behind compute
  },
  "train_micro_batch_size_per_gpu": 4,
  "fp16": {"enabled": True}
}

# In code: wrap the model and let DeepSpeed place the tiers for you.
import deepspeed
engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config=zero_config)
for batch in loader:
    loss = engine(batch)               # forward; params streamed up as needed
    engine.backward(loss)              # grads computed on GPU
    engine.step()                      # Adam runs on CPU over offloaded state

With state spread across the memory hierarchy, a single node can now hold a model far beyond its GPU, but a run spread across many such nodes and many slow links becomes much harder to see into when something goes wrong. Knowing whether the bottleneck is compute, the PCIe link, or a straggler node requires instrumentation, which is where the chapter turns next, in Section 18.8 on monitoring and debugging distributed training.