NVIDIA发布Reshard集合通信,开源NCCL M2N实现加速RL权重同步
NVIDIA技术博客介绍,将强化学习训练与推理之间的权重重分发(refit)重新表述为DTensor到DTensor的重分片问题,并在NeMo RL中给出参考实现;在DGX H100与DGX GB200的评测中,延迟较传统gather-and-broadcast方式快 2.7 至 5.7 倍。
AI解读:NVIDIA技术博客发文介绍Reshard集合通信,用于加速强化学习后训练中训练GPU向推理GPU同步最新权重的refit过程;作者为Tanya Lenz。
文章指出,现有gather-and-broadcast做法会把完整参数广播给每张推理GPU,通信量约为模型权重的两倍;按BF16计算,Kimi K3和Qwen 3.8 2.4T A95B单次refit逻辑载荷分别约 11.2 TB和 9.6 TB,且每个训练步都要执行。
Reshard把refit建模为DTensor到DTensor的重分片,只传输源分片与目标分片实际需要的交集区域;NCCL M2N Reshard作为首个开源参考实现,已在NVIDIA/nccl-extensions GitHub仓库发布。
NeMo RL中的参考实现采用混合路径:占比约 95% 的前馈投影权重走reshard,其余小权重仍用gather-and-broadcast,以覆盖更多层类型。
博客给出的基准显示,在DGX H100和DGX GB200上,reshard refit相比传统方式加速 2.7 至 5.7 倍;Qwen3.8 2.4T BF16模型的refit-only微基准耗时为 4.3 秒。
启用方式为在NeMo RL配置中设置generation.colocated.enabled: false与refit_transport: nccl_reshard;故障恢复时的通信组调整目前采用销毁重建,基于NCCL shrink/grow API的支持列为后续计划。
NVIDIA技术博客发布文章,介绍一种名为Reshard的新集合通信类型,主要面向强化学习后训练中的权重refit(把训练GPU上更新的策略权重同步到推理GPU)。作者Tanya Lenz称,其开源参考实现NCCL M2N Reshard已通过NVIDIA/nccl-extensions GitHub仓库提供。
文章将refit的难点归结为两点:训练与推理使用不同的并行映射,权重分片方式不一致;且每次refit要更新全部模型权重,通信载荷与模型权重总量同量级。文章称,现有gather-and-broadcast方式先all-gather再广播,推理GPU即使只需要一个分片也会收到完整参数,因而产生冗余通信。
把refit改写成DTensor到DTensor
文章提出,refit应只传输理论最少量的权重数据,并把这一过程重新表述为将一个分布式张量(DTensor)从训练布局重分片到推理布局。DTensor布局由全局形状、描述GPU逻辑排布的device mesh,以及各维度分片或复制方式组成。
以 8×8 权重矩阵为例:训练侧 2 张GPU按行切分,推理侧 4 张GPU按列切分,每个目标分片可从每个源接收一个 4×2 区域,这些区域由源分片与目标分片的交集直接得出。Reshard调用只需给出源和目标两个DTensor描述。
- Reshard被定义为一种可像all-reduce、broadcast一样调用的集合通信操作,传输逻辑张量从源分布到目标分布的取值。
- NCCL M2N Reshard将分布在M张源GPU上的张量转换为N张目标GPU所需的布局,训练GPU为源组、推理GPU为目标组。
带宽利用方式与NeMo RL实现
NCCL M2N Reshard先计算源分片与目标分片的精确交集,再映射到物理网络拓扑。它按NVLink域对副本分组,域内尽量均匀分配网络流入的数据,接收rank之间再通过NVLink交换,以避免同一份数据在网络上重复传输;网络传输与NVLink复制做流水线重叠。
在NeMo RL框架的参考实现中,团队观察到前馈投影权重在最新LLM中约占参数的 95%,因此让这部分主导载荷走reshard路径,QKV投影、mamba权重等其他小权重仍走传统gather-and-broadcast路径。文章称,这样兼顾了大部分载荷的通信效率与传统路径的功能覆盖。
- NCCL M2N Reshard的示例:将TP4分片的张量转换为 4 张推理GPU(DP4)上的副本。
- 博客中以TP4到TP2×DP2为例:直接点对点传输会把每个源分片向两个数据并行副本各发一次,冗余一倍;M2N Reshard让每张目标GPU经网络收到一半数据,另一半经NVLink从同域目标GPU获取。
基准数据与使用条件
文章给出的延迟数据覆盖DGX H100(NVL8)与DGX GB200(NVL72)两类系统。在NeMo RL中,Qwen 3 235B从TP4×PP8×EP16转到TP8×DP16,传统方式 6.4 秒、reshard 1.9秒;DeepSeek v3 671B在H100上从PP16×EP16转到TP32×DP8,从 14.9 秒降至 2.6 秒,在GB200上从 6.1 秒降至 2.3 秒;Nemotron Ultra v3从 8.2 秒降至 2.3 秒。文章称整体加速 2.7 至 5.7 倍。
文章还称,面向万亿参数模型的refit-only微基准中,Qwen3.8 2.4T BF16的Reshard refit耗时 4.3 秒。这些数字为NVIDIA在自家框架与硬件上的测量结果,用于说明该实现的延迟表现。
- 在NeMo RL中启用reshard路径需配置policy.generation.colocated.enabled: false和refit_transport: nccl_reshard。
- 布局在节点增减或故障恢复时可能变化;NeMo RL当前通过销毁并重建通信组来应对,基于NCCL shrink和grow API的弹性支持列为未来更新。
- 文章致谢Ke Wen、Ching-Hsiang Chu、Guyue Huang、Seonjin Na和Nitin对NCCL M2N Reshard及NeMo RL中nccl-reshard权重refit的贡献。