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

新闻中心 NEWS CENTER

2025-11-20

浏览次数：次

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

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

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

WeightWriter：在每个训练进程内运行，负责当前训练进程的权重 Shard 的元数据收集上报、权重转换、权重发送 Resharding 计划构建、权重发送等功能；
WeightReader：在每个推理实例的控制进程上运行，其会在推理实例管理的每张 GPU 上面启动一个 WorkerWeightsReader，与训练进程的 WeightWriter 相对应，负责每个推理进程的权重 Shard 的元数据收集上报、权重转换、权重接受 Resharding 计划构建、权重接受等功能；
MetaServer：Job 级别全局 Server，用于训推引擎的服务发现和权重元数据交换，以及共卡情况下的事件通知等功能；

☞☞☞AI 智能聊天, 问答助手, AI 智能搜索, 免费无限量使用 DeepSeek R1 模型☜☜☜

权重交换的核心功能模块主要由 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 框架，秒级完成 TB 级参数交换 -