NVIDIA 推出 Reshard 集合通信加速 RL 权重重载
原标题:Accelerating Reinforcement Learning Weight Refit with the Reshard Collective
AI 摘要
NVIDIA 技术博客介绍了一种名为 Reshard 的新型集合通信原语,专门用于加速强化学习后训练中的权重重载(refit)。传统 gather-and-broadcast 方法在训练与推理 GPU 布局不同时会产生约两倍于模型权重的冗余通信量,对 Kimi K3 等万亿参数模型每次 refit 可达约 10 TB。Reshard 将 refit 重新表述为 DTensor 从训练布局到推理布局的 reshard 问题,只传输理论最小数据量,其开源参考实现 NCCL M2N Reshard 已在 NVIDIA/nccl-extensions 仓库发布。
正文节选
During reinforcement learning (RL) post-training, training trajectories are generated on the fly by running inference with the current policy model. For training efficiency, it is common to run inference and training in parallel on separate GPU resources in a disaggregated setup. Although the policy model is updated only on the training GPUs, the system must periodically refresh the weights on the inference GPUs with the latest parameters generated on the training GPUs. This transfer of model we