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

Large language model (LLM) training workloads increasingly run into GPU memory limits before compute is fully used. Model weights, gradients, optimizer states,...