蚂蚁开源高性能权重交换框架 Awex,支持万亿参数强化学习
蚂蚁集团宣布开源高性能权重交换框架 Awex(Asystem Weight Exchange Framework),实现了在数千卡 GPU 集群上 5~10 秒内完成万亿参数级别的权重同步。 在千卡规模的集群上,Awex 使用 NCCL 传输数据可以在一秒内完成 10B 规模的模型权限交换,二十秒内完成 1T规模的模型权重交换,使用 RDMA 进行传输,1T 模型权重交换耗时可以进一步缩短到六秒钟。 据介绍,Awex 是为极致性能打造的训练推理引擎权重同步框架,解决 RL 流程中训练权重参数同步到推理模型的核心难题,可在秒级完成 TB 级大规模参数交换,显著降低 RL 模型训练延迟,主要特点如下: 极速同步性能:千卡集群万亿参数模型 6 秒内全量同步,性能领先; 统一模型适配层:自动处理训推引擎并行策略与引擎间的 Tensor 格式 / 布局差异,兼容多种模型架构; 零冗余 Resharding 传输与原地更新:仅传输必要分片(Shard),推理侧原地更新显存,避免重分配与拷贝开销; 多模式传输支持:支持 NCCL、RDMA、共享内存多种传输模式,充分发挥 NVLink / NVSwi...
