IT之家 11 月(yue) 20 日消息，螞(ma)蟻集團今(jin)日宣布開(kai)源萬(wan)億參數強化學習高(gao)性(xing)能權重交(jiao)換框(kuang)架 Awex。
據介紹，Awex 是為極致性能打造的訓練推理引擎權重同步框架，解決 RL 流程中訓練權重參數同步到推理模型的核心難題，可在秒級完成 TB 級大規模參數交換，顯著(zhu)降低 RL 模型訓練(lian)延(yan)遲，主要(yao)特點(dian)如下：
極速同步性能：千卡集群萬億(yi)參(can)數模(mo)型 6 秒內全量同步，性(xing)能(neng)領(ling)先；
統一模型適配層：自動(dong)處(chu)理訓推引擎并行策(ce)略與(yu)引擎間的 Tensor 格式 / 布(bu)局(ju)差(cha)異，兼(jian)容多種模型架(jia)構；
零冗余 Resharding 傳輸與原地更新：僅傳(chuan)輸必(bi)要分片（Shard），推(tui)理側原地更新顯存，避(bi)免重(zhong)分配與拷(kao)貝開銷；
多模式傳輸支持：支持 NCCL、RDMA、共享內存多種傳輸模(mo)式，充分發揮(hui) NVLink / NVSwitch / RDMA 帶寬并減少長尾延遲；
異構部署兼容：適(shi)配共卡 / 分(fen)卡模(mo)式，支(zhi)持同步和(he)異步 RL 算法訓練場(chang)景，同時 RDMA 傳(chuan)輸模(mo)式支(zhi)持推理實例動態擴縮(suo)容；
靈活可插拔架構：支(zhi)持(chi)對不(bu)同模型定制化(hua)權重 Sharing 和(he) Layout 行為，同時支(zhi)持(chi)新的訓練和(he)推理引(yin)擎(qing)接入。
Awex 權重交換框架整體主要由三個組件組成：
WeightWriter：在每個訓練(lian)進程內運行，負責當前(qian)訓練(lian)進程的(de)權(quan)重 Shard 的(de)元數據(ju)收集上報、權(quan)重轉換、權(quan)重發送(song) Resharding 計劃(hua)構建(jian)、權(quan)重發送(song)等(deng)功能；
WeightReader：在每個推理實例的(de)(de)控制進(jin)程上運行(xing)，其會在推理實例管理的(de)(de)每張 GPU 上面啟動(dong)一個 WorkerWeightsReader，與訓練進(jin)程的(de)(de) WeightWriter 相對應(ying)，負責(ze)每個推理進(jin)程的(de)(de)權(quan)重 Shard 的(de)(de)元數據收集上報、權(quan)重轉(zhuan)換、權(quan)重接受 Resharding 計劃構(gou)建、權(quan)重接受等(deng)功(gong)能；
MetaServer：Job 級別(bie)全(quan)局 Server，用于訓推引(yin)擎(qing)的服務發現和權重元數據交(jiao)換(huan)，以及共卡(ka)情況下(xia)的事件通知等功能(neng)；
權重交換的核心功能模塊主要由 5 個部分組成：
訓(xun)推權(quan)(quan)(quan)重統一(yi)轉(zhuan)換：負(fu)責將(jiang)不(bu)同并行策略和 Tensor 布(bu)局的訓(xun)練引(yin)擎和推理(li)引(yin)擎的權(quan)(quan)(quan)重轉(zhuan)換成(cheng)統一(yi)的格式，用于后續的權(quan)(quan)(quan)重元數據計算和權(quan)(quan)(quan)重傳輸；
全局權重元數據計算與交換：將訓推權重轉換成統一的格式后，收集每個 Worker 的所有權重 Shard 元數據，并上報到 Meta Server，用于接下(xia)來的的權(quan)重傳輸計劃構建(jian)；
P2P 權重傳(chuan)輸(shu)執行計劃：訓練(lian)和(he)推(tui)理引擎拿到全(quan)局所有 Worker 的訓練(lian)和(he)推(tui)理權重 Shard 元數(shu)據，然后分別(bie)各自構建對等的發送和(he)接受確(que)定性傳(chuan)輸(shu)計劃；
NCCL 權(quan)重(zhong)傳(chuan)輸：使用(yong) NCCL 的(de)(de) send / recv API 基于構建的(de)(de)傳(chuan)輸計劃進行對等的(de)(de)權(quan)重(zhong)發送(song)與接(jie)收；
RDMA 權重傳輸：使用(yong) NUMA 親和，面向全局負載均衡傳輸計劃的 RDMA 通信(xin)來進行權重的更新；
同時(shi) Awex 也支持對權重(zhong)進行(xing) Tensor 級(ji)別的校(xiao)驗，將通過文件系統模(mo)式加載(zai)的權重(zhong)跟通過傳(chuan)輸模(mo)式加載(zai)的權重(zhong)進行(xing) Tensor 級(ji)別的細粒度比對，逐(zhu)個判斷差異，保證傳(chuan)輸模(mo)式的正確性。
在千卡規模的集群上，Awex 使用 NCCL 傳輸數據可以在一秒內完成 10B 規模的模型權限交換，二十秒內完成 1T 規模的模型權重交換，使用 RDMA 進行傳輸(shu)，1T 模型權重交換耗時可以(yi)進一步縮短到(dao)六秒(miao)鐘。
未來展望，Awex 是螞蟻 ASystem 強化學習系統的核心組件之一，而 ASystem 是百靈萬億模型訓練的堅實基礎。我們將在未來陸續開源 ASystem 的其他核心 RL 組件，進一步完善開源強化學習訓練生態。目前 Awex 開源版已支持 Megatron 和 SGLang 引擎。
IT之(zhi)家附(fu)開源地址：
廣(guang)告(gao)聲明：文內含(han)有(you)的對(dui)外跳轉鏈(lian)接（包括不限于超鏈(lian)接、二維(wei)碼、口令等(deng)形式），用于傳遞更多信息，節省甄選時間，結果僅供參(can)考，IT之家所有(you)文章均包含(han)本聲明。