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.…