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

Abstract: 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.

• The paper demonstrates that long-context RL post-training (up to 2M tokens) is feasible using fixed GPU budgets by decoupling prompt evaluation from response replay.
• The methodology leverages explicit tensor lifetime management and GRPO to efficiently scale training despite challenges from quadratic attention and massive memory demands.
• Empirical results on Qwen3.6-27B and GLM-5.2 validate the approach, achieving million-token updates with minimal memory overhead and outlining practical system limitations.

LongStraw: Architecture-Aware RL Post-Training for Million-Token Contexts Under Fixed-GPU Constraints

Introduction

The paper "LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget" (2607.14952) systematically addresses the challenge of reinforcement learning (RL) post-training for LLMs with million-token contexts, constrained by a fixed number of GPUs. Existing deployment-level LLMs routinely perform inference over extremely long sequences (1M+ tokens), but RL post-training workloads lag, typically capped at 256k tokens. This gap is problematic for agentic applications where the accumulated trajectory history—tools, observations, documents, and prior actions—routinely exceeds practical RL training sequence limits. Training is bottlenecked by quadratic attention computation, massive prompt state retention for backpropagation (backward paths), and exponential GPU memory usage when scoring and differentiating over multiple group responses sharing common context.

The key insight of LongStraw is the explicit formalization and optimization of tensor lifetime and physical buffer ownership, decoupled from mere algorithmic sparsity or attention kernel innovations. LongStraw demonstrates post-training RL updates on two modern LLM architectures with context lengths up to 2.1M tokens using only 8 or 32 H20 GPUs, using architectural replay scheduling and targeted state residency, rather than the brute-force hardware scaling of recent long-context works.

Methodology: Prompt-State Residency and Suffix Replay

Group Relative Policy Optimization (GRPO) Dependency Graph

LongStraw works with the GRPO objective [shao2024deepseekmath], which normalizes reward within a group of responses sampled from the same prompt context. The core technical challenge is that, unlike inference, GRPO backward must retain and differentiate the entire prompt and every response path, making the context group size ($G$) and prompt length ($P$) scale linearly with the number of live activations.

LongStraw introduces a boundary in the computation graph:

• The prompt is executed once with gradients disabled, capturing only the minimal model-specific state required for subsequent suffix replay.
• For each group response (member), only the short suffix (response tokens: $R_i$) is replayed under autograd. Gradients are accumulated locally, with a single optimizer update after all member backward passes.
• This transition reduces the live autograd graph size from $O(G \cdot (P+R_i))$ to $O(\max_i R_i)$. There is no attempt to recompute the gradient with respect to the prompt state ($\partial z_P /\partial\theta$); the update is thus a conditional-response gradient, not a fully correct sequence gradient.

LongStraw serializes group members, managing time/memory trade-offs for extremely large contexts without needing additional devices.

Architecture-Specific State Management

• Qwen3.6-27B: Hybrid recurrent and full-attention stack. The recurrent Gated DeltaNet (GDN) layers carry fixed-size recurrent state across the prompt boundary, while the 16 full-attention layers retain compact, CP-sharded key/value (KV) pages. All storage is physically allocated on GPU, requiring careful page table and memory management.
• GLM-5.2: 78-layer MoE Transformer with Multi-Head Latent Attention (MLA) and Dynamic Sparse Attention (DSA). Prompt state for MLA and DSA keys is offloaded to CPU RAM, with only one layer staged at a time during replay. Index reuse (IndexShare) is implemented, with DSA selection and expert routing performed per local context shard.

Core technical innovations include:

• Detachment and right-sized allocation of context-sharded pages, ensuring that allocator fragmentation or parent views do not block memory release.
• Response suffix replay with whole-layer checkpointing (especially for MoE tails), suppressing activation retention of the long prefix.
• Explicit device and ownership management, with clear separation of context parallelism (CP—token/history axis) and expert parallelism (EP—MoE parameters/routing), necessary for GLM's large parameter footprint.

Empirical Results

Qwen3.6-27B Results (8 H20 GPUs, CP8)

• Achieved grouped post-training with context length 2,097,152 tokens (2.1M): 2,088,960 prompt + 8,192 response.
• Wall-time: 5,198s ($G=2$), 6,785s ($G=8$); peak allocated memory grew only 0.2 GB (0.21%) when $G$ increased from 2 to 8, verifying live autograd scaling with suffix length, not group size.
• Succeeded in processing up to 4.25M tokens (4,456,448) in a further scaling test, before running out of memory at 4,542,464 tokens.
• Global partitioned attention forward: all CP8 shards contribute via a numerically stable log-sum-exp (LSE) merge using AllReduce operations. Caveat: only $dQ$ is AllReduced; $P$0 (KV gradients) and LoRA adapter gradients are not synchronized, resulting in incoherent optimizer steps across ranks.

GLM-5.2 Results (32 H20 GPUs, TP1/CP32/EP32)

• Achieved grouped execution (two response members, $P$1) with prompt length 2,097,152 tokens.
• Each of the 32 ranks processed 65,536 tokens per local context shard and routed expert assignments.
• CPU-resident prompt-state storage: $P$25.8GB per rank.
• Layer-wise suffix replay through 78 layers matches architectural state, with full MoE computation. Index reuse and per-forward IndexShare handling are implemented.
• Limitations: The reported run realizes only local DSA—sparse attention and selection are performed per CP-shard, no cross-rank candidate merging or selected-value exchanges. Megatron gradient finalize is bypassed, leaving distributed gradient reduction and coherent model updates incomplete.

Systems Lessons and Theoretical Implications

• Physical state lifetime is the critical system bottleneck, not attention sparsity or parameter-efficient design alone. Allowing prompt activations (FFN, MoE, router, attention scratch) to die immediately after capture, with only necessary conditional state (e.g., compact KV or MLA/DSA latent pages) retained, is key to supporting million-token updates.
• Parallelism axes (CP, EP) are orthogonal and must be handled distinctly: CP divides the token sequence for attention state, EP divides the expert set for FFN computation. Mismanaging their composition leads to silent failures in distributed gradient aggregation.
• Sparsity moves the bottleneck, does not remove it: DSA shifts cost from dense computation to candidate selection and value handling, MoE shifts from parameter count to buffer residency and dispatch/combination communication overheads.
• Gradient parity and correct distributed update are distinct: Execution capacity, correct forward operator, synchronized distributed update, and full-sequence gradient equivalence must each be separately validated. The presented receipts achieve only the first and partial second level.

Practical Considerations and Future Developments

LongStraw's strategy of separating prompt evaluation from response computation and serializing response replay enables long-context RL post-training at context lengths previously requiring 10-100x more hardware (Liu et al., 2023, Jacobs et al., 2023, Ge et al., 28 Feb 2025, Fang et al., 2024). The open limitations are:

• Missing distributed gradient synchronization (especially for CP-replicated LoRA adapters).
• Local-only DSA sparse attention for GLM-5.2; missing candidate merge and value exchange for exact global sparse attention in the context parallel regime.
• Detached prefix state implies that gradient computation omits components with respect to changes in the prompt.
• Evaluation only with synthetic data, not end-to-end RL loops or real reward models.

Future work should prioritize:

1. Implementing and verifying selective distributed gradient reduction for all replicated adapter families (experimentally required for full semantic correctness).
2. Restoring proper global DSA operator fidelity via cross-rank candidate merging and atomic value exchanges.
3. Establishing full-sequence gradient parity and optimizer delta correctness via matched, short-context reference tests.
4. Extending to repeated updates, online rollout and checkpointing to demonstrate practical RL performance at million-token contexts.

These advancements will be crucial for enabling research teams with modest accelerator resources to conduct effective RL for truly long-context LLMs, democratizing agentic AI RL studies.

Conclusion

LongStraw establishes that million-token RL post-training can be achieved without hardware overprovisioning, provided that model- and architecture-specific state boundary management, physical buffer residency, and lifetime control are explicitly targeted. For Qwen, prompt-state compaction, CP-sharded AllReduce for attention, and response-graph serialization enable scaling to 4.25M tokens on 8 H20 GPUs. For GLM, CPU-resident page management, one-layer checkpointing, and replay isolation make 2M-token group updating feasible on 32 H20 GPUs. The adopted approach exposes the strict systems requirements for distributed RL updates and separates execution and update-path engineering from pure model algorithmics. These findings reframe long-context RL as primarily an execution and physical-ownership challenge under fixed-resource regimes, and chart a clear path for scaling RL for LLMs to realistic trajectory lengths under resource constraints.

---

Reference:

Changhai Zhou et al., "LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget" (2607.14952)