JAX Host Offloading: A Practical Guide to Breaking Through LLM Training Memory Bottlenecks

The Memory Wall: The Hidden Ceiling on LLM Training

As large language model (LLM) scales continue to expand, training workloads are increasingly hitting the memory ceiling before GPU compute is even fully utilized. Model weights, gradients, optimizer states, and activations relentlessly consume high-bandwidth memory (HBM), making single-device memory capacity the core bottleneck constraining both model scale and batch size.

This reflects a fundamental structural tension: modern GPUs' compute capabilities are growing far faster than their memory capacity. GPU performance gains are driven primarily by increases in transistor density and the expansion of parallel compute units — roughly doubling every two years or faster — while HBM capacity is constrained by package area, stacking layers, and yield rates, growing at a far more modest pace. Take NVIDIA's flagship GPUs as an example: the A100 ships with 80 GB of HBM2e, the H100 steps up to 80 GB of HBM3, and only the H200 jumps to 141 GB — while compute throughput (FLOPS) increased by a far greater margin over the same period. In other words, when we try to run larger models on GPU, the problem isn't usually insufficient compute — it's that the memory simply can't hold it all.

Mixed-precision training puts numbers to this tension. A model with 7 billion parameters (7B) requires roughly 14 GB just for the model weights (FP16); the Adam optimizer's first and second moments (FP32) consume an additional ~56 GB; and gradients (FP16/FP32 mixed) require another 14–28 GB. Combined, this already exceeds 80 GB before accounting for the activations generated during batch training. The industry commonly uses "parameter count × 16–20 bytes" as a rough estimate for full-parameter Adam training memory usage — which is precisely why 70B-scale models, even in multi-GPU settings, must rely on techniques like offloading, pipeline parallelism, or ZeRO (Zero Redundancy Optimizer) to break through the memory wall. JAX-based host offloading is a systematic solution to exactly this problem.

What Is Host Offloading and How Does It Work

The core idea behind host offloading is straightforward: since GPU HBM capacity is limited, temporarily migrate tensors that aren't actively involved in computation to the larger, lower-cost host memory (CPU DRAM), then transfer them back to the GPU when needed.

Which Data Types Are Good Candidates for Offloading

During LLM training, the following data categories are prime candidates for host offloading:

• Optimizer states: Optimizers like Adam maintain first and second moments for every parameter, consuming several times the memory of the model weights themselves. This data is only accessed during parameter updates and can safely reside in host memory the rest of the time.
• Activations: Intermediate activations produced during the forward pass but not consumed until the backward pass are major contributors to memory usage. Offloading them to the host frees up space for deeper networks or larger batch sizes.
• Gradient accumulation intermediates: In gradient accumulation scenarios, some gradient data can similarly be temporarily offloaded.

How Host Offloading Differs from Activation Recomputation

These are fundamentally different technical approaches. Activation recomputation saves memory by discarding activations and recomputing them during the backward pass — at the cost of additional compute. Host offloading trades PCIe/NVLink bandwidth for memory capacity — at the cost of data transfer latency. The two are complementary and can be combined to find the optimal balance between compute overhead and transfer bandwidth.

It's worth noting that host offloading is conceptually related to Microsoft DeepSpeed's ZeRO (Zero Redundancy Optimizer) family, though the approaches differ. ZeRO-1/2/3 reduce peak per-device memory by sharding optimizer states, gradients, and parameters across multiple GPUs; ZeRO-Infinity goes further by offloading sharded data to host memory or even NVMe SSDs. JAX's host offloading shares a similar philosophy, but leverages the XLA compiler's static analysis capabilities to perform more fine-grained scheduling at compile time, rather than relying on runtime dynamic decisions.

Why JAX Is Naturally Well-Suited for Host Offloading

