LongStraw:在固定GPU預算下超越200萬Token的長上下文強化學習
LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget
July 16, 2026
作者: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin
cs.AI
摘要
推論上下文長度與強化學習後訓練之間的差距日益擴大:推論系統正接近百萬級標記的上下文,而後訓練負載通常仍停留在256K標記或以下,並在部署時依賴長度泛化。這一差距對AI智能體尤為重要,因其觀察結果、工具輸出、文件以及先前的決策會在長時間軌跡中不斷累積。LongStraw是一個基於固定GPU預算、專為百萬級標記強化學習後訓練設計的架構感知執行棧,並以群體相對策略優化(Group Relative Policy Optimization, GRPO)實例化。該方法對共享提示進行不帶自動梯度的評估,僅保留後續標記所需的模型特定狀態,並逐一重播短回應分支,以額外重播時間為代價縮減即時訓練圖。我們針對混合循環與全注意力機制的Qwen3.6-27B以及壓縮注意力混合專家模型GLM-5.2進行了實作。在八張H20 GPU上,LongStraw針對群組大小為2和8的場景,完成了為2.1M位置分組的Qwen評分與回應反向傳播;增加群組大小僅額外增加0.21 GB的峰值記憶體配置量,而獨立的壓力測試則達到了4.46M位置。在三十二張H20 GPU上,我們驗證了GLM-5.2全部78層、2.1M標記提示的端到端LongStraw執行路徑。由於所捕獲的提示狀態為分離狀態,且部分分散式前向傳播與梯度組合路徑尚未完整,這些實驗僅建立了執行容量,而非完整的訓練正確性。
English
A growing gap separates inference context lengths from RL post-training: inference systems are approaching million-token contexts, while post-training workloads often remain at 256K tokens or below and rely on length generalization at deployment. The gap is especially important for AI agents, whose observations, tool outputs, documents, and prior decisions accumulate over long trajectories. LongStraw is an architecture-aware execution stack for million-token RL post-training under a fixed GPU budget, instantiated with Group Relative Policy Optimization (GRPO). It evaluates the shared prompt without autograd, retains only model-specific state needed by later tokens, and replays short response branches one at a time, reducing the live training graph at the cost of additional replay time. We implement it for the hybrid recurrent and full-attention Qwen3.6-27B and the compressed-attention mixture-of-experts GLM-5.2. On eight H20 GPUs, LongStraw completes grouped Qwen scoring and response backward at 2.1M positions for groups of 2 and 8; increasing the group size adds only 0.21 GB of peak allocated memory, while a separate stress test reaches 4.46M positions. On 32 H20 GPUs, we validate the end-to-end LongStraw execution path for a 2.1M-token prompt across all 78 layers of GLM-5.2. These experiments establish execution capacity rather than complete training correctness because the captured prompt state is detached and some distributed forward and gradient composition paths remain incomplete.