Reducing GPU HBM Pressure in JAX Training via Host Offloading

Training large language models (LLMs) increasingly encounters GPU high‑bandwidth memory (HBM) limits before compute capacity is fully used. Model weights, gradients, optimizer state, communication buffers and intermediate activations all contend for HBM. As model size, sequence length and batch size grow, HBM capacity often becomes the primary scaling bottleneck.

What is host offloading in JAX?

Host offloading in the open‑source JAX library moves selected activations into pinned host memory during the forward pass and streams them back to the GPU when needed in the backward pass. This is an alternative to activation rematerialization: instead of recomputing certain activations during the backward pass, they are reloaded from host memory.

Why is this especially useful on NVIDIA Grace Blackwell systems?

Host offloading benefits from very high CPU–GPU connectivity. On NVIDIA Grace Blackwell systems the NVIDIA Grace CPU and NVIDIA Blackwell GPU are connected via NVLink‑C2C with 900 GB/s bidirectional bandwidth, making pinned host memory a practical staging area for large activations. The Vera CPU and Rubin GPU configurations further raise coherent bidirectional bandwidth to 1.8 TB/s. High bandwidth alone is necessary but not sufficient: activation transfers must be overlapped with useful GPU work to avoid stalling training.

Performance results on MaxText workloads (NVIDIA GB200 NVL72, 128 GPUs)

The experiments used MaxText, a JAX LLM training framework that runs on the XLA compiler. All measurements reported below were taken on NVIDIA GB200 NVL72 systems with 128 GPUs. Two workloads were evaluated:

• Llama 3.1 405B — a dense decoder‑only transformer used to study targeted QKV activation offloading at a fixed batch.
• DeepSeek‑V3 671B — a sparse mixture‑of‑experts (MoE) model with multihead latent attention (MLA), used to study throughput and memory capacity effects.

DeepSeek‑V3 671B offloading policy

DeepSeek‑V3 671B contains 61 decoder layers: the first three use dense MLP blocks, the remaining layers use MoE blocks. The offload policy targets large MLA query and key/value projection intermediates and selected MoE up projection intermediates in the repeated MoE decoder layer. These intermediates are large enough to determine whether larger batch configurations fit in GPU memory.

Throughput improvements

MaxText reports throughput in TFLOPs/s/device (model FLOPs per training step divided by measured step time).

• DeepSeek‑V3 671B: with host offloading, Latency Hiding Scheduler (LHS), and pipelined transfers enabled, the workload reached 908.2 TFLOPs/s/device. That is 57% faster than activation rematerialization at the same batch configuration and 67.7% faster than offloading without LHS or pipelining.

For DeepSeek’s large MoE/MLA activation footprint, pipelined transfers provided a notable positive impact on total throughput; for some dense workloads, LHS alone may suffice to hide transfer latency. On Blackwell systems, XLA scheduling flags coordinate with hardware to perform asynchronous data movement over dedicated copy streams, enabling batch configurations that would stall on commodity clusters.

Increasing feasible batch size (DeepSeek example)

A comparison across activation placement policies shows that host offloading changes feasible batch configurations. Key results (DeepSeek‑V3 671B):

• Host offload + LHS + pipelined: micro batch 8, global batch 1024, 908.2 TFLOPs/s/device, GPU peak 165.2 GiB, host memory 145.1 GiB.
• No offload + LHS + rematerialization: micro batch 8, global batch 1024, 578.3 TFLOPs/s/device, GPU peak 151.3 GiB, host memory 0 GiB.
• Host offload without LHS/pipelining: micro batch 8, global batch 1024, 541.6 TFLOPs/s/device, GPU peak 145.6 GiB, host memory 145.1 GiB.
• No offload, saving activations on device: micro batch 2, global batch 256, 425.3 TFLOPs/s/device, GPU peak 113.3 GiB.
• No offload, save on device with micro batch 8/global 1024: OOM.