JAX is fundamentally a Python library for numerical computing. At its core, it transforms Python functions into XLA (Accelerated Linear Algebra) computation graphs via jit (just-in-time compilation). XLA is a domain-specific compiler developed internally at Google that performs global optimization across the entire computation graph — including operator fusion, memory reuse planning, and asynchronous scheduling. Unlike PyTorch's dynamic graph execution, JAX's functional programming paradigm requires functions to be pure (free of side effects), which allows the compiler to safely reorder, fuse, and pre-schedule operations.

This property provides a critical foundation for host offloading: when generating an execution plan, the compiler knows the complete lifetime of every tensor, and can insert offload and prefetch operations into the gaps between computations — achieving true pipeline overlap rather than bolting on asynchronous calls after the fact. Through APIs like jax.device_put, developers can explicitly control where tensors reside; combined with the compiler's global analysis of the computation graph, offload and prefetch operations can be planned at compile time, enabling pipelined overlap of data transfers and computation.

In the ideal case, while the GPU is executing the forward computation for one layer, the data needed for the next layer is already being quietly prefetched in the background. This overlap of compute and communication is the key to whether host offloading actually delivers real benefits.

Bandwidth Is the Core Trade-Off

Host offloading is not without cost. Frequent data movement between device memory and host memory consumes PCIe or NVLink bandwidth. A poorly designed strategy can turn transfer latency into a new bottleneck.

Understanding this constraint requires confronting the raw bandwidth gap head-on: PCIe 4.0 x16 provides roughly 64 GB/s of bidirectional bandwidth; PCIe 5.0 doubles that to approximately 128 GB/s. Meanwhile, GPU HBM bandwidth sits at 3–4 TB/s — a difference of roughly 30–60×. This means offloading operations must be sufficiently overlapped with computation; otherwise, the transfers themselves become the new bottleneck.

In practice, the following principles apply:

1. Prioritize offloading infrequently accessed data: Optimizer states are accessed only once per step, making them the ideal offloading target.
2. Use asynchronous transfers to hide latency: Overlap data movement with computation to avoid leaving the GPU idle.
3. Deploy on hardware with high-bandwidth host-device interconnects: Next-generation platforms like the Grace Hopper architecture with NVLink-C2C offer dramatically better host-device interconnect bandwidth, significantly reducing the cost of offloading.

The NVIDIA Grace Hopper Superchip tightly couples the CPU and GPU via NVLink-C2C (Chip-to-Chip), a high-bandwidth coherent interconnect delivering up to 900 GB/s of bidirectional bandwidth — nearly an order of magnitude higher than PCIe — while supporting a unified memory address space between CPU and GPU. This eliminates the need for explicit data copies; tensors in host memory are nearly transparently accessible to the GPU, fundamentally changing the cost structure of host offloading. On such platforms, host memory can be treated almost as a slower extension of GPU memory.

When Host Offloading Is the Right Choice

The greatest practical value of host offloading is this: you can train larger models or use larger batch sizes without adding more GPUs. For workloads that are memory-bound but compute-rich, it's an extremely cost-effective optimization.

It's particularly well-suited for:

• Training large models on a single machine, constrained by per-device memory capacity;
• Scenarios where optimizer states consume excessive memory (e.g., full-parameter fine-tuning of large models);
• Hardware equipped with high-bandwidth host-device interconnects.

Conversely, for compute-intensive workloads where bandwidth is already the limiting factor, blindly introducing host offloading may do more harm than good. Technology selection should be guided by model architecture, hardware configuration, and profiling data — balancing memory capacity, compute throughput, and bandwidth to find the optimal operating point.

Conclusion

As the arms race in model scale continues to escalate, the memory bottleneck will only become more acute. Host offloading offers a pragmatic path forward: rather than simply throwing more expensive GPUs at the problem, it cleverly leverages the host memory resources already present in the system. With JAX and XLA's compilation optimization capabilities — particularly the global computation graph analysis enabled by functional purity — combined with the high-bandwidth interconnects of next-generation hardware, host offloading is becoming an indispensable tool in the large-scale LLM training toolkit. For teams striving to balance training efficiency and cost, developing a deep understanding of this technique and applying it thoughtfully is a key to breaking through the memory wall.

Key Takeaways