蚂蚁集团宣布开源高性能权重交换框架 Awex（Asystem Weight Exchange Framework），实现了在数千卡 GPU 集群上 5~10 秒内完成万亿参数级别的权重同步。

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

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

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