在充分利用计算能力之前，大语言模型 (LLM) 训练工作负载越来越多地遇到 GPU 显存限制。模型权重、梯度、优化器状态、通信缓冲区和中间激活函数都在争夺 GPU 高带宽内存 (HBM) 。随着模型大小、序列长度和批量大小的增长，HBM 容量通常成为主要的扩展瓶颈。
本文将介绍开源 Python 库中的主机卸载如何减轻 HBM 压力。此过程在 NVIDIA Blackwell 上尤为有利，它在前向传递期间将选定的激活函数移动到固定主机内存，并在后向传递需要时将其流回。主机卸载是 激活重物质化 的替代方案。训练步骤不是重新计算选定的激活函数，而是从主机内存中重新加载它们，特别是在开源 Python 库 JAX 中。
为什么主机卸载在 NVIDIA Grace Blackwell 系统上具有优势？
主机卸载在 NVIDIA Grace Blackwell 系统上尤其有利。NVIDIA Grace CPU 和 NVIDIA Blackwell GPU 通过 NVLink-C2C 连接，双向带宽高达 900 GB/s，使固定主机内存成为选定激活的实用暂存区。Vera CPU 和 Rubin GPU 通过将双向速度翻倍至 1.8 TB/s 的一致性带宽，进一步改善了这一点。
高带宽 CPU-GPU 连接有助于实现主机卸载，但仅凭带宽是不够的。为提高性能，激活传输必须与有用的 GPU 工作重叠。
MaxText 工作负载的性能结果
这些实验使用 MaxText，一个 JAX LLM 训练框架，使用加速线性代数 (XLA) 编译器在 NVIDIA GPU 上进行大规模训练。所有结果均在使用 128 个 GPU 的 NVIDIA GB200 NVL72 系统上测量得出。
评估使用两个 MaxText 工作负载。 Llama 3.1 405B 是一个仅使用密集解码器的 Transformer 模型，用于研究固定批量大小的定向查询、键和值 (QKV) 激活卸载。 DeepSeek-V3 671B 是具有多头潜在注意力 (MLA) 的稀疏混合专家 (MoE) 模型，用于研究吞吐量和内存容量效应。
DeepSeek-V3 671B 卸载策略
DeepSeek-V3 671B 拥有 61 个解码器层：前三层使用密集多层感知器 (MLP) 块，其余层使用 MoE 块。
图 1 显示了主导堆栈的重复 MoE 解码器层的激活卸载策略。系统未显示前三层密集 MLP 层；它们采用类似策略，卸载了选定的上下投影输出。
在图 1 所示的 MoE 层中，策略卸载了选定的 MLA 查询和键值投影中间体，以及选定的 MoE 向上投影中间体。这些激活函数足够大，足以影响较大的批量配置是否适合。
图 1. DeepSeek-V3 671B 重复 MoE 解码器层激活卸载策略训练吞吐量提升
MaxText 将吞吐量报告为 TFLOPS/s/ 设备，计算方法为每个训练步骤的模型 FLOPS 除以测量的步长时间。
DeepSeek-V3 671B
图 2 比较了 DeepSeek-V3 671B 在激活放置策略中的吞吐量。启用卸载、LHS 和流水线传输后，工作负载达到 908.2 TFLOPS/s/ 设备。这在相同批量配置下比激活重物质化快 57%，在不使用 LHS 或工作流的情况下比卸载快 67.7%。
在密集的 Llama 工作负载中，仅 LHS 足以隐藏延迟，而 DeepSeek-V3 MoE 和 MLA 层的大量激活占用意味着，启用流水线传输会对总吞吐量产生明显的积极影响。
这些性能飞跃凸显了 NVIDIA 的一项基本优势：软件和硬件的紧密协同设计。在典型的通用集群中，将大量 MoE 激活层串流到主机内存会使整个训练工作流陷入停顿。在 Blackwell 上，XLA 自定义调度标志与硬件协同工作，可确保数据在专用的复制流中异步移动。
这使得 NVIDIA 平台能够解锁大规模批量配置，而对于缺乏这种紧密的编译器到互连集成的架构来说，这些配置仍然完全遥不可及。
图 2. NVIDIA GB200 上的 DeepSeek-V3 671B 吞吐量，跨激活放置配置增加可行的批量大小
主机卸载还可以更改可行的批量配置。表 1 显示了 DeepSeek-V3 671B 在激活放置策略中的结果。
| 否 | 配置 | 微型 批量 |
全球 批量 |
吞吐量 ( TFLOPS/ 秒/ 设备) |
GPU 峰值 ( GiB) |
主机内存 ( GiB) |
| 1 | 主机卸载、LHS、流水线卸载 | 8 | 1024 | 908.2 | 165.2 | 145.1 |
| 2 | 无卸载、LHS、激活重物质化 | 8 | 1024 | 578.3 | 151.3 | 0.0 |
| 3 | 主机卸载、无 LHS、无流水线卸载 | 8 | 1024 | 541.6 | 145.6 | 145.1 |
| 4 | 无卸载，LHS，保存在设备上 | 2 | 256 | 425.3 | 113.3 | 0.0 |
| 5 | 无卸载，LHS，保存在设备上 | 8 | 1024 | – | OOM | 0.0 |
第 4 行和第 5 行显示容量比较。将选定的激活函数保存在设备拟合微批量 2 和全局批量 256 上，同时优化主机卸载拟合微批量 8 和全局批量 1024。如果不使用卸载或重物质化，设备在尝试使用微批量大小为 8、全局批量大小为 1024 时将遇到内存不足 (OOM) 错误。
主机卸载通过将选定的激活存储移出 GPU 显存，使微批量 8 (全局批量 1024) 配置成为可能。在此 DeepSeek 配置中，卸载策略针对的是来自 MLA、MoE 和 MLP 块的大型中间激活函数。在设备上保留这些激活函数会限制可行的批量配置，同时将其卸载，从而使更多 HBM 可用于模型状态、通信缓冲区、运行时工作空间和主动计算。
启用 LHS 和工作流传输后，卸载运行将使用 165.2 GiB 的 GPU 显存，而在不进行这些优化的情况下，这一比例为 145.6 GiB。增加的原因是在 GPU 显存中保留更多的复制缓冲区和预取激活函数，同时传输与计算重叠。这种额外的 HBM 使用会交换一些内存容量，以实现更好的重叠和更高的吞吐量。
Llama 3.1 405B
Llama 3.1 405B 实验在合成数据上运行了 10 个步骤，其中批量大小为 2，序列长度为 8192，全分片数据并行 (FSDP) 设置为 128，并使用 NVFP4 4 位权重量化进行 bfloat16 激活。
如表 2 所示，使用延迟隐藏调度器 (LHS) 的 QKV 激活卸载将吞吐量从 2669 TFLOPS/s/ 设备提高到 2746 TFLOPS/s/ 设备，在不卸载的情况下比基准增加了 2.9%。禁用 LHS 可将 QKV 卸载吞吐量降低至 2569 TFLOPS/s/ 设备，这凸显出主机卸载依赖于与其他 GPU 工作的有效重叠。
| 配置 | LHS | 流水线卸载 | 吞吐量 ( TFLOPS/ 秒/ 设备) |
GPU 峰值显存 ( GiB) |
主机内存 ( GiB) |
| 无卸载 | 开启 | 关闭 | 2,669 | 149.6 | 0 |
| QKV 卸载 | 开启 | 关闭 | 2,746 | 149.9 | 70.9 |
| QKV 卸载 | 关闭 | 关闭 | 2,569 | 139.7 | 70.9 |
| QKV 卸载 | 开启 | 开启 | 2,718 | 151.0 | 70.9 |
对于本次 Llama 3.1 405B 的运行，仅 LHS 就提供了最佳吞吐量：2746 TFLOPS/s/ 设备，无需管线，而使用管线的设备为 2718 TFLOPS/s/ 设备。在此配置中，LHS 已将大部分传输延迟隐藏在计算和通信背后，将很少的暴露延迟留给流水线隐藏。
70.9 GiB 主机显存值是所有 126 层中 QKV 激活存储的总量，而不是单个时刻节省的 GPU 显存量。在批量大小为 2 和序列长度为 8192 时，单层 Bfloat16 QKV 激活大约需要 576MiB：512MiB 用于查询，32MiB 用于键和值。
为层启用扫描循环 (scan_layers = True) 后，反向传递一次处理一层，因此 GPU 上一次只需激活一层的 QKV。在此工作负载中，QKV 卸载主要是一种性能优化：它使用可与计算和通信重叠的传输来取代反向传递 QKV 重物质化。GPU 峰值显存仍由模型状态、通信缓冲区和运行时工作空间主导。
Llama 3.1 405B 提供了一个密集模型固定批量示例。该增益小于 DeepSeek V3 671B，但它展示了相同的机制：定向 QKV 卸载用与计算和通信重叠的传输取代了反向传递再物质化。
主机卸载何时最有用？
当 GPU 显存限制模型大小、序列长度或批量大小，以及选定的张量足够大以减轻 HBM 压力时，主机卸载最有用。当卸载可以取代昂贵的激活重材质化或使更大的批量配置可行时，它尤其有用。
性能取决于重叠。当工作负载具有足够的计算、通信或其他独立工作来隐藏传输延迟时，主机卸载效果最佳。在 NVIDIA GPU 上，XLA 通过使用专用的复制流、使用 LHS 调度传输以及启用流水线主机卸载来帮助创建这种重叠。
在 NVIDIA Blackwell GB200 和 NVIDIA Blackwell Ultra GB300 等利用专用 NVLink-C2C 互连技术完全绕过传统 PCIe 瓶颈的平台上，该流程尤其具有优势。由于 NVIDIA Vera Rubin 平台 可提供更高的互连性能，因此主机卸载提供了一个可预测的架构杠杆，用于将训练吞吐量与严格的物理内存限制解。
当张量较小、几乎没有可用于重叠传输的独立工作，或者工作负载在内存以外的某个地方遇到瓶颈时，主机卸载不太可能发挥作用。运行时内存仍应通过实际运行进行验证，因为静态估计值可能不包括 NVIDIA 集合通信库 (NCCL) 通信暂存空间、NVIDIA cuDNN 注意力工作空间和框架管理缓冲区。
如何开始主机卸载
从小型代表性 JAX 训练开始。从昂贵的前向路径中选择大型激活函数，启用卸载，并测量运行时 GPU 内存、主机内存使用情况和端到端步骤时间。
有关开源 JAX API，请参阅 JAX 主机卸载教程。它涵盖使用 jax.remat、检查点策略和 memory_kind=”pinned_host” 进行激活卸载，以及使用 jax.device_put() 进行参数和优化器状态卸载。
本文中的 MaxText 实验使用了 NGC JAX 容器。DeepSeek-V3 671B 运行使用 ghcr.io/nvidia/jax:deepseek_v3_maxtext，这是一个自定义容器，具有针对 Transformer 引擎 MoE 排列优化的额外 MaxText 集成。
优化的卸载配置使用以下 XLA 标志：
--xla_gpu_enable_latency_hiding_scheduler=true
--xla_gpu_enable_pipelined_host_offloading=true
--xla_gpu_experimental_parallel_async_compute_limit=8
最后一个设置增加了可以运行的异步工作量，为 LHS 提供了更多的空间来重叠激活副本和 NCCL 群集。使用 NVIDIA Nsight Systems 等分析工具确认设备到主机和主机到设备的复制与计算和 NCCL 通信重叠。
了解详情
当 HBM 容量限制批量大小、上下文长度或模型规模时，以及当选定的激活值足够大以降低 GPU 显存压力时，主机卸载是很强的拟合。当它取代昂贵的激活重物质化或使更大的批量配置成为可能时，效果最为显著。
将主机卸载视为应通过测量验证的内存放置选项。选择大型激活函数，通过 LHS 和流水线主机卸载实现重叠，并分析运行时内存和步骤时间。
虽然本文中介绍的结果重点介绍了 Llama 3.1 405B 和 DeepSeek-V3 671B，但当选定的激活函数足够大，并且传输成本可能与有用的计算重叠时，相同的方法可以帮助其他 JAX 工作负载。
为了在 NVIDIA GPU 上运行 JAX，NVIDIA JAX-Toolbox 提供了经过维护的容器、文档以及优化的 JAX 和 MaxText 示例。
致谢
我们想感谢 Jaroslav Sevcik、Sevin Varoglu、Tj Xu、Haixin Liu、Abhinav Goel、Md Fahim Faysal Khan、Stefano Bosisio 和 Jinxin Yang 所做的技术贡献。
