蚂蚁集团开源Awex框架：高效TB级数据交换

11月20日，蚂蚁集团正式宣布开源其万亿参数强化学习高性能权重交换框架——Awex。

作为一款专为极致性能打造的训练推理引擎权重同步框架，Awex解决了强化学习（RL）过程中训练权重参数同步到推理模型的关键难题。该框架实现了突破性进展：在大规模TB级数据传输中，可实现秒级完成，显著降低了RL模型的训练延迟。

以下是Awex的核心优势：

• 卓越性能表现： 在千卡集群环境下，处理万亿参数模型仅需6秒即可完成全量同步。

• 统一的模型适配能力： 自动化处理训练推理引擎间的并行策略差异，并兼容不同Tensor格式与布局，支持多种模型架构。

• 高效传输机制： 采用零冗余Resharding技术进行传输和显存原地更新，仅传输必要数据分片（Shard），避免不必要的内存分配和拷贝开销。

• 多模式传输支持： 兼容NCCL、RDMA及共享内存等多种传输协议，充分释放NVLink/NVSwitch/RDMA带宽潜力，有效减少网络延迟。

• 灵活部署能力： 支持共卡与分卡部署模式，并兼容同步和异步RL算法训练场景。RDMA传输模式下可实现推理实例的动态扩展。

• 模块化设计： 允许针对不同模型进行定制化权重共享与布局行为，同时支持新训练和推理引擎的接入。

Awex框架主要由三个核心组件构成：

• WeightWriter: 在每个训练进程中运行，负责收集并上报权重Shard的元数据，执行权重转换、传输计划构建及发送操作。

• WeightReader: 运行于推理实例的控制进程，为每张GPU启动WorkerWeightsReader，与训练进程中的WeightWriter对应交互。负责目标权重Shard的元数据收集、转换和接收。

• MetaServer: 作为Job级全局服务器，承担训推引擎服务发现、权重元数据交换功能，并在共卡模式下处理事件通知。

Awex的核心功能模块包括：

• 训推权重统一转换: 将训练引擎和推理引擎基于不同并行策略和Tensor布局的权重，转化为一致格式以供后续处理。

• 全局元数据计算与交换: 收集各Worker的所有权重Shard元数据，并上报至MetaServer，用于构建传输计划。

• P2P传输执行计划: 训练和推理引擎根据收集到的元数据，分别制定确定性的发送和接收传输计划。

• NCCL权重传输: 利用NCCL的send/recv API按构建的传输计划执行点对点传输。

• RDMA权重传输: 通过NUMA亲和性和全局负载均衡传输计划，利用RDMA通信进行高效更新。

Awex还提供了一套完整的权重验证机制：将文件系统加载的权重与传输模式加载的权重进行逐个Tensor级别的比对，确保数据传输的准确性。

性能实测显示： 在千卡规模集群中，使用NCCL传输时，10B规模模型仅需一秒完成同步；而1T规模模型则在二十秒内完成。采用RDMA传输模式下，1T模型权重交换耗时更可缩短至六秒钟。

未来规划： Awex作为蚂蚁集团”A系统”中的核心组件，是支持百灵万亿模型训练的重要基石。蚂蚁计划陆续开源该系统的其他关键RL组件，持续丰富强化学习训练的开源生态。目前，Awex已支持Megatron和SGLang引擎。

项目开源地址： https://github.com/inclusionAI/asystem-awex