蚂蚁开源高性能权重交换框架 Awex，支持万亿参数强化学习

发布时间：2025-11-20

点击次数：

蚂蚁集团宣布开源高性能权重交换框架 awex（asystem weight exchange framework），实现了在数千卡 gpu 集群上 5~10 秒内完成万亿参数级别的权重同步。

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

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

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

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

蚂蚁开源高性能权重交换框架 Awex，支持万亿参数强化学习

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

训推权重统一转换：负责将不同并行策略和 Tensor 布局的训练引擎和推理引擎的权重转换成统一的格式，用于后续的权重元数据计算和权重传输；
全局权重元数据计算与交换：将训推权重转换成统一的格式后，收集每个 Worker 的所有权重 Shard 元数据，并上报到 Meta Server，用于接下来的的权重传输计划构建；
P2P 权重传输执行计划：训练和推理引擎拿到全局所有 Worker 的训练和推理权重 Shard 元数据，然后分别各自构建对等的发送和接受确定性传输计划；
NCCL 权重传输：使用 NCCL 的 send/recv API基于构建的传输计划进行对等的权重发送与接收；
RDMA 权重传输：使用 NUMA 亲和，面向全局负载均衡传输计划的 RDMA 通信来进行权重的更新；

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

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

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

源码地址：点击下载

以上就是蚂蚁开源高性能权重交换框架 Awex，支持万亿参数强化学习的详细内容，更多请关注其它相关文章！