Generated by Codex with GPT-5

NVIDIA’s official Technical Blog published Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading, a July 10, 2026 post about treating CPU-attached memory as an active tier in the training memory hierarchy rather than leaving scarce GPU high-bandwidth memory to hold every activation needed by backpropagation.

The pressure comes from several consumers sharing HBM at once: model weights, gradients, optimizer state, collective-communication buffers, runtime workspaces, and intermediate activations. As models, batches, and sequences grow, training can therefore become capacity-bound before the GPU’s arithmetic units are fully occupied. The usual response is activation rematerialization, which discards selected forward-pass results and recomputes them during the backward pass. NVIDIA examines another trade: move carefully chosen activations into pinned host memory during the forward pass, then copy them back just before their gradients are calculated.

The mechanism is useful only when transfer time can be hidden. Grace Blackwell connects its CPU and GPU over coherent NVLink-C2C with 900 GB/s of bidirectional bandwidth, while NVIDIA says Vera Rubin doubles that figure to 1.8 TB/s. That bandwidth makes host memory a plausible staging tier, but it does not make movement free. The training stack must overlap device-to-host and host-to-device copies with useful computation and communication, or offloading simply replaces a recomputation cost with a data-transfer stall.

JAX and XLA supply the scheduling layer for that overlap. The evaluated MaxText configurations combine XLA’s Latency Hiding Scheduler with dedicated asynchronous copy streams and, for the heaviest workload, pipelined host offloading. Selected activations flow out while later forward work continues, then are prefetched back while the backward pass is busy elsewhere. XLA is also allowed to keep more asynchronous work in flight so copies can overlap with GPU kernels and NCCL collectives. In this design, host offloading is less a swap operation than a compiler-scheduled dataflow transformation.

The placement policy is selective rather than indiscriminate. For DeepSeek-V3 671B, a 61-layer sparse mixture-of-experts model with multihead latent attention, MaxText offloads large query and key/value projection intermediates plus selected MoE up-projection activations. The first three dense MLP layers use a similar targeted policy. These tensors are attractive because they are large enough to constrain feasible batch size and expensive enough to make recomputation costly. Smaller tensors or values with little independent work around their transfers would be poor candidates.

NVIDIA evaluated the design on 128 GPUs in a GB200 NVL72 environment. On DeepSeek-V3, optimized host offloading reached 908.2 TFLOPs/s/device at microbatch 8 and global batch 1024. That was 57% faster than activation rematerialization at the same batch configuration and 67.7% faster than host offloading without latency hiding or pipelined transfers. The contrast is the main result: moving data to the CPU is not inherently an optimization. It becomes one only when the runtime schedules the memory traffic as part of the training graph.

Offloading also changed which configurations could run. Saving activations on the device fit microbatch 2 and global batch 256, but the larger 8/1024 configuration ran out of memory. The optimized offload policy made that larger batch feasible by putting 145.1 GiB of activation storage in host memory. This gives the method two possible payoffs: it can replace expensive recomputation at a fixed batch, or it can free enough HBM to increase batch size, sequence length, or model scale.

The memory measurements expose an important subtlety. The fastest offloaded DeepSeek run peaked at 165.2 GiB of GPU memory, more than the 145.6 GiB used by the slower offload configuration without latency hiding and pipelining. Prefetch buffers and concurrent transfers consume HBM so that the GPU can remain busy. Offloading therefore does not monotonically minimize device memory; it creates a tunable trade between capacity and overlap. The best configuration may deliberately spend some recovered HBM on staging to gain much more throughput.

The dense Llama 3.1 405B experiment shows why the result should not be generalized into a universal speedup. With a fixed batch size of 2 and sequence length 8,192, offloading QKV activations and enabling the latency-hiding scheduler improved throughput by a modest 2.9%, from 2,669 to 2,746 TFLOPs/s/device. Pipelining transfers reduced that result slightly because the scheduler was already hiding most of the copy latency. Without latency hiding, throughput fell below the no-offload baseline. The right policy depends on model shape, activation footprint, recomputation cost, and the amount of independent work available for overlap.

This also clarifies how the system differs from naive virtual memory. Placement is explicit, model-aware, and measured. Engineers choose activations from expensive forward paths, observe real peak GPU and host memory rather than relying only on static estimates, and profile the timeline to verify that copies overlap with compute and collectives. Runtime workspaces and NCCL scratch buffers can invalidate paper capacity calculations, while an apparently asynchronous configuration can still expose transfer bubbles. NVIDIA specifically recommends using Nsight Systems to inspect the schedule.

The broader engineering takeaway is that accelerator memory capacity is becoming a whole-system scheduling problem. Once CPU-GPU links are fast enough, host DRAM can participate in training in the same way caches and storage tiers participate in other data systems: not as a transparent overflow area, but as a tier whose latency must be anticipated and hidden. Compiler scheduling, framework-level placement policy, copy engines, interconnect bandwidth, and model architecture all determine whether the extra tier increases Goodput.

The post also makes a useful case for workload-specific optimization. DeepSeek’s large MoE and attention intermediates produced a dramatic win; Llama’s targeted QKV policy produced a small one. A single feature flag cannot erase that difference. The disciplined approach is to identify the current limiter, compare transfer against recomputation, choose only tensors that materially affect capacity, and validate end-to-end step time. Optimizing memory in isolation can easily reduce throughput.

Host offloading is therefore best understood as a way to reshape the boundary between memory and compute. It can unlock configurations that HBM alone cannot hold and turn otherwise idle CPU memory into a productive resource, but only when the system treats movement as schedulable work. The durable lesson extends beyond JAX or Blackwell: new memory tiers create value when the software stack exposes placement, coordinates transfers with the execution graph, and measures the full pipeline rather than assuming that more apparent capacity automatically means faster training.