LongStraw:固定GPU预算下超200万token的长上下文强化学习后训练

LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget

精选理由

LongStraw 能让你在固定 GPU 预算下做百万 token 级别的强化学习后训练,比现有 256K 的方法多处理 8 倍上下文,而且内存增长极低。

AI 摘要

LongStraw 是一种针对百万token级别强化学习后训练的架构感知执行栈,基于 GRPO 实现,在固定 GPU 预算下运行。它通过共享提示无自动求导评估、仅保留模型特定状态并逐次重放短响应分支,将实时训练图缩小以换取重放时间。在 8 块 H20 GPU 上,LongStraw 完成了 2.1M 位置的分组 Qwen 评分和响应反向传播;分组数从 2 增至 8 时,峰值显存仅增加 0.21 GB。在 32 块 H20 GPU 上,端到端执行路径验证了 GLM-5.2 全部 78 层的 2.1M token 提示处理能力。这些实验证明了执行容量,但完整训练正确性尚未全面验证。

原文 · arXiv cs.LG

LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget

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.