蚂蚁集团开源Awex框架,万亿参数模型实现秒级同步突破
11月20日,蚂蚁集团宣布开源其自主研发的万亿参数强化学习高性能权重交换框架Awex。这一突破性技术将极大推动大规模强化学习模型的训练效率,为解决训练权重参数同步到推理模型的核心难题提供了全新方案。
Awex框架最引人注目的成就是其卓越的同步性能。在千卡集群环境下,该框架能够在短短6秒内完成万亿参数模型的全局同步,相比传统方法实现了数量级的提升。这一突破意味着,研究人员在训练超大规模强化学习模型时,不再需要为权重同步耗费数小时甚至更长的等待时间,从而显著降低了模型训练的整体延迟。
该框架的核心由三个精密配合的组件构成。WeightWriter组件负责训练进程中的权重分片元数据收集和发送;WeightReader组件则在推理实例中对应地进行权重接收;而MetaServer作为全局调度中心,协调整个权重交换过程。这种分工明确的架构设计,确保了系统在处理海量参数时的稳定性和效率。
Awex框架通过多项创新技术实现了性能的极致优化。其“零冗余Resharding传输”技术能够智能识别并仅传输必要的参数分片,同时在推理侧实现显存的原位更新,避免了重复分配与拷贝的开销。框架还支持NCCL、RDMA和共享内存多种传输模式,能够充分利用硬件带宽并有效减少传输延迟。
特别值得关注的是,该框架具备出色的兼容性和灵活性。它不仅能够自动处理训练与推理引擎间的并行策略差异,还支持共卡和分卡的不同部署模式。研究人员可以根据具体需求,对权重共享和布局行为进行定制化配置,展现出框架强大的适应能力。
蚂蚁集团此次选择将Awex框架开源,体现了其推动人工智能技术发展的开放态度。作为蚂蚁ASystem强化学习系统的核心组件,Awex目前已经支持Megatron和SGLang等主流引擎。集团表示,未来还将陆续开源ASystem的其他核心强化学习组件,进一步完善开源强化学习训练生态。