XYZ-AI-Lab/axrl

解決的問題

AxisRL 是為代理型強化學習（RL）後訓練所設計，模型必須在多回合中與環境互動、呼叫工具，並根據長期結果獲得獎勵。它解決了管理多回合回放、環境狀態、工具呼叫與權重同步的協調挑戰，同時確保訓練過程穩定且可觀測。

如何運作

AxisRL 作為一個系統層，連結高吞吐量的回放引擎（SGLang）與大規模分散式訓練引擎（Megatron）。它管理一個循環：回放代理執行代理工作流程，SGLang 工作者產生文字，環境產生獎勵，Megatron 工作者執行策略最佳化（如 PPO 或 GRPO）。

關鍵技術機制包括：

• 白箱與黑箱整合：支援內部代理迴圈（白箱）或透過 OpenAI 相容代理捕捉外部框架的互動（黑箱）。
• 部分回放：允許已完成的樣本立即傳送至訓練器，以減少因長尾軌跡延遲所造成的閒置時間。
• 基於句柄的資料平面：透過句柄而非中央驅動程式傳輸大型資料（如 MoE 路由資料），以避免瓶頸。
• 上下文管理：使用前綴樹合併與 MagiAttention，優化多回合情境中的重複注意力運算。
• 一致性工具：包含不匹配分析與「脈衝重播」，用於調試回放與訓練器執行路徑之間的差異。

適用對象

適用於訓練大型代理型 LLM（數百億參數）的研究人員與工程師，這些模型需要複雜的多回合互動與高吞吐量 RL 後訓練。

主要特色

• 基於 SGLang 實作回放，基於 Megatron 實作分散式訓練。
• 支援多種最佳化目標，包括 PPO、GRPO/GRPO2、GSPO、TOPR 與 TIS。
• 可處理超過 300 回合的軌跡。
• 提供 Rollout Routing Replay（R3）以在訓練期間穩定 MoE 專家路由。
• 提供可重現的除錯套件，用於分析損失突增與 token 級別不匹配。