Host offloading enabled the micro batch 8 / global batch 1024 configuration that otherwise produced an out‑of‑memory error. The offload policy moves large intermediates from MLA, MoE and MLP blocks into host memory, freeing HBM for model state, communication buffers, runtime workspaces and active computation. Enabling LHS and pipelined transfers increases GPU peak memory (165.2 GiB vs 145.6 GiB) because additional copy buffers and prefetched activations are retained on the GPU to improve overlap and throughput — a trade‑off of capacity for performance.

Llama 3.1 405B results

The Llama 3.1 405B experiment ran 10 steps on synthetic data: batch size 2, sequence length 8192, FSDP=128, bfloat16 activations and NVFP4 4‑bit weight quantization. Results:

• No offload, LHS on: 2,669 TFLOPs/s/device, GPU peak 149.6 GiB, host 0 GiB.
• QKV offload + LHS: 2,746 TFLOPs/s/device (2.9% increase), GPU peak 149.9 GiB, host 70.9 GiB.
• QKV offload, LHS off: 2,569 TFLOPs/s/device, GPU peak 139.7 GiB, host 70.9 GiB.
• QKV offload + LHS + pipelining: 2,718 TFLOPs/s/device, GPU peak 151.0 GiB, host 70.9 GiB.

For this dense fixed‑batch example, LHS alone hid most transfer latency and provided the best throughput. QKV offloading primarily replaced backward‑pass rematerialization with transfers that could overlap compute and communication. The reported 70.9 GiB host memory is the total QKV activation storage across 126 layers, not the instantaneous GPU memory saved; with a scan loop enabled the backward pass only needs one layer’s QKV activations on GPU at a time.

When is host offloading most useful?

Host offloading is most beneficial when:

• GPU memory limits model size, sequence length or batch size, and
• Selected tensors are large enough to meaningfully reduce HBM pressure, and
• Offloading can replace expensive rematerialization or enable a larger batch configuration.

Performance depends on overlap: host offloading works best when the workload has sufficient compute, communication or other independent work to hide transfer latency. On NVIDIA GPUs, XLA features such as dedicated copy streams, LHS and pipelined host offloading help create this overlap. The approach is particularly advantageous on platforms like NVIDIA Blackwell GB200 and Blackwell Ultra GB300 that use NVLink‑C2C to bypass PCIe bottlenecks; the NVIDIA Vera Rubin platform offers still higher interconnect performance.

Host offloading is less helpful when tensors are small, when little independent work exists to overlap transfers, or when the workload is bottlenecked by resources other than memory. Always validate runtime memory with real runs because static estimates may exclude NCCL communication scratch space, cuDNN attention workspace and framework‑managed buffers.

How to get started

• Run a small representative JAX training job. Select large activations on expensive forward paths, enable offloading and measure runtime GPU memory, host memory use and step time.
• See the JAX host offloading tutorial for APIs (jax.remat, checkpoint policies, memory_kind="pinned_host") and parameter/optimizer state offloading with jax.device_put().

The MaxText experiments used the NGC JAX container. DeepSeek‑V3 671B runs used ghcr.io/nvidia/jax:deepseek_v3_maxtext, a customized container with MaxText integration and Transformer Engine MoE permutation optimizations. Recommended XLA flags used in the optimized configurations:

• --xla_gpu_enable_latency_hiding_scheduler=true
• --xla_gpu_enable_pipelined_host_offloading=true
• --xla_gpu_experimental_parallel_async_compute_limit=8

Use profiling tools such as NVIDIA Nsight Systems to confirm that device→host and host→device copies overlap with compute and NCCL communication.

Conclusion

Host offloading is a practical memory placement strategy when HBM capacity limits batch size, context length or model scale and when selected activations are large enough to matter. It is most effective when it replaces costly activation rematerialization or makes larger batch configurations feasible, and when transfers can be overlapped with compute and communication via scheduler and interconnect features.

Acknowledgments

Thanks to Jaroslav Sevcik, Sevin Varoglu, Tj Xu, Haixin Liu, Abhinav Goel, Md Fahim Faysal Khan, Stefano Bosisio and Jinxin Yang for technical contributions.