Awex: 实现万亿参数模型秒级权重同步的强化学习训练推理框架

Awex是由InclusionAI开源的高性能强化学习权重同步框架，核心目标是解决RLHF等强化学习训练中训练端与推理端之间的参数更新延迟问题。该框架已在千卡规模GPU集群上验证，支持万亿参数模型在10秒内完成全量权重同步，为大规模强化学习训练提供高效协同能力。

在大语言模型的强化学习训练（如RLHF、DPO等）过程中，传统权重同步方式需先将权重写入存储系统再由推理端加载，耗时数分钟甚至更久。这种延迟严重制约了算法迭代效率，尤其在在线RL场景中，推理端需频繁使用最新模型生成响应，同步瓶颈会显著影响训练吞吐量与收敛速度。

Awex的核心技术特性包括：

1. 极致同步速度：NCCL模式下100亿参数同步0.8秒、万亿参数20秒；RDMA模式万亿参数仅需6秒；
2. 统一权重适配层：自动处理训练（如Megatron）与推理（如vLLM）的并行策略和张量布局差异；
3. 零冗余传输与原地更新：仅传输必要权重分片，推理端GPU内存原地更新避免额外开销；
4. 多模式传输支持：兼容NCCL、RDMA、共享内存等高速互联技术；
5. 异构部署兼容：支持同地/分离部署，适配同步与异步RL算法需求。

架构组件

• WeightWriter：训练节点收集权重分片元数据、转换格式、构建传输计划；
• WeightReader：推理实例接收权重数据并完成本地更新；
• MetaServer：全局元数据交换与协调枢纽。

权重交换流程

1. 统一格式转换：将不同引擎（Megatron、vLLM等）的权重转为标准格式；
2. 全局元数据交换：收集分片元数据并上报MetaServer；
3. P2P传输计划构建：基于元数据生成点对点传输计划；
4. 传输执行：使用NCCL/RDMA进行数据传输；
5. 张量级验证：比对传输与文件加载的权重确保正确性。

Awex的性能在基准测试中表现领先，能有效解决同步瓶颈。适用场景包括：

• 在线RLHF训练：需频繁同步最新模型生成高质量训练数据；
• 多轮迭代优化：快速迭代场景下减少训练周期；
• 大规模集群训练：千卡/万卡规模下高效协同；
• 实时推理服务：快速部署最新模型版本的生产环境。

Awex通过创新架构与高效传输机制，成功解决了大规模RL训练中的权重同步瓶颈。其秒级同步能力使万亿参数模型的在线强化学习训练成为可能，为大语言模型持续优化提供坚实支撑。未来随大模型规模增长，这类专门优化的权重同步框架将在AI基础设施领域发挥更重要作用。