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

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.

Takeaways

1. 1

   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.

2. 2

   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.

3. 3

   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.