蚂蚁 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+ | 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
源码下载地址：点击获取
