How JAX Host Offloading Breaks the GPU HBM Wall in LLM Training

The Bottleneck Nobody Talks About: HBM, Not FLOPs

Every team scaling LLMs eventually hits the same wall. Your GPUs are not compute-bound—they are memory-bound. Model weights, gradients, optimizer states, communication buffers, and intermediate activations all fight for the same HBM pool. The moment you try to push sequence length or batch size, you get an OOM.

This is where host offloading enters the picture. Instead of recomputing activations during the backward pass (rematerialization), you stream them to pinned host memory during the forward pass and pull them back on demand. The trade-off flips from "extra compute" to "extra bandwidth."

On commodity clusters, this trade is a losing proposition—PCIe is too slow. On NVIDIA Grace Blackwell, it becomes a superpower. The Grace CPU and Blackwell GPU are wired through NVLink-C2C at 900 GB/s bidirectional, with Vera Rubin doubling that to 1.8 TB/s. Suddenly, host memory is a legitimate staging area, not a penalty box.

If you're tracking how infrastructure vendors are rethinking memory hierarchies under real workloads, this analysis of Cloudflare's headless browser migration is a great parallel case study.

The Results: DeepSeek-V3 671B and Llama 3.1 405B

All benchmarks ran on NVIDIA GB200 NVL72 with 128 GPUs, using MaxText (JAX + XLA) as the training framework.

DeepSeek-V3 671B (MoE + MLA)

The policy offloads selected MLA query/key-value projections and MoE up-projection intermediates—activations large enough to dictate whether a batch config even fits.

That's 57% faster than rematerialization and 67.7% faster than naive offloading without LHS or pipelining. Critically, row 5 shows the real win: without offloading, micro batch 8 / global batch 1024 simply does not fit. Offloading makes it feasible.

Llama 3.1 405B (Dense)

A 2.9% gain over baseline—smaller than DeepSeek, but revealing. Here LHS alone hides most latency, so pipelining adds nothing. For dense models, offloading is a performance optimization. For MoE models with massive activation footprints, it's a capacity unlock.

The XLA Flags That Make It Work

# Enable the latency hiding scheduler
--xla_gpu_enable_latency_hiding_scheduler=true

# Enable pipelined host offloading (critical for MoE workloads)
--xla_gpu_enable_pipelined_host_offloading=true

# Increase in-flight async work so LHS can overlap copies with NCCL collectives
--xla_gpu_experimental_parallel_async_compute_limit=8

The third flag is underrated. It gives the scheduler headroom to overlap device-to-host copies with collective communication. Without it, you're leaving throughput on the table.

Limitations and Caveats

Host offloading is not a silver bullet. Here's where it fails:

• Small tensors: If offloading individual activations doesn't meaningfully reduce HBM pressure, the transfer cost dominates.
• Low-overlap workloads: If there's not enough independent compute or communication to hide transfer latency, you lose performance. Llama 3.1 405B shows this—LHS was already sufficient, so pipelining was neutral.
• Non-memory bottlenecks: If your workload is bound by NCCL collectives or kernel launch overhead, offloading won't help.
• Static estimates lie: Runtime memory includes NCCL scratch, cuDNN attention workspace, and framework buffers. Always profile with real runs.

One more thing worth noting: the offload run used more GPU memory (165.2 vs 145.6 GiB) because it keeps copy buffers and prefetched activations resident. You're trading HBM capacity for throughput. That's a deliberate choice, not a free lunch.

How to Get Started

1. Start with a small representative JAX training run—not a toy, not production.
2. Pick large activations from expensive forward paths (QKV projections, MoE up-projections).
3. Enable offloading via jax.remat with memory_kind="pinned_host".
4. Measure three things: runtime GPU memory, host memory, and end-to-end step time.
5. Profile with Nsight Systems to confirm D2H and H2D copies actually overlap with compute and NCCL.

For teams also thinking about how cloud platforms handle workload scaling decisions, the Google Cloud Workbench VS Code integration is worth a look—it's a different layer of the same "where does the work actually run" question.

The Takeaway

Host offloading reframes a classic trade-off. Instead of paying with compute (rematerialization), you pay with bandwidth (streaming). On NVLink-C2C-class interconnects, that bandwidth is cheap enough to win.

The DeepSeek-V3 result is the headline: 908 TFLOPs/s/device at micro batch 8 / global batch 1024, a configuration that literally OOMs without offloading. That's not a 5% tuning gain—it's a new operating regime.

But don't cargo-cult it. Profile first. Offloading is a memory placement decision that must be validated with measurements, not assumptions. If your activations are small or your workload already overlaps well, you'll gain nothing. If your HBM is the wall between you and the next batch size, this might be the lever you've been looking for.

Next steps:

• Read the JAX host offloading tutorial (jax.remat, checkpoint policies, memory_kind="pinned_host").
• Try the NVIDIA JAX-Toolbox containers for a maintained starting point.
• Profile your current workload with Nsight Systems to see if transfer latency is actually exposed.

Source: NVIDIA Developer Blog — Reducing HBM Bottlenecks in JAX-Based LLM Training with Host Offloading