Prime Intellect 發佈 prime-rl 0.6.0，用於訓練萬億參數 MoE 模型的智能體強化學習任務

Prime Intellect 發佈了 prime-rl 版本 0.6.0，這是一個用於萬億參數混合專家（MoE）模型強化學習的開源框架。該框架專注於處理重智能體工作負載，例如長週期軟件工程任務。研究團隊在 SWE 任務上訓練了 GLM-5 模型，序列長度高達 131k，步時保持在 5 分鐘以內，批量大小為 256 次 rollout，僅使用了 28 個 H200 節點。

關鍵特性 prime-rl 0.6.0 的核心能力在於異步強化學習，它將訓練器和推理系統分離，使兩者可以獨立運行和擴展。在智能體任務中，某些編碼 rollout 可能持續數小時，異步 RL 避免了在每次策略更新前等待所有 rollout 完成，從而避免 GPU 空閒。唯一的同步點是策略更新：一旦優化器步驟完成，prime-rl 立即推送新權重。已分派的 rollout 保留其活動前綴緩存，因此單個 rollout 可能混合來自多個策略版本的令牌。

推理優化 推理通常是 RL 系統的吞吐瓶頸。prime-rl 通過多種優化提升吞吐量，同時保持延遲可控。

• FP8 推理：使用 DeepEP 和 DeepGEMM 內核降低預填充和解碼的計算精度。
• 寬專家並行（Wide EP）：將專家分佈在 32 個或更多 GPU 上，每個 GPU 持有獨立的專家並作為端點，通過調度和組合操作進行同步。
• 預填充/解碼分離（P/D 分離）：當預填充與解碼令牌比例達到 4:1 時，分離工作負載避免端到端延遲膨脹。
• KV 緩存管理：支持分層的 CPU 和磁盤卸載，使用 vLLM 原生卸載或 Mooncake Store 集中池化。
• 請求路由：默認使用 vllm-router 分支，也可插入 NVIDIA Dynamo 路由器，根據 KV 緩存重用率、隊列深度和實時負載評分。
• 路由器重放（R3）：捕捉推理路由決策並直接重放於訓練器，將 KL 偏差降低約一個數量級。

訓練優化 訓練器基於 torchtitan，採用 3D 並行：FSDP、上下文並行（CP）和專家並行（EP）。對於 GLM-5 案例，所有三種策略均被使用。

• FSDP2：按層按需收集權重，分攤內存。
• 專家並行：通過 all2all 調度和組合（torch-native 或 DeepEP）縮小活動層內存。
• 上下文並行：處理長上下文激活內存，默認使用 Ulysses 或 Ring Attention。對於 GLM-5 使用的 DSA，prime-rl 提供了自定義上下文並行實現。
• FP8 訓練：使用 DeepGEMM 塊縮放 FP8，匹配訓練器和推理精度，減少 KL 不匹配並穩定訓練。

應用場景

• 長週期 SWE 智能體：訓練模型處理真實倉庫問題，rollout 可包含數百輪和工具調用。P/D 分離保持解碼延遲可預測。
• 在更少節點上進行萬億參數級後訓練：GLM-5 僅用 28 個 H200 節點運行，寬 EP 和 KV 卸載提高了併發和吞吐。
• 大規模穩定智能體 RL：路由器重放和 FP8 訓練共同降低訓練器-推理器的 KL 不匹配，實現更穩定的訓練。

更多技術細節請查閲原文。