Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading | NVIDIA Technical Blog

Technical Blog

Subscribe

Related Resources

Agentic AI / Generative AI

Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading

Jul 10, 2026

By Jane (Zhenying) Liu, Ming Huang, Johannes Reifferscheid, Michael Goldfarb and Tejash Shah

Like

Discuss (0)

AI-Generated Summary

Like

Dislike

Host offloading in JAX significantly alleviates HBM pressure during large language model training by moving selected activations to pinned host memory and streaming them back for the backward pass, enabling larger model, batch, and sequence configurations on NVIDIA Grace Blackwell and Vera Rubin platforms due to their high-bandwidth NVLink-C2C interconnects.<br>MaxText experiments on NVIDIA GB200 NVL72 with DeepSeek-V3 671B and Llama 3.1 405B demonstrated that host offloading, when combined with Latency Hiding Scheduler and pipelined transfers, delivered up to 57% throughput improvement over activation rematerialization and unlocked batch sizes otherwise limited by GPU memory, with the greatest gains observed for large sparse MoE models.<br>Achieving optimal performance with host offloading requires careful overlap of activation transfers with compute and communication, which is managed via XLA custom scheduling flags and dedicated copy streams, and should be validated using profiling tools such as NVIDIA Nsight Systems to ensure asynchronous data movement and memory utilization align with expected performance gains.

AI-generated content may summarize information incompletely. Verify important information. Learn more

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.

High-bandwidth CPU-GPU connectivity helps make host offloading practical, but bandwidth alone is not enough. To improve performance, activation transfers must overlap with useful GPU work.

Performance results on MaxText workloads

The experiments use MaxText, a JAX LLM training framework that uses the Accelerated Linear Algebra (XLA) compiler for training at scale on NVIDIA GPUs. All results were measured on NVIDIA GB200 NVL72 systems using 128 GPUs.

The evaluation uses two MaxText workloads. Llama 3.1 405B is a dense decoder-only transformer model used to study targeted query, key, and value (QKV) activation offloading at a fixed batch size. DeepSeek-V3 671B is a sparse mixture-of-experts (MoE) model with multihead latent attention (MLA), used to study both throughput and memory-capacity effects.

DeepSeek-V3 671B offloading policy

DeepSeek-V3 671B has 61 decoder layers: the first three use dense multilayer perceptron (MLP) blocks, while the remaining layers use MoE blocks.

Figure 1 shows the activation offloading policy for the repeated MoE decoder layer, which dominates the stack. The first three dense MLP layers are not shown; they use a similar policy, with selected up and down projection outputs offloaded.

In the MoE layer shown in Figure 1, the policy offloads selected MLA query and key/value projection intermediates and selected MoE up projection intermediates. These activations are large enough to impact whether larger batch configurations fit.

Figure 1. DeepSeek-V3 671B repeated MoE decoder layer activation offloading policy

Training throughput improvements

MaxText reports throughput as TFLOPs/s/device, computed from model FLOPs per training step divided by measured step time.

DeepSeek-V3 671B

Figure 2 compares DeepSeek-V3 671B throughput across activation placement policies. With offloading, LHS, and pipelined transfers enabled, the workload reached 908.2 TFLOPs/s/device. This was 57% faster than activation rematerialization at the same batch configuration and...