蚂蚁集团开源 Awex 框架，秒级完成 TB 级参数交换

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

据介绍，Awex 是为极致性能打造的训练推理引擎权重同步框架，解决 RL 流程中训练权重参数同步到推理模型的核心难题，可在秒级完成 TB 级大规模参数交换，显著降低 RL 模型训练延迟，主要特点如下：

• 极速同步性能：千卡集群万亿参数模型 6 秒内全量同步，性能领先；

• 统一模型适配层：自动处理训推引擎并行策略与引擎间的 Tensor 格式 / 布局差异，兼容多种模型架构；

• 零冗余 Resharding 传输与原地更新：仅传输必要分片（Shard），推理侧原地更新显存，避免重分配与拷贝开销；

• 多模式传输支持：支持 NCCL、RDMA、共享内存多种传输模式，充分发挥 NVLink / NVSwitch / RDMA 带宽并减少长尾延迟；

• 异构部署兼容：适配共卡 / 分卡模式，支持同步和异步 RL 算法训练场景，同时 RDMA 传输模式支持推理实例动态扩缩容；

• 灵活可插拔架构：支持对不同模型定制化权重 Sharing 和 Layout 行为，同时支持新的训练和推理引擎接入。

Awex 权重交换框架整体主要由三个组件组成：

• WeightWriter：在每个训练进程内运行，负责当前训练进程的权重 Shard 的元数据收集上报、权重转换、权重发送 Resharding 计划构建、权重发送等功能；

• WeightReader：在每个推理实例的控制进程上运行，其会在推理实例管理的每张 GPU 上面启动一个 WorkerWeightsReader，与训练进程的 WeightWriter 相对应，负责每个推理进程的权重 Shard 的元数据收集上报、权重转换、权重接受 Resharding 计划构建、权重接受等功能；

• MetaServer：Job 级别全局 Server，用于训推引擎的服务发现和权重元数据交换，以及共卡情况下的事件通知等功能；

权重交换的核心功能模块主要由 5 个部分组成：

• 训推权重统一转换：负责将不同并行策略和 Tensor 布局的训练引擎和推理引擎的权重转换成统一的格式，用于后续的权重元数据计算和权重传输；

• 全局权重元数据计算与交换：将训推权重转换成统一的格式后，收集每个 Worker 的所有权重 Shard 元数据，并上报到 Meta Server，用于接下来的的权重传输计划构建；

• P2P 权重传输执行计划：训练和推理引擎拿到全局所有 Worker 的训练和推理权重 Shard 元数据，然后分别各自构建对等的发送和接受确定性传输计划；

• NCCL 权重传输：使用 NCCL 的 send / recv API 基于构建的传输计划进行对等的权重发送与接收；

• RDMA 权重传输：使用 NUMA 亲和，面向全局负载均衡传输计划的 RDMA 通信来进行权重的更新；

同时 Awex 也支持对权重进行 Tensor 级别的校验，将通过文件系统模式加载的权重跟通过传输模式加载的权重进行 Tensor 级别的细粒度比对，逐个判断差异，保证传输模式的正确性。

在千卡规模的集群上，Awex 使用 NCCL 传输数据可以在一秒内完成 10B 规模的模型权限交换，二十秒内完成 1T 规模的模型权重交换，使用 RDMA 进行传输，1T 模型权重交换耗时可以进一步缩短到六秒钟。

未来展望，Awex 是蚂蚁 ASystem 强化学习系统的核心组件之一，而 ASystem 是百灵万亿模型训练的坚实基础。我们将在未来陆续开源 ASystem 的其他核心 RL 组件，进一步完善开源强化学习训练生态。目前 Awex 开源版已支持 Megatron 和 SGLang 引擎。

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