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

Summary (Overview)

• LongStraw enables million-token GRPO post-training on a fixed GPU budget (8×H20 for Qwen, 32×H20 for GLM) by separating the long prompt from the short response computation.
• The shared prompt is evaluated once without automatic differentiation; only the architecture-specific state (recurrent/attention/MoE) needed by later tokens is retained. Response branches are replayed serially one at a time under autograd, reducing the live training graph from the full sequence to a single response branch.
• Two substantially different model families are implemented: the hybrid recurrent + full-attention Qwen3.6-27B and the compressed-attention Mixture-of-Experts GLM-5.2.
• On eight H20 GPUs, Qwen completes grouped scoring and response backward at 2.1M context positions with group sizes 2 and 8; a stress test extends the envelope to 4.46M positions. On 32 H20 GPUs, GLM validates a 2.1M-token prompt across all 78 layers.
• The key finding is that state lifetime and physical ownership are the primary determinants of the practical context limit, rather than attention sparsity or parameter count alone.

Introduction and Theoretical Foundation

AI agents operate over long trajectories (ReAct, LongCat-Flash-Thinking), creating a growing gap between inference-time context lengths (approaching millions) and post-training lengths (often ≤256K tokens). While inference can cache prompt state and discard computational graphs, RL post-training must score and backpropagate through multiple responses conditioned on the same shared prompt - a much heavier memory burden.

Group Relative Policy Optimization (GRPO) compares responses within a group through normalized advantages. For a prompt of $P$ tokens, group member $i$ with $R_i$ scored tokens, and $G$ responses, the policy objective is:

$$L_{\text{policy}} = -\frac{1}{G}\sum_{i=1}^{G}\frac{1}{R_i}\sum_{t=1}^{R_i}\min\big(\rho_{i,t}A_i,\;\text{clip}(\rho_{i,t},1-\epsilon,1+\epsilon)A_i\big)$$

where $\rho_{i,t}(\theta)=\exp\big(\log\pi_\theta(y_{i,t}|x_{1:P},y_{i,<t})-\log\pi_{\text{old}}(y_{i,t}|x_{1:P},y_{i,<t})\big)$ is the importance ratio, and $A_i$ is the normalized group-relative advantage.

Existing techniques (FlashAttention, LoRA, QLoRA, activation checkpointing) reduce individual costs but leave the live autograd graph spanning the full prompt and all responses competing for device memory. Scale-out approaches (Ring Attention, DeepSpeed-Ulysses, ByteScale) distribute work across many GPUs; LongStraw asks the complementary question: how far can GRPO go when the accelerator count stays fixed?

Methodology

LongStraw Design

LongStraw separates the long prompt from short response computation through a four-phase transaction:

1. Prompt capture: Evaluate $x_{1:P}$ under current parameters with autograd disabled. Save only the model-specific conditional state needed by future tokens; release all transient hidden states, attention scratch, FFN activations, and routing buffers.
2. Pre-step scoring: Materialize old-policy and reference-policy response log-probabilities before any policy backward, with parameters fixed through the group.
3. Policy replay: For each group member, rebuild only the short response path under autograd, backpropagate, and immediately free that member's graph. Reuse the same read-only prompt state.
4. Optimizer transaction: Accumulate local gradients over all $G$ members, synchronize, and issue one optimizer call per worker.

The dominant activation scale changes from $P + R$ to $R$, with total resource accounting:

$$M_{\text{live}} \approx M_{\text{fixed}} + M_{\text{prompt}}(P) + M_{\text{grad}} + \max_i M_{\text{branch}}(R_i) + M_{\text{score}}\left(\sum_i R_i\right)$$ $$T_{\text{update}} = T_{\text{prompt}}(P) + \sum_{i=1}^G T_{\text{score+replay}}(R_i)$$

This bounds the live policy graph by the largest member rather than by the number of members. The critical trade-off: serial replay increases elapsed time but avoids keeping the prompt graph and every response graph live simultaneously.

State Inventory and Ownership

The captured prompt state is architecture-specific:

For Qwen at 2,088,960 prompt tokens with CP8, page size $S=64$:

$$N_{\text{page}} = \frac{2,088,960}{64} = 32,640,\quad N_{\text{page/rank}} = \frac{32,640}{8} = 4,080$$

BF16 KV storage per rank across 16 full-attention layers:

$$B_{\text{KV/rank}} = 16 \times 4,\!080 \times 2 \times 64 \times 4 \times 256 \times 2\;\text{bytes} \approx 15.94\;\text{GiB}$$

Parallel Layout and Key Operations

Global Full-Attention Forward Composition (Qwen)

For response query $t$, rank $r$ returns local output $o_{r,t}$ and log-normalizer $\ell_{r,t}$. Global output is reconstructed without moving KV pages:

