蚂蚁 ASystem 团队开源 NCCL 功能扩展库：AMem NCCL-Plugin

蚂蚁 asystem 团队近日宣布开源其强化学习系统中的关键组件——amem nccl-plugin，该工具旨在突破 rl（强化学习）训练过程中显存资源紧张与通信连接耗时的瓶颈，助力实现更高效的分布式计算性能。

NCCL，全称为 NVIDIA Collective Communications Library（英伟达集合通信库），是多 GPU 与多节点环境下深度学习任务的核心通信支撑库，提供如 AllReduce、AlltoAll 等高性能集体通信操作，广泛应用于大规模模型训练。

据官方介绍，AMem NCCL-Plugin 是由蚂蚁 ASystem 自主研发的 NCCL 功能增强插件，核心贡献在于新增了 ncclPause() 和 ncclResume() 两个用于显存管理的 API 接口。这一设计专门应对在 RL 训练流程中，NCCL 所分配的显存难以被有效卸载 的长期难题。

通过轻量级插件形式集成，AMem 在不中断 NCCL 通信连接的前提下，实现了对训练/推理引擎中 NCCL 显存区域的透明卸载（offload）与快速恢复（reload）。此项技术已在 Ring-1T 超大规模万亿参数推理模型 的强化学习训练场景中成功验证，展现出显著优势。

AMem NCCL-plugin总体架构图

该插件的主要优势集中体现在以下两方面：

• 显存优化能力突出：通过精准识别并解决 NCCL 内部跨 rank 显存间的交叉引用问题，确保显存可安全释放与重建。在训推模式切换期间，可在维持通信组活跃状态的同时，单卡（基于 Hopper 架构）释放超过 10GB 的显存空间；
• 切换效率极致提升：由于无需重新建立通信连接（传统方式重建耗时可达数秒），仅需卸载和恢复 NCCL 元数据，训推转换过程更加迅捷高效，大幅降低整体延迟。

下表展示了 AMem NCCL-Plugin 与其他主流方案在 Hopper 架构 GPU 上的能力对比：

**组件**	**方案**	**内存节省情况**	**每step卸载恢复耗时**
Slime注2	通过销毁和重建训练引擎通信组清理 NCCL 显存	推理不节省：残留 2GB 训练节省 10GB+	数秒
OpenRLHF	不支持卸载 NCCL 显存	推理不节省：残留 2GB 训练不节省：残留 10GB+ Remusic Remusic - 免费的AI音乐、歌曲生成工具 514 查看详情	0s
AMem	通过 Plugin 卸载 NCCL 显存	推理节省 2GB 训练节省 10GB+

AMem NCCL-plugin 功能对比

注 1: 显存释放指将显存归还操作系统；显存卸载指将显存数据暂存至 CPU pinned buffer 后释放显存；显存恢复则是重新分配显存并将数据从 CPU 拷贝回显存的过程。

注 2: slime 介绍, Slime V0.1.0. https://www.php.cn/link/5da4215028c3f94fbb1c42dab58ab173

项目代码已全面开源：https://www.php.cn/link/8552de0210a22e464a92e2d550533910

源码下载地址：点击获取

以上就是蚂蚁 ASystem 团队开源 NCCL 功能扩展库：AMem NCCL-Plugin的详细内容，更多请关注php中文网其它相关文章！