Loading…
Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading
NVIDIA Developer BlogTanya Lenz
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.
Related reading
NVIDIA Blackwell Tops MLPerf Training 6.0 with Industry-Leading Scale and Performance
NVIDIA submitted results across all MLPerf Training v6.0 benchmarks, setting performance records on workloads including the DeepSeek-V3 and GPT-OSS-20B Mixture of Experts models. To scale training up to 8,192 Blackwell GPUs, NVIDIA combined hardware cluster designs with scale-out networking via Spectrum-X Ethernet and Quantum InfiniBand. The engineering team eliminated CPU-GPU synchronization bottlenecks in token-dropless MoEs by implementing full-iteration CUDA graphs with synchronization-free operators and paged stashing. Further software optimizations included CuTe DSL kernel fusions, an MXFP8 attention block in the Transformer Engine, and transitioning router elementwise math to FP32. Across these benchmarks, the GB300 NVL72 platform achieved the fastest time to train at scale and the highest normalized per-accelerator performance.
Farshad Ghodsian