$$m_t = \max_r \ell_{r,t},\quad a_t = \sum_r e^{\ell_{r,t} - m_t},\quad n_t = \sum_r e^{\ell_{r,t} - m_t} o_{r,t}$$

Then $o_t = n_t / a_t$, $\ell_t = m_t + \log a_t$. This requires one MAX all-reduce and two SUM all-reduces per composed output.

Empirical Validation / Results

Audited Execution Receipts

Table 5: Fixed-budget execution receipts on 8 H20s (Qwen) and 32 H20s (GLM), 2.1M context positions.

Qwen Group Scaling

• $G=2$ to $G=8$: peak memory increases only 0.208 GB (0.213%); wall time grows 1,586.445 s. Post-prefix cost per member remains ∼266–271 s.
• The shared prefix (∼4,656 s) dominates total time; amortizing it over more members reduces mean wall time per response from 2,599.39 to 848.15 s.
• 4.25M stress test: 4,456,448 positions (4,448,256 prompt + 8,192 response) complete G=8 response replay at 82.960 GB peak. Under a prefix-frozen objective, 8 consecutive G=8 optimizer steps (64 member replays) run at 83.894 GB peak.

GLM Terminal Manifest

• 32/32 ranks report complete: 2 old score phases, 2 live 78-layer policy backwards, 1 optimizer call per worker.
• Trace inventory: 128 JSONL files, 35,584 events. Policy traces record 78 forward and 78 backward layer-end events per rank.
• CPU prompt state: 5.8125 GiB per rank (72 MiB per MLA layer + 16 MiB per index-key set). Per-layer GPU staging: 72–88 MiB.
• Gradient composition gap: custom path bypasses finalize_model_grads, so CP-replicated non-expert adapter gradients remain unreduced before the distributed optimizer call.

Memory Accounting and Evidence Levels

The paper defines four levels of evidence:

1. Execution capacity: All requested backwards, collectives, optimizer events complete (✓ both paths).
2. Response-operator fidelity: Qwen CP8 global merge ✓; GLM DSA CP-local only (×).
3. Distributed-update consistency: Qwen dK/dV reduction missing; GLM gradient finalization skipped (×).
4. Full-gradient parity: Prompt state detached in both; no numerical reference exists (×).

Theoretical and Practical Implications

Key systems lessons:

• Fixed-budget capacity comes from lifetime management, not sparsity alone. Both models evaluate every prompt token through every layer; the gain comes from allowing prompt scratch, intermediates, and routing buffers to die before response backward. DSA reduces attention arithmetic but retains long-context indexing; MoE activates only selected experts but creates dispatch/combine overhead.
• Physical ownership must match logical ownership. A context shard is only useful when its physical allocation is also sharded. Copying selected pages into right-sized tensors (rather than retaining view parent allocations) is critical.
• Dense and MoE architectures move the peak to different places. For Qwen, context-growing storage is concentrated in 16 full-attention KV sets (48 GDN layers remain fixed-size). For GLM, full-sequence storage is dominated by expanded routed hidden buffers (6 GiB for one balanced BF16 buffer at 65K-token shard), not by attention scratch.
• Context parallelism and expert parallelism are orthogonal. CP partitions token history; EP partitions model parameters. CP collectives cannot replace EP collectives and vice versa.
• Group scaling is a scheduling result. Serial replay makes group size primarily a time dimension, not a memory dimension. $G=2$ vs $G=8$ changes peak memory by only 0.2%.
• Forward fidelity, update consistency, and gradient parity are distinct claims. The paper achieves level 1 (execution capacity) for both models, level 2 for Qwen only, and neither reaches levels 3 or 4.

Practical implications: By reducing reliance on ever-larger GPU clusters, LongStraw enables more researchers and smaller teams to explore long-context training under limited accelerator budgets. The explicit identification of missing gradient composition paths provides a concrete roadmap for reaching full training correctness.

Conclusion

LongStraw demonstrates that long-context GRPO under a fixed GPU budget is a tensor-lifetime and ownership problem, not a context-length race or a single-kernel optimization. By evaluating the shared prompt once without autograd, retaining only architecture-specific state, and replaying response branches serially, LongStraw composes million-token GRPO-shaped execution on 8 H20 GPUs for Qwen and 32 H20 GPUs for GLM – a four- to eightfold increase over the native context settings of these models.

The current results establish execution capacity, not full training correctness. The prompt state is detached from the backward pass, and both implementations have incomplete distributed gradient composition:

• Qwen: shard-local dK/dV contributions to replicated KV projection adapters are not synchronized across CP8 ranks.
• GLM: the custom replay path bypasses Megatron gradient finalization, leaving CP-replicated non-expert adapter gradients unreduced.

Closing these gaps – restoring gradient finalization, validating global DSA candidate/output composition, and running full-gradient parity comparisons – is the necessary next step before complete online RL training. The paper provides a clear dependency-ordered validation roadmap toward that goal.