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

🕒 Published (v1): 2026-07-16 13:00 UTC · Source: Arxiv · link

Why this paper was selected

Closes RL/inference context gap (2M tokens); novel approach for long-horizon agent training

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LongStraw is an architecture-aware execution stack that enables GRPO-based RL post-training at million-token contexts on a fixed GPU budget by evaluating the shared prompt once without autograd, retaining only architecture-specific conditional state, and replaying short response branches serially under autograd. It is instantiated for two model families—Qwen3.6-27B (hybrid recurrent/full-attention) and GLM-5.2 (compressed-attention MoE)—and demonstrates 2.1M-token grouped execution on 8 H20 GPUs for Qwen and 32 H20 GPUs for GLM, with a stress-test envelope reaching 4.46M positions.

Problem

A growing gap exists between inference-time context lengths (approaching millions of tokens) and RL post-training context lengths (typically ≤256K). GRPO post-training is memory-intensive because it must score and backpropagate through multiple responses conditioned on the same long prompt: the live training graph must span \(P + \sum_i R_i\) tokens simultaneously, causing quadratic attention memory and long-lived backward states to exceed fixed GPU budgets. Scale-out approaches (Ring Attention, DeepSpeed-Ulysses, ByteScale) use hundreds to thousands of GPUs; LongStraw asks how far a fixed-accelerator budget can reach instead.

Method

LongStraw separates the shared prompt computation from response computation via a capture-once, replay-suffix schedule:

1. Prompt capture (no-grad): Evaluate the full \(P\)-token prompt once without autograd. Retain only the architecture-specific conditional state \(\bar{z}_P = \text{stopgrad}(z_P(\theta))\) needed to condition response tokens: for Qwen, compact GDN recurrent state + CP-sharded KV pages; for GLM, CPU-offloaded MLA latent pages + DSA index-key pages.
2. Pre-step scoring: Compute old-policy and reference-policy log-probabilities from the frozen state without changing parameters.
3. Serial response replay (under autograd): Rebuild one response graph at a time (length \(R_i\)), backpropagate, accumulate gradients into the shared adapter, free the graph, then proceed to member \(i+1\).
4. Single optimizer call: After all \(G\) members, synchronize and step once.

This bounds live autograd memory to a single response branch rather than \(P + GR\). The GRPO objective is the standard clipped surrogate with group-relative advantages: $\(L_{\text{policy}} = -\frac{1}{G}\sum_{i=1}^G \frac{1}{R_i}\sum_{t=1}^{R_i} \min(\rho_{i,t} A_i,\ \text{clip}(\rho_{i,t}, 1-\epsilon, 1+\epsilon) A_i)\)$

Critically, the approach computes only \(\partial\ell/\partial\theta|_{z_P}\) and omits \((\partial\ell/\partial z_P)(\partial z_P/\partial\theta)\)—the gradient through the prompt state is detached by design.

Qwen implementation (8 H20s, CP8): QLoRA (NF4) with 116M trainable parameters; 48 GDN layers contribute compact fixed-shape recurrent state; 16 full-attention layers contribute CP-sharded KV pages. Response replay uses a global CP8 LSE/output merge for full-attention layers. Dense FFN intermediates from the prompt are discarded immediately.

GLM implementation (32 H20s, CP32/EP32): MLA latent pages and 21 indexer-key pages (for index-computing layers; 57 IndexShare layers reuse their selections) are offloaded to CPU. Whole-layer checkpointing is applied across all 78 response decoder layers. EP32 and CP32 are folded onto the same 32 GPUs to reduce process groups.

Key Contributions

• Capture-once/replay-suffix execution design for GRPO that reduces the live training graph from the full \(P+R\) sequence to a single response branch \(R_i\), trading time for memory.
• Dual implementation covering two substantially different model architectures: hybrid recurrent/full-attention dense (Qwen3.6-27B) and compressed-attention MoE (GLM-5.2).
• Formal taxonomy of evidence levels distinguishing execution capacity, response-operator fidelity, distributed-update consistency, and full-gradient parity—clarifying exactly what the current system does and does not prove.
• Architecture-aware state inventory identifying which tensors must survive the prompt boundary (GDN state, KV pages, MLA latents, DSA index keys) and which can be immediately freed (FFN intermediates, MoE route buffers, attention scratch).

Results

• Qwen on 8 H20 GPUs: Completes grouped scoring and response backward at 2.1M positions with group sizes \(G \in \{2, 8\}\); increasing \(G\) from 2 to 8 adds only 0.21 GB of peak allocated memory—confirming group cardinality is primarily a scheduling/time dimension, not a memory dimension.
• Qwen stress test on 8 H20 GPUs: Execution envelope extended to 4.46M positions (4.25M context + response).
• GLM on 32 H20 GPUs: End-to-end LongStraw execution path validated for a 2.1M-token prompt across all 78 layers of GLM-5.2.
• Comparison baseline (scale-out): Ring Attention achieves 4.096M positions for a 7B model on 32 A100s; DeepSpeed-Ulysses reaches 1M tokens on 256 A100s; ByteScale reaches 2M on 1,024 GPUs. LongStraw reaches 2.1M on 8–32 H20s by trading time for devices.
• These are execution receipts, not convergence results; training correctness (full distributed gradient parity, repeated-loop stability) is not yet demonstrated.

Limitations

• Detached prompt gradient: The gradient \((\partial\ell/\partial z_P)(\partial z_P/\partial\theta)\) is not computed; only the conditional response gradient is used. This is an approximation whose effect on learned policy quality is unquantified.
• Incomplete distributed gradient composition: Qwen does not all-reduce \(dK/dV\) contributions to replicated projection adapters; the eight AdamW instances step independently from unreduced gradients. GLM skips finalize_model_grads and CP-replicated non-expert adapters step from unreduced gradients.
• GLM DSA fallback is CP-shard local: The accepted GLM receipt uses each rank's 65,536-token shard for sparse index selection rather than global full-context DSA, so response-operator fidelity is not claimed for GLM.
• Execution probe only: Runs use supplied deterministic responses and rewards; online generation and reward-model execution are excluded. No repeated-loop evidence; prompt state becomes stale after one optimizer step.
• No full-gradient parity test: Neither implementation has been compared parameter-by-parameter against a conventional full-sequence reference at shorter context.
• Resource and reproducibility gaps: Hardware configuration details limit external reproducibility.

Relevance to Agentic AI / LLM Agents

Long-context RL post-training is directly relevant to training AI agents, whose inputs are inherently long-form—combining tool outputs, environment observations, document contents, and multi-turn interaction histories into single trajectories. LongStraw lowers the hardware barrier for GRPO fine-tuning at the context lengths agents actually require at inference, narrowing the gap between training and deployment distribution. By demonstrating that group size affects memory only marginally (0.21 GB for \(G: 2 \to 8\)), the system enables richer reward-contrast signals without proportional GPU scaling—critical for RLHF and RLAIF pipelines for agent behavior. The explicit acknowledgment of the missing prompt-gradient term and incomplete distributed updates frames a concrete open problem for the community: correct long-context policy gradient estimation without requiring prohibitively large GPU clusters.