在强化学习（RL）后训练过程中，训练轨迹是通过使用当前策略模型进行推理实时生成的。出于训练效率考虑，通常会在分离式架构中将推理和训练分别运行在不同的GPU资源上。尽管策略模型只在训练GPU上更新，系统仍需定期用训练GPU生成的最新参数刷新推理GPU上的权重。这种在训练和推理GPU之间传输模型权重的过程被称为"重载"（refit）。
重载面临两个关键挑战。首先，推理和训练的计算需求不同，因此权重在推理和训练实例中的分布方式也不同，二者采用不同的并行映射策略（数据并行、上下文并行、张量并行、专家并行和流水线并行）。
其次，由于重载必须更新推理GPU上的所有模型权重，通信负载量级与模型权重总大小相当。对于当前最先进的模型，这一负载已经非常庞大，并且预计将持续增长。综合这两个特点，重载需要在多种分布式源端和目标端布局之间实现高性能通信。
本文将介绍AI训练框架如何利用Reshard集合通信——一种主要为加速RL权重重载而设计的新型集合通信原语——来应对这些重载挑战。NCCL M2N Reshard作为Reshard集合通信的开源参考实现，可通过NVIDIA/nccl-extensions GitHub仓库获取。
随着模型规模扩大，当前权重重载方法面临的挑战
现有的重载实现通常采用"收集-广播"方法来处理训练和推理布局之间的差异。
对于每个模型参数，典型的"收集-广播"权重重载流程如下：
训练GPU通过全收集（all-gather）操作收集各自的本地分片，重建完整参数。
将完整的参数张量广播到推理GPU。
每个推理GPU保留其本地布局所需的分片。
使用完整张量作为中间表示，可以简化不同布局之间的权重传输。但这样一来，即使推理GPU只需要一个分片，也会收到完整的参数，这在收集和广播阶段都造成了冗余通信，效率欠佳。
这种通信量会随模型规模的扩大而增长，而模型规模本身仍在持续扩张。例如，Kimi K3拥有2.8万亿参数，而Qwen 3.8 2.4T A95B拥有2.4万亿参数。假设使用BF16权重（每个参数2字节），二者的权重大小分别约为5.6 TB和4.8 TB。
综合考虑全收集和广播两个阶段，估计逻辑负载约为模型权重大小的两倍（见表1）。
对于当今最大的开放权重模型，每次重载的通信负载已达到约10 TB，而这一过程需要在每个训练步骤中执行。随着大语言模型规模持续扩大，这种大负载带来的开销将变得愈发显著。因此，需要一种更精细的方法来执行权重重载，在妥善处理多样化权重布局的同时，降低实际传输的负载大小。
将重载重新表述为DTensor到DTensor的传输问题
为克服这些局限，需要从一个基本原则出发：重载应仅传输理论上所需的最小量权重数据。然而，实现这一目标存在若干挑战。训练和推理实例可能使用不同的并行配置和GPU数量，从而产生不同的分片边界。因此，目标端的一个分片可能需要从多个源端分片的部分内容组装而成，反之亦然。
复制机制进一步使映射关系复杂化，因为相同的权重可能存在于多个源GPU上，或被多个目标GPU所需要。要确定哪些源GPU应向每个目标发送哪些张量区域，并在所有可能的配置中实现高性能，这并非易事。
解决这一挑战的关键是为这种独特的通信模式找到合适的抽象。研究团队借助了PyTorch现有的分布式张量（DTensor）抽象，该抽象旨在表示在不同形式的数据并行和模型并行下，跨GPU分片或复制的张量。关键的洞见在于：无论两个实例使用何种并行配置，权重重载都可以重新表述为将DTensor从训练布局重新分片（reshard）为推理布局的过程。
DTensor布局由张量的全局形状、指定GPU逻辑排列方式的设备网格（device mesh），以及指定张量如何在该网格上分片或复制的放置方式（placements）来描述。重载可以使用源端描述来指定更新后的权重所在位置，并使用目标端描述来指定每个推理GPU所需的权重。
举例来说，考虑一个8×8的权重矩阵，在两个训练GPU上按行分割，在四个推理GPU上按列分割。假设该张量在训练GPU上按行维度分片，在推理GPU上按列维度分片，那么每个训练GPU持有一个4×8的分片，而每个推理GPU需要一个8×2的分片。每个目标端可以通过从每个源端接收一个4×2的区域来获得其所需的分片。这些区域直接源自源端和目标端分片之间的重叠部分，如图1所示。
如图所示，任何一种分片张量都可以轻松地表示为DTensor格式，重载问题可以重新定义为DTensor到DTensor的重新分片问题。这一发现为Reshard集合通信奠定了基础。
用于DTensor到DTensor传输的Reshard集合通信
Reshard是一种新型集合通信调用，用于任意DTensor到DTensor的传输。与全归约（all-reduce）、归约-散射（reduce-scatter）、全收集（all-gather）、广播（broadcast）等操作类似，reshard定义了一种通信模式，应用程序可以将其作为单个操作来调用。它将逻辑张量的数值从源端分布形式转换为目标端分布形式。
Reshard调用只需通过源端和目标端的DTensor描述来指定。这种抽象不仅简化了重载问题，还将通信逻辑封装在库中，使底层算法可以独立于应用程序不断演进和改进。
NCCL M2N Reshard作为开源项目，是Reshard集合通信的首个参考实现，面向社区开放。该实现基于NCCL构建，可将分布在M个源GPU上的张量转换为N个目标GPU所需的布局。在RL重载场景中，训练GPU构成源端组，推理GPU构成目标端组。
图2展示了该实现如何接受布局描述和本地缓冲区，以执行Reshard操作。
NCCL M2N Reshard如何释放最大网络带宽
NCCL M2N Reshard首先推导出源端和目标端分片之间的精确交集，然后将这些传输映射到物理GPU网络拓扑结构上。该实现并非将每个源端贡献独立发送给每个目标端副本，而是按NVIDIA NVLink域对副本进行分组。在每个域内，它尽可能均匀地将接收到的贡献分配给各个副本节点，使每个接收节点通过网络接收到总数据中不同的子集。
随后，接收节点通过NVLink交换这些子集，使每个副本都能组装出完整的分片，而无需重复网络流量。如果副本跨越多个NVLink域，则每个贡献的一份副本会被转发至各个域，并在每个域内进行同样的均衡本地交换。网络传输与NVLink复制通过流水线方式并行执行，从而互相隐藏延迟。
图3展示了这种数据移动方案：张量从源端的TP4移动到目标端的TP2×DP2，且所有四个目标GPU位于同一NVLink域内。采用直接点对点传输方式时（图3a），每个源端分片会被发送给两个数据并行（DP）副本，导致每个分片都要在网络上冗余传输两次。这意味着（假设每个分片都能占满网络带宽的情况下）每个分片实际只能获得约一半的网络带宽。
NCCL M2N Reshard则只通过网络传输每个源端分片一次，为每个目标GPU提供其所需数据的一半（图3b）。目标GPU之间通过NVLink交换缺失的另一半。例如，目标GPU 0通过网络从源GPU 0接收S0，并通过NVLink从目标GPU 2接收S1。网络传输和NVLink交换同时进行。这样，每个分片就能获得完整的网络带宽。
使用NVIDIA NeMo RL演示Reshard集合通信
为展示Reshard集合通信在权重重载方面的有效性，NVIDIA NeMo RL框架中提供了一个参考实现。NeMo RL是一个开源框架，旨在简化和扩展RL后训练流程。Reshard集合通信清晰的抽象使得在该框架中的实现相对简单。对于给定的并行配置，框架只需为训练和推理实例的每个参数提供DTensor描述，而这应该很容易实现，因为DTensor正是为此目的而设计的。
然而在实践中，层类型的多样性以及训练和推理引擎在每种并行配置下分片参数方式的差异，使得通过Reshard集合通信处理所有边缘情况颇具挑战。
为设计出高效且可持续的方案，团队基于一个关键观察展开：Transformer模块中的前馈投影权重在权重重载负载中占主导地位。在最新的大语言模型中，这些权重约占模型参数总量的95%。
基于这一观察，设计采用了一种混合方法：占主导地位的前馈投影权重走reshard路径，而其他所有小型权重（例如QKV投影或mamba权重）则使用传统的"收集-广播"重载路径，该路径针对广泛的功能覆盖进行了优化。这种方法将高效的reshard通信应用于大部分负载，同时保持与传统"收集-广播"重载相同的功能覆盖范围。
分片布局可能在运行时因节点在故障恢复过程中的增减而发生变化。为适应GPU成员变化，框架可以选择销毁并重建通信组，或使用NCCL的收缩（shrink）和扩展（grow）API来调整现有的通信组。NeMo RL目前采用第一种方法来支持弹性恢复。基于NCCL收缩和扩展API的弹性支持计划在未来版本中提供。
RL权重重载性能基准测试结果
表2总结了使用上述NeMo RL实现测得的基于reshard的重载延迟数据。基准测试覆盖了在NVIDIA DGX H100（NVL8）和NVIDIA DGX GB200（NVL72）系统上运行的多种模型。在各项基准测试中，reshard重载相较于传统"收集-广播"重载实现了2.7倍至5.7倍的加速。
这种新的重载方式专为万亿参数级模型设计。在仅测试重载性能的微基准测试中，针对Qwen3.8 2.4T BF16模型，Reshard重载仅耗时4.3秒，证明了其在设计目标——万亿参数规模上的有效性。
开始使用NeMo RL体验reshard权重重载
NVIDIA NeMo RL框架中的参考实现使您能够为自己的RL工作负载采用并扩展这一方法。
要在NeMo RL中启用reshard重载路径，请使用以下配置：
policy: generation: colocated: enabled: false refit_transport: nccl_reshard
NeMo RL重载指南和NCCL Reshard配置要求文档详细说明了该功能支持的后端、布局和精度模式。
致谢
衷心感谢Ke Wen、Ching-Hsiang Chu、Guyue Huang、Seonjin Na和Nitin为NCCL M2N Reshard以及NeMo RL中nccl-reshard权重重载功能所做出的贡献。
