螞蟻集團宣佈開源高性能權重交換框架 Awex（Asystem Weight Exchange Framework），實現了在數千卡 GPU 集羣上 5~10 秒內完成萬億參數級別的權重同步。

據介紹，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，用於訓推引擎的服務發現和權重元數據交換，以及共卡情況下的事件通知等功能；

Awex 是螞蟻 ASystem 強化學習系統的核心組件之一，而 ASystem 是百靈萬億模型訓練的堅實基礎。團隊稱將在未來陸續開源 ASystem 的其他核心 RL 組件，進一步完善開源強化學習訓練生態。目前 Awex 開源版已支持 Megatron 和 SGLang 引擎。

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