---
格式版本: 2
标题: "Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading | NVIDIA Technical Blog"
原文链接: "https://developer.nvidia.com/blog/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading/"
发布日期: "2026-07-10"
发布时间校准状态: "found"
发布时间来源: "llm:strict_original_body"
发布时间证据: "div class=post-info: Jul 10, 2026"
发布时间校准原因: "日期位于标题下方的 post-info 区域，符合文章发布时间的典型位置，且无证据表明其为会议、活动或事件日期。"
发布时间校准置信度: "100"
发布时间候选数量: 8
发布时间严格候选数量: 2
发布时间原页读取状态: "原页面来自已抓取 HTML"
发布时间未找到原因: "候选日期无效或 LLM 未确认"
发布时间校准时间: "2026-07-20T12:07:12+08:00"
发现时间: "2026-07-20T11:40:00+08:00"
入库时间: "2026-07-20T04:15:44.303Z"
来源平台: "固定入口"
搜索渠道: "fixed_url"
搜索词: "https://developer.nvidia.com/blog/"
匹配关键词:
  - "GPU"
  - "HBM"
  - "NVL72"
  - "Nvlink"
  - "Vera Rubin"
相关厂家:
  - "NVIDIA"
  - "Meta"
相关专家:
  []
内容类型: "网页"
抓取工具: "AgentKey Scrape"
清洗工具: "AgentKey Markdown + LLM 正文裁剪"
原始附件:
  []
AI优质: "否"
AI打分: 68
AI分档: "召回候选"
AI质检状态: "不通过"
AI打分理由: "NVIDIA官方技术博客，来源权威且内容完整。文章聚焦JAX训练中的HBM瓶颈缓解与Host Offloading技术，虽提及GB200 NVL72、NVLink-C2C及Vera Rubin平台，但核心为软件/算法优化而非超节点硬件架构…"
AI质检模型: "qwen3.6-plus"
AI质检时间: "2026-07-20T12:15:44+08:00"
AI主题相关性: 12
AI来源权威性: 14
AI新颖性: 16
AI技术细节: 14
AI商业部署信号: 4
AI完整性: 8
图片摘要:
  - "✗ ./assets/img-011e5156.png | decorative | 装饰性配图，无具体技术信息"
  - "★ ./assets/img-7deec1c6.webp | diagram | DeepSeek-V3 671B MoE层激活卸载策略图，展示前向传播中Query/Key/Value/MoE投影激活卸载至Host内存，反向传播时回拷计算梯…"
  - "★ ./assets/img-854fd23a.webp | chart | DeepSeek-V3 671B吞吐量对比图，显示采用Host Offload结合LHS和流水线技术后，吞吐量达908.2 TFLOPs/s/device，显…"
  - "✗ ./assets/img-bab70b95.webp | ad | 会议广告"
  - "✗ ./assets/img-9d3a9443.webp | ad | 会议广告"
采集批次: "2026年7月20日11点39分58秒"
采集批次ID: "20260720-113958-114"
去重键: "https://developer.nvidia.com/blog/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, communication buffers, and intermediate activations all compete for GPU high-bandwidth memory (HBM). As model size, sequence length, and batch size grow, HBM capacity often becomes the primary scaling bottleneck.

