NVIDIA published research on reducing high-bandwidth memory bottlenecks in JAX-based LLM training using host memory offloading techniques.

Large language model training workloads increasingly encounter GPU memory limits before compute is fully utilized. 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.

Host offloading in the open-source JAX library reduces HBM pressure by moving selected activations to pinned host memory during the forward pass and streaming them back when needed in the backward pass. This approach offers an alternative to activation rematerialization, which requires recomputing selected activations instead of reloading them from host memory.

Host offloading is particularly advantageous on NVIDIA Grace Blackwell systems, where the NVIDIA Grace CPU and NVIDIA Blackwell GPU connect 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 performance, doubling the bidirectional speed to 1.8 TB/s of coherent bandwidth. However, high-bandwidth CPU-GPU connectivity alone is insufficient; to improve performance, activation transfers must overlap with useful GPU work.

Experiments were conducted using MaxText, a JAX LLM training framework that leverages the Accelerated Linear Algebra (XLA) compiler for large-scale training on NVIDIA GPUs. All results were measured on NVIDIA GB200 NVL72 systems using 128 GPUs across two workloads: Llama 3.1 405B, a dense decoder-only transformer model used to study targeted query, key, and value (QKV) activation offloading at fixed batch sizes, and DeepSeek-V3 671B, a sparse mixture-of-experts model with multihead latent attention (MLA), used to study both throughput and memory-capacity effects.

DeepSeek-V3 671B contains 61 decoder layers: the first three use dense multilayer perceptron (MLP) blocks, while the remaining layers employ MoE blocks. The activation offloading policy for the repeated MoE decoder layer, which dominates the stack, offloads selected MLA query and key/value projection intermediates and selected MoE up projection intermediates. These activations are large enough to determine whether larger batch configurations fit in memory.

With offloading, Latency Hiding Scheduler (LHS), and pipelined transfers enabled, DeepSeek-V3 671B reached 908.2 TFLOPs/s/device—57% faster than activation rematerialization at the same batch configuration and 67.7% faster than offloading without LHS or pipelining. Unlike dense Llama workloads where LHS alone suffices to hide latency, the massive activation footprint of DeepSeek-V3 MoE and MLA layers means that pipelined transfers provide a distinct, positive impact on total throughput. This performance advantage reflects NVIDIA's tight co-design of software and hardware: on Blackwell systems, XLA custom scheduling flags work in concert with dedicated copy streams to ensure data moves asynchronously. This integration enables platforms to unlock massive batch configurations that remain inaccessible for architectures lacking compiler-to-interconnect integration.

For capacity comparison, saving selected activations on device allowed micro batch 2 and global batch 256, while optimized host offloading enabled micro batch 8 and global batch 1024. Without offloading or rematerialization, the device encounters an Out-of-Memory error at micro batch 8 and global batch 1024. Host offloading made this configuration feasible by moving selected activation storage out of GPU memory, leaving more HBM available for model state, communication buffers, runtime workspaces, and active computation. With LHS and pipelined transfers enabled, the offload configuration uses 165.2 GiB of GPU memory, compared with 145.6 GiB without those optimizations. This increase reflects keeping more copy buffers and prefetched activations in GPU memory to overlap transfers with computation, trading some memory capacity for better overlap and higher throughput.

The Llama 3.1 405B experiment ran 10 steps on synthetic data with batch size 2, sequence length 8,192, fully sharded data parallelism set to 128, and bfloat16 activations with NVFP4 4-bit weight quantization. QKV activation offloading with Latency Hiding Scheduler improved throughput from 2,669 to 2,746 TFLOPs/s/device—a 2.9% increase over the baseline without offloading. Disabling LHS reduced QKV offload throughput to 2,569 TFLOPs/s/device, demonstrating that host offloading depends on effective overlap with other GPU work. For this configuration, LHS alone provided the best throughput at 2,746 TFLOPs/s/device without pipelining, compared with 2,718 TFLOPs/s/device with pipelining, since LHS already hides most transfer latency behind compute and communication.

The 70.9 GiB host memory value represents the total QKV activation storage across all 126 layers, not the amount of GPU memory saved at a single moment. At batch size 2 and sequence length 8,192, one layer's bfloat16 QKV activations require approximately 576 MiB: 512 MiB for query and 32 MiB each for key and value. With the scan loop enabled for layers, the backward pass processes one layer at a time, requiring only one layer's QKV activations on the GPU at once. In this workload, QKV offloading primarily optimizes performance by replacing backward-pass QKV rematerialization with transfers that overlap with compute and communication, while GPU peak memory remains dominated by model state, communication buffers, and runtime workspaces. The dense Llama 3.1 405B model demonstrates smaller gains than the sparse DeepSeek V3 671B, but reveals the same underlying mechanism: targeted QKV offloading replaces backward-pass rematerialization with transfers that overlap with compute.