Large language model (LLM) training workloads increasingly run into GPU memory limits before compute is fully used. Model weights, gradients, optimizer states, communication buffers, and intermediate activations all compete for GPU high-bandwidth memory (HBM). As model size, sequence length, and batch size grow, HBM capacity often becomes the primary scaling bottleneck.

This post explains how host offloading in the open source Python library JAX reduces HBM pressure. This process, which is especially advantageous on NVIDIA Blackwell, moves selected activations to pinned host memory during the forward pass and streams them back when needed in the backward pass. Host offloading is an alternative to activation rematerialization. Instead of recomputing selected activations, the training step reloads them from host memory.

Why is host offloading advantageous on NVIDIA Grace Blackwell systems?

Host offloading is especially advantageous on NVIDIA Grace Blackwell systems. The NVIDIA Grace CPU and NVIDIA Blackwell GPU are connected through NVLink-C2C with 900 GB/s of bidirectional bandwidth, making pinned host memory a practical staging area for selected activations. The Vera CPU and Rubin GPU further improve on this by doubling the bidirectional speed to 1.8 TB/s of coherent bandwidth.