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

[NVIDIA Developer Blog](https://yomu.fyi/company/nvidia-developer-blog) · Tanya Lenz · Jul 10, 2026

**Type:** Benchmark

## Summary

Large language model training workloads frequently face GPU high-bandwidth memory (HBM) capacity constraints before compute limits are reached. To alleviate this memory pressure in JAX training workloads, host offloading transfers selected forward-pass activations to pinned host memory over high-speed NVLink-C2C interconnects and reloads them during the backward pass instead of recomputing them. Benchmarks executed on 128 GPUs within NVIDIA GB200 NVL72 systems using MaxText demonstrated significant improvements. For DeepSeek-V3 671B, host offloading combined with the Latency Hiding Scheduler and pipelined offloading reached 908.2 TFLOPs/s/device, surpassing activation rematerialization throughput by 57% and unlocking larger batch configurations. Similarly, QKV activation offloading on Llama 3.1 405B achieved a 2.9% throughput increase when transfer latency was hidden behind computation.

## Context

In large language model training, model weights, gradients, optimizer states, communication buffers, and intermediate activations compete for GPU high-bandwidth memory, making HBM capacity the primary scaling bottleneck as model sizes, batch sizes, and sequence lengths increase.

## Approach / What changed

JAX workloads offload selected activations to pinned host memory over high-bandwidth CPU-GPU interconnects during the forward pass and reload them for the backward pass, utilizing XLA custom flags including Latency Hiding Scheduler and pipelined offloading to overlap transfers with computation.

## Takeaways

- Host offloading with LHS and pipelined transfers enabled DeepSeek-V3 671B to reach 908.2 TFLOPs/s/device, a 57% throughput improvement over activation rematerialization.
- Host offloading made a micro batch 8 and global batch 1024 configuration feasible for DeepSeek-V3 671B, whereas saving activations on-device caused out-of-memory errors.
- QKV activation offloading with Latency Hiding Scheduler on Llama 3.1 405B increased throughput from 2,669 to 2,746 TFLOPs/s/device, relying on effective overlap with computation.

**Tags:** [LLMs](https://yomu.fyi/topic/llm), [Machine Learning](https://yomu.fyi/topic/machine-learning), [Performance](https://yomu.fyi/topic/performance), [Python](https://yomu.fyi/topic/python), [Scalability](https://yomu.fyi/topic/scalability)

[Read original post](https://developer.nvidia.com/blog/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading)