LongStraw Pushes Long-Context RL Past 2M Tokens on a Fixed GPU Budget

Lead

Inference systems are rapidly moving toward million-token context windows, but RL post-training has not kept pace. Many training workloads still operate at 256K tokens or below, leaving deployment to rely on length generalization. That gap matters most for AI agents, where observations, tool outputs, retrieved documents and previous decisions can accumulate over long trajectories.

LongStraw addresses this mismatch from an execution perspective. It is not introduced as a new model, but as an architecture-aware stack for running long-context RL post-training under a fixed GPU budget. The paper instantiates the approach with Group Relative Policy Optimization, or GRPO.

Key ideas

• Separate the shared prompt from response branches: LongStraw first evaluates the long shared prompt without automatic differentiation, avoiding a full training graph over the entire prompt.
• Keep only the state needed by later tokens: Instead of retaining every activation for the full context, it stores model-specific state that later tokens depend on.
• Replay short branches under autograd: Response branches are replayed one at a time when gradients are needed. This reduces peak memory, at the cost of additional replay time.
• Support different model architectures: The implementation covers Qwen3.6-27B, which mixes recurrent and full-attention components, as well as GLM-5.2, a compressed-attention mixture-of-experts model.

The reported results focus on execution capacity. On eight H20 GPUs, LongStraw completes grouped Qwen scoring and response backward passes at 2.1M positions for group sizes 2 and 8. Increasing the group size adds only 0.21 GB of peak allocated memory. A separate stress test reaches 4.46M positions. For GLM-5.2, the authors validate the end-to-end execution path for a 2.1M-token prompt across all 78 layers on 32 H20 GPUs.

Why it matters

The contribution is important because it shifts the conversation from “Can a model read a million tokens?” to “Can we post-train behavior over trajectories that long?” For agents, this is a practical requirement: long-running tasks often involve evolving context, repeated tool use and accumulated state rather than a single static prompt.

The caveat is equally important. The paper states that the experiments establish execution capacity rather than complete training correctness. The captured prompt state is detached, and some distributed forward and gradient composition paths remain incomplete. In that sense, LongStraw is best read as a systems blueprint for million-token RL, not a finished training recipe. Its next major test will be whether the same idea works when agent trajectories grow unpredictably and are not cleanly split into prompt and generation.

Source: Hugging Face Daily Papers