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

DevFeed: [Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading](<https://devfeed.tech/articles/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading-6925.md>)

Original publisher: [Read original article](<https://developer.nvidia.com/blog/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading/>)

Author: Tanya Lenz

Published: 2026-07-10T18:17:40Z

Content type: article

Language: en

Sources: [NVIDIA Developer](<https://devfeed.tech/sources/nvidia-developer.md>), [NVIDIA Technical Blog](<https://devfeed.tech/sources/nvidia-technical-blog.md>)

Topics: [LLM Techniques](<https://devfeed.tech/topics/llm-techniques.md>), [Large Language Model](<https://devfeed.tech/topics/llm.md>), [GPU](<https://devfeed.tech/topics/gpu.md>), [Mixture of Experts (MoE)](<https://devfeed.tech/topics/mixture-of-experts-moe.md>), [cpu](<https://devfeed.tech/topics/cpu.md>), [NVLink](<https://devfeed.tech/topics/nvlink.md>), [systems](<https://devfeed.tech/topics/systems.md>), [Nvidia](<https://devfeed.tech/topics/nvidia.md>)

Tags: [agentic-ai-generative-ai](<https://devfeed.tech/tags/agentic-ai-generative-ai.md>), [blackwell](<https://devfeed.tech/tags/blackwell.md>), [compute](<https://devfeed.tech/tags/compute.md>), [cpu](<https://devfeed.tech/tags/cpu.md>), [developer-tools-techniques](<https://devfeed.tech/tags/developer-tools-techniques.md>), [featured](<https://devfeed.tech/tags/featured.md>), [gb200](<https://devfeed.tech/tags/gb200.md>), [gpu](<https://devfeed.tech/tags/gpu.md>), [grace-cpu](<https://devfeed.tech/tags/grace-cpu.md>), [large-language-model](<https://devfeed.tech/tags/large-language-model.md>), [llm](<https://devfeed.tech/tags/llm.md>), [llm-techniques](<https://devfeed.tech/tags/llm-techniques.md>), [mixture-of-experts-moe](<https://devfeed.tech/tags/mixture-of-experts-moe.md>), [nvidia](<https://devfeed.tech/tags/nvidia.md>), [nvl72](<https://devfeed.tech/tags/nvl72.md>), [nvlink](<https://devfeed.tech/tags/nvlink.md>), [simulation-modeling-design](<https://devfeed.tech/tags/simulation-modeling-design.md>), [training-ai-models](<https://devfeed.tech/tags/training-ai-models.md>), [vera-rubin](<https://devfeed.tech/tags/vera-rubin.md>), [work](<https://devfeed.tech/tags/work.md>)

## AI overview

This article explains how host offloading in JAX-based large language model training reduces GPU high-bandwidth memory pressure by moving selected activations to pinned host memory and streaming them back during the backward pass. It discusses activation-transfer overlap, NVIDIA Grace Blackwell and GB200 NVL72 systems, and experiments involving Llama 3.1 405B and DeepSeek-V3 671B.

## Source excerpt

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