大型语言模型（LLM）的训练任务，经常在计算资源尚未耗尽时，便率先遭遇GPU显存的瓶颈。模型权重、梯度、优化器状态、通信缓冲区以及中间激活值，都在争夺GPU的高带宽显存（HBM）。随着模型规模、序列长度和批次大小的持续增长，HBM容量往往成为限制扩展的首要障碍。

本文将深入探讨开源Python库JAX中的一项关键技术——主机卸载（host offloading），它如何有效缓解HBM的压力。这一机制在NVIDIA Blackwell平台上尤为突出：在前向传播过程中，它会将选定的激活值传输至固定主机内存（pinned host memory），待反向传播需要时再流回GPU。与激活重计算（activation rematerialization）不同，主机卸载并非重新计算这些激活值，而是直接从主机内存中加载它们。

为何在NVIDIA Grace Blackwell系统上，主机卸载独具优势？

NVIDIA Grace CPU与NVIDIA Blackwell GPU之间通过NVLink-C2C互联，双向带宽高达900 GB/s，这使得固定主机内存成为存放选定激活值的实用中间缓存区。到了Vera CPU与Rubin GPU这一代，双向带宽更是翻倍至1.8 TB/s，并且具备连贯内存（coherent）带宽。

高带宽的CPU-GPU连接是主机卸载得以实现的基础，但仅有带宽还不够。若要真正提升性能，激活值的传输必须能够与GPU上的有效计算任务实现重叠。

MaxText工作负载上的性能表现

实验所用的MaxText，是一个基于JAX的LLM训练框架，底层依赖加速线性代数（XLA）编译器，用于在NVIDIA GPU上进行大规模训练。所有结果均在NVIDIA GB200 NVL72系统上，使用128块GPU测得。

评估采用了两个MaxText工作负载。Llama 3.1 405B是一个纯解码器（dense decoder-only）Transformer模型，用于研究在固定批次大小下对查询、键、值（QKV）激活值的定向卸载。DeepSeek-V3 671B则是一个稀疏混合专家（MoE）模型，配备多头潜在注意力（MLA），用于同时考察吞吐量与内存容量效应。

DeepSeek-V3 671B 卸载策略

DeepSeek-V3 671B包含61个解码层：前3层使用密集的多层感知机（MLP）块，后续层均为MoE块。

图1展示的是重复MoE解码层的激活卸载策略（该层占据了堆叠的绝大部分）。前3个密集MLP层未在图中显示，但其策略类似，也是将选中的up和down投影输出进行卸载。

在图1的MoE层中，该策略卸载了选中的MLA查询和键/值投影中间结果，以及选中的MoE up投影中间结果。这些激活值规模足够大，会直接影响更大批次配置的可行性。

训练吞吐量提升

MaxText以TFLOPs/s/device为单位报告吞吐量，计算公式为：每个训练步的模型FLOPs除以实测步时间。

DeepSeek-V3 671B

图2对比了DeepSeek-V3 671B在不同激活放置策略下的吞吐量。启用卸载、LHS（延迟隐藏调度器）和流水线传输后，工作负载达到了908.2 TFLOPs/s/device。这一结果比相同批次配置下使用激活重计算快了57%，比未启用LHS和流水线化的卸载快了67.7%。

与密集的Llama工作负载不同（LHS单独即可隐藏延迟），DeepSeek-V3的MoE和MLA层会产生巨大的激活值，因此启用流水线传输对总吞吐量具有明显的正面影响。

这些性能飞跃的背后，是NVIDIA一个根本性的优势：软硬件的紧密协同设计。在典型的普通集群上，将如此庞大的MoE激活层流式传输至主机内存，会拖垮整个训练流水线。但在Blackwell上，XLA的自定义调度标志与硬件协同工作，确保数据通过专用的拷贝流进行异步移动。

这使得NVIDIA平台能够解锁那些在缺乏这种编译器到互联紧密集成的架构上完全不可行的大型批次配置。

增大可行的批次大小

主机卸载还能改变哪些批次配置是可行的。表1展示了DeepSeek-V3 671B在不同激活放置策略下的结果。

第4行和第5行展示了容量对比。在设备上保存选中的激活值，可容纳微批次2和全局批次256；而优化后的主机卸载，则可容纳微批次8和全局批次1024。若既不使用卸载也不使用重计算，当尝试微批次8、全局批次1024时，设备将直接报OOM（内存不足）。

主机卸载通过将选中的激活值存储移出GPU内存，使微批次8、全局批次1024的配置变得可行。在这个DeepSeek配置中，卸载策略针对的是MLA、MoE和MLP块中的大型中间激活值。若将这些激活值留在设备上，将限制可行的批次配置；而卸载它们，则能为模型状态、通信缓冲区、运行时工作空间和活跃计算腾出更多HBM。

启用LHS和流水线传输后，卸载运行的GPU内存占用为165.2 GiB，相比之下，没有这些优化时为145.6 GiB。这一增长是因为在传输与计算重叠的同时，GPU内存中保留了更多的拷贝缓冲区和预取激活值。多占用的这部分HBM，是为了换取更好的重叠度和更高的吞吐量。

Llama 3.1 405B

Llama 3.1 405B的实验运行了10步，使用合成数据，批次大小2，序列长度8192，全分片数据并行（FSDP）设为128，激活值采用bfloat16，权重采用NVFP4 4-bit量化。

如表2所示，QKV激活卸载配合延迟隐藏调度器（LHS），将吞吐量从2,669提升至2,746 TFLOPs/s/device，相比不卸载的基线提升了2.9%。关闭LHS后，QKV卸载的吞吐量降至2,569 TFLOPs/s/device，这凸显了主机卸载依赖与其他GPU任务的有效重叠。

对于Llama 3.1 405B的这次运行，LHS单独就提供了最佳吞吐量：不用流水线时为2,746 TFLOPs/s/device，而用流水线时为2,718 TFLOPs/s/device。在此配置下，LHS已将大部分传输延迟隐藏到计算和通信之后，留给流水线去隐藏的暴露延迟已所剩无几。

70.9 GiB的主机内存值，是所有126层QKV激活值的总存储量，而非某一时刻节省的GPU内存量。在批次大小2、序列长度8192的情况下，一层bfloat16的QKV激活值大约需要576 MiB：查询占512 MiB，键和值各占32 MiB。

由于启用了层的扫描循环（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，这是一个定制化的容器，额外集成了MaxText对Transformer Engine MoE排列优化的支持。

优化后的卸载配置使用了以下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的技术贡献。