蚂蚁 ASystem 团队宣布开源其强化学习系统的关键组件：AMem NCCL-Plugin，用于解决 RL 训练中的显存瓶颈和通信连接的耗时难题，以实现高性能计算。

NCCL 是 NVIDIA Collective Communications Library（英伟达集合通信库）的缩写，它是多 GPU 和多节点分布式深度学习的核心通信库，提供了包括 AllReduce、AlltoAll 等多种高效集体通信操作。

据介绍，AMem NCCL-Plugin 是蚂蚁 ASystem 团队自研的 NCCL 功能扩展库，主要提供了 ncclPause() 和 ncclResume() 两个显存管理 API，旨在解决 RL 流程中，通信库 NCCL 分配的显存无法被高效卸载的这一难题。

通过轻量级插件方式，在保留 NCCL 通信连接的情况下，实现对训推引擎 NCCL 显存的透明卸载（offload）与恢复（reload）注1，这些优势特点已在 Ring-1T 万亿参数推理模型的强化学习训练中得到了验证。

AMem NCCL-plugin总体架构图

AMem NCCL-Plugin 的优势体现在如下两个方面：

• 显存节约：通过识别并解决 NCCL 通信库中 cross-rank 的显存交叉引用问题，实现正确的显存透明释放与恢复。在训推切换时，可在保持通信组连接的情况下，单卡（Hopper 架构卡）释放出 10GB+ 显存；

• 极致高效：因为保留了通信组的连接，训推转换仅卸载和恢复 NCCL 的元数据，无需重建通信连接（典型耗时为数秒钟），从而实现典型耗时 <1 秒的极致优化。

与社区已知方案的在 Hopper 架构卡上的能力对比：

AMem NCCL-plugin 功能对比

注 1: 显存释放：把显存交还操作系统；显存卸载：把显存中的信息放入CPU pinned buffer, 然后释放显存；显存恢复：把显存重新分配回来，把暂存 CPU pinned buffer 的信息拷贝回显存中。

注 2: slime 介绍, Slime V0.1.0. https://zhuanlan.zhihu.com/p/1945237948166547268

代码地址：https://github.com/inclusionAI/asystem-amem