蚂蚁开源Awex框架,秒级同步万亿模型权重

Source

品玩11月21日讯,据百灵大模型官方消息,蚂蚁集团开源高性能权重交换框架Awex,该框架旨在解决大模型强化学习训练中,训练权重参数同步至推理引擎的核心延迟难题。

Awex可在秒级完成TB级大规模参数交换。在千卡GPU集群上,实现了万亿参数模型5-10秒内完成权重同步,显著降低训练延迟。其核心技术包括统一模型适配层,可自动处理训练与推理引擎间的并行策略与权重格式差异;采用零冗余Resharding传输,仅传输必要分片;支持NCCL、RDMA等多种传输模式。

该框架已成功支撑蚂蚁百灵万亿参数模型Ring-1T的训练。代码已在GitHub开源,目前支持Megatron训练引擎与SGLang推理引擎。