This post explains how host offloading in the open source Python library [JAX](https://github.com/jax-ml/jax) reduces HBM pressure. This process, which is especially advantageous on [NVIDIA Blackwell](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/), moves selected activations to pinned host memory during the forward pass and streams them back when needed in the backward pass. Host offloading is an alternative to [activation rematerialization](https://docs.jax.dev/en/latest/gradient-checkpointing.html). Instead of recomputing selected activations, the training step reloads them from host memory.

## Why is host offloading advantageous on NVIDIA Grace Blackwell systems?

Host offloading is especially advantageous on [NVIDIA Grace Blackwell](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/) systems. The NVIDIA Grace CPU and NVIDIA Blackwell GPU are connected through NVLink-C2C with 900 GB/s of bidirectional bandwidth, making pinned host memory a practical staging area for selected activations. The Vera CPU and Rubin GPU further improve on this by doubling the bidirectional speed to 1.8 TB/s of coherent bandwidth.

High-bandwidth CPU-GPU connectivity helps make host offloading practical, but bandwidth alone is not enough. To improve performance, activation transfers must overlap with useful GPU work.

## Performance results on MaxText workloads

The experiments use [MaxText](https://github.com/AI-Hypercomputer/maxtext), a JAX LLM training framework that uses the [Accelerated Linear Algebra (XLA)](https://github.com/openxla/xla) compiler for training at scale on NVIDIA GPUs. All results were measured on [NVIDIA GB200 NVL72](https://www.nvidia.com/en-us/data-center/gb200-nvl72/) systems using 128 GPUs.

The evaluation uses two MaxText workloads. [Llama 3.1 405B](https://ai.meta.com/blog/meta-llama-3-1/) is a dense decoder-only transformer model used to study targeted query, key, and value (QKV) activation offloading at a fixed batch size. [DeepSeek-V3 671B](https://arxiv.org/abs/2412.19437) is a sparse mixture-of-experts (MoE) model with multihead latent attention (MLA), used to study both throughput and memory-capacity effects.

### DeepSeek-V3 671B offloading policy

DeepSeek-V3 671B has 61 decoder layers: the first three use dense multilayer perceptron (MLP) blocks, while the remaining layers use MoE blocks.

Figure 1 shows the activation offloading policy for the repeated MoE decoder layer, which dominates the stack. The first three dense MLP layers are not shown; they use a similar policy, with selected up and down projection outputs offloaded.

In the MoE layer shown in Figure 1, the policy offloads selected MLA query and key/value projection intermediates and selected MoE up projection intermediates. These activations are large enough to impact whether larger batch configurations fit.

![Diagram showing the DeepSeek-V3 671B repeated MoE decoder layer activation offloading policy. The top row shows the forward pass, where Query Projection a, Key and Value Projection a, and MoE Up Projection activations are offloaded to host memory. The bottom row shows the backward pass, where those activations are copied back to compute gradients.](./assets/img-7deec1c6.webp)

Figure 1. DeepSeek-V3 671B repeated MoE decoder layer activation offloading policy

### Training throughput improvements

MaxText reports throughput as TFLOPs/s/device, computed from model FLOPs per training step divided by measured step time.

#### DeepSeek-V3 671B

Figure 2 compares DeepSeek-V3 671B throughput across activation placement policies. With offloading, LHS, and pipelined transfers enabled, the workload reached 908.2 TFLOPs/s/device. This was 57% faster than activation rematerialization at the same batch configuration and 67.7% faster than offloading without LHS or pipelining.

Unlike the dense Llama workload where LHS alone was sufficient to hide latency, the massive activation footprint of DeepSeek-V3 MoE and MLA layers means that enabling pipelined transfers provides a distinct, positive impact on total throughput.

These performance leaps highlight a fundamental NVIDIA advantage: the tight co-design of software and hardware. In a typical commodity cluster, streaming massive MoE activation layers to host memory would stall the entire training pipeline. On Blackwell, XLA custom scheduling flags work hand-in-hand with the hardware to ensure data is moved asynchronously over dedicated copy streams.

This enables NVIDIA platforms to unlock massive batch configurations that remain completely out of reach for architectures lacking this tight compiler-to-interconnect integration.

![Bar chart comparing DeepSeek-V3 671B throughput for three configurations. Host offload with LHS and pipelined offloading reaches 908.2 TFLOPs/s/device, compared with 578.3 TFLOPs/s/device for no offload with LHS and activation recompute, and 541.6 TFLOPs/s/device for host offload without LHS or pipelined offloading.](./assets/img-854fd23a.webp)

Figure 2. DeepSeek-V3 671B throughput on NVIDIA GB200 across activation placement configurations

### Increasing feasible batch size

Host offloading can also change which batch configurations are feasible. Table 1 shows DeepSeek-V3 671B results across activation placement policies.

| **No** | **Config** | **Micro** **batch** | **Global** **batch** | **Throughput** **(TFLOPs/s/device)** | **GPU peak** **(GiB)** | **Host memory** **(GiB)** |
| --- | --- | --- | --- | --- | --- | --- |
| **1** | **Host offload, LHS, pipelined offloading** | **8** | **1024** | **908.2** | **165.2** | **145.1** |
| 2 | No offload, LHS, activation rematerialization | 8 | 1024 | 578.3 | 151.3 | 0.0 |
| 3 | Host offload, no LHS, no pipelined offloading | 8 | 1024 | 541.6 | 145.6 | 145.1 |
| 4 | No offload, LHS, save on device | 2 | 256 | 425.3 | 113.3 | 0.0 |
| 5 | No offload, LHS, save on device | 8 | 1024 | – | OOM | 0.0 |

*Table 1. DeepSeek-V3 671B comparison across activation placement policies*

Rows 4 and 5 show the capacity comparison. Saving selected activations on device fit micro batch 2 and global batch 256, while optimized host offloading fit micro batch 8 and global batch 1024. Without using either offloading or rematerialization, the device will hit an Out-of-Memory (OOM) error when attempting to use a micro batch size of 8 and global batch size of 1024.

Host offloading made the micro batch 8, global batch 1024 configuration feasible by moving selected activation storage out of GPU memory. In this DeepSeek configuration, the offload policy targets large intermediate activations from MLA, MoE, and MLP blocks. Keeping those activations on device limits the feasible batch configuration, while offloading them leaves more HBM available for model state, communication buffers, runtime workspaces, and active computation.

With LHS and pipelined transfers enabled, the offload run uses 165.2 GiB of GPU memory, compared with 145.6 GiB without those optimizations. The increase comes from keeping more copy buffers and prefetched activations in GPU memory while transfers overlap with computation. This extra HBM use trades some memory capacity for better overlap and higher throughput.

#### Llama 3.1 405B

The Llama 3.1 405B experiment ran 10 steps on synthetic data with batch size 2, sequence length 8,192, fully sharded data parallelism (FSDP) set to 128, and bfloat16 activations with [NVFP4](https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/) 4-bit weight quantization.

As shown in Table 2, QKV activation offloading with Latency Hiding Scheduler (LHS) improved throughput from 2,669 to 2,746 TFLOPs/s/device, a 2.9% increase over the baseline without offloading. Disabling LHS reduced QKV offload throughput to 2,569 TFLOPs/s/device, highlighting that host offloading depends on effective overlap with other GPU work.

| **Config** | **LHS** | **Pipelined offloading** | **Throughput** **(TFLOPs/s/device)** | **GPU peak memory** **(GiB)** | **Host memory** **(GiB)** |
| --- | --- | --- | --- | --- | --- |
| No offload | ON | OFF | 2,669 | 149.6 | 0 |
| QKV offload | ON | OFF | 2,746 | 149.9 | 70.9 |
| QKV offload | OFF | OFF | 2,569 | 139.7 | 70.9 |
| QKV offload | ON | ON | 2,718 | 151.0 | 70.9 |

*Table 2. Llama 3.1 405B throughput and memory across activation placement configurations*

For this Llama 3.1 405B run, LHS alone provided the best throughput: 2,746 TFLOPs/s/device without pipelining compared with 2,718 TFLOPs/s/device with pipelining. In this configuration, LHS already hides most transfer latency behind compute and communication, leaving little exposed latency for pipelining to hide.

The 70.9 GiB host memory value is the total QKV activation storage across all 126 layers, not the amount of GPU memory saved at a single moment. At batch size 2 and sequence length 8,192, one layer’s bfloat16 QKV activations require about 576 MiB: 512 MiB for query and 32 MiB each for key and value.

With the scan loop enabled for layers (scan\_layers=True), the backward pass processes one layer at a time, so QKV activations of only one layer are needed on the GPU at once. In this workload, QKV offloading is mainly a performance optimization: it replaces backward pass QKV rematerialization with transfers that can overlap with compute and communication. GPU peak memory remains dominated by model state, communication buffers, and runtime workspaces.

Llama 3.1 405B provides a dense-model fixed-batch example. The gain is smaller than DeepSeek V3 671B, but it shows the same mechanism: targeted QKV offloading replaces backward-pass rematerialization with transfers that overlap with compute and communication.

## When is host offloading most useful?

Host offloading is most useful when GPU memory limits model size, sequence length, or batch size, and when selected tensors are large enough to reduce HBM pressure. It is especially useful when offloading can replace expensive activation rematerialization or make a larger batch configuration feasible.

Performance depends on overlap. Host offloading works best when the workload has enough compute, communication, or other independent work to hide transfer latency. On NVIDIA GPUs, XLA helps create this overlap by using dedicated copy streams, scheduling transfers with LHS, and enabling pipelined host offloading.

This pipeline is particularly advantageous on platforms like the NVIDIA Blackwell GB200 and NVIDIA Blackwell Ultra GB300, which leverage dedicated NVLink-C2C interconnect to bypass traditional PCIe bottlenecks entirely. With the [NVIDIA Vera Rubin platform](https://www.nvidia.com/en-us/data-center/technologies/rubin/) delivering even higher interconnect performance, host offloading provides a predictable architectural lever for decoupling training throughput from strict physical memory limits.

Host offloading is less likely to help when tensors are small, when little independent work is available to overlap transfers, or when the workload is bottlenecked somewhere other than memory. Runtime memory should still be validated with real runs because static estimates may not include [NVIDIA Collective Communications Library (NCCL)](https://developer.nvidia.com/nccl) communication scratch space, [NVIDIA cuDNN](https://developer.nvidia.com/cudnn) attention workspace, and framework-managed buffers.

## How to get started with host offloading

Start with a small representative JAX training run. Choose large activations from expensive forward paths, enable offloading, and measure runtime GPU memory, host memory use, and end-to-end step time.

For open source JAX APIs, see the [JAX host offloading tutorial](https://jax.readthedocs.io/en/latest/notebooks/host-offloading.html). It covers activation offloading with jax.remat, checkpoint policies, and memory\_kind=”pinned\_host”, as well as parameter and optimizer state offloading with jax.device\_put().

The MaxText experiments in this post used the [NGC JAX container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax). The DeepSeek-V3 671B runs used [ghcr.io/nvidia/jax:deepseek\_v3\_maxtext](http://ghcr.io/nvidia/jax:deepseek_v3_maxtext), a customized container with additional MaxText integration for [Transformer Engine](https://github.com/NVIDIA/TransformerEngine) MoE permutation optimizations.

The optimized offload configurations used the following XLA flags:

```
--xla_gpu_enable_latency_hiding_scheduler=true
--xla_gpu_enable_pipelined_host_offloading=true
--xla_gpu_experimental_parallel_async_compute_limit=8
```

The last setting increases the amount of asynchronous work that can be in flight, giving LHS more room to overlap activation copies and NCCL collectives. Use profiling tools such as [NVIDIA Nsight Systems](https://developer.nvidia.com/nsight-systems) to confirm that device-to-host and host-to-device copies overlap with compute and NCCL communication.

## Learn more

Host offloading is a strong fit when HBM capacity limits batch size, context length, or model scale, and when selected activations are large enough to reduce GPU memory pressure. It is most effective when it replaces expensive activation rematerialization or makes a larger batch configuration feasible.

Treat host offloading as a memory placement choice that should be validated with measurements. Choose large activations, enable overlap with LHS and pipelined host offloading, and profile both runtime memory and step time.

While the results presented in this post focus on Llama 3.1 405B and DeepSeek-V3 671B, the same approach can help other JAX workloads when selected activations are large enough to matter and transfer cost can overlap with useful computation.

To run JAX on NVIDIA GPUs, [NVIDIA JAX-Toolbox](https://github.com/NVIDIA/JAX-Toolbox) provides maintained containers, documentation, and optimized JAX and MaxText examples.

### Acknowledgments

*We would like to thank Jaroslav Sevcik, Sevin Varoglu, Tj Xu, Haixin Liu, Abhinav Goel, Md Fahim Faysal Khan, Stefano Bosisio, and Jinxin Yang for their technical contributions.*

![Diagram showing the DeepSeek-V3 671B repeated MoE decoder layer activation offloading policy. The top row shows the forward pass, where Query Projection a, Key and Value Projection a, and MoE Up Projection activations are offloaded to host memory. The bottom row shows the backward pass, where those activations are copied back to compute gradients.](./assets/img-7deec1c6.webp)_Figure 1. DeepSeek-V3 671B repeated MoE decoder layer activation offloading policy_

![Bar chart comparing DeepSeek-V3 671B throughput for three configurations. Host offload with LHS and pipelined offloading reaches 908.2 TFLOPs/s/device, compared with 578.3 TFLOPs/s/device for no offload with LHS and activation recompute, and 541.6 TFLOPs/s/device for host offload without LHS or pipelined offloading.](./assets/img-854fd23a.webp)_Figure 2. DeepSeek-V3 671B throughput on NVIDIA GB200 across activation placement configurations_

- ![图片](./assets/img-bab70b95.webp)

- ![图片](./assets/img-9d3a9443.webp)