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

> LongStraw enables million-token GRPO on fixed GPUs by separating prompt from response computation, achieving over 2M context positions.

- **Source:** [arXiv](https://arxiv.org/abs/2607.14952)
- **Published:** 2026-07-18
- **Permalink:** https://picx.dev/p/745msU
- **Whiteboard:** https://picx.dev/p/745msU/image

## Summary

# 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**:

| State component | Placement & ownership | Scaling with $P$ |
|----------------|----------------------|------------------|
| **Qwen GDN recurrent state** | GPU per rank (CP8) | Fixed size (independent of $P$) |
| **Qwen full-attention KV pages** | Compact GPU pages across CP8 for 16 layers | $O(P/8)$ per rank |
| **GLM MLA latent pages** | CPU; Megatron-zigzag CP32 shards for 78 layers | $O(P/32)$ per rank |
| **GLM DSA indexer-key pages** | CPU; CP32 shards for 21 index-computing layers | $O(P/32)$ at those layers |

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

| Path | Layout | Response-time collectives | Audited semantic boundary |
|------|--------|---------------------------|---------------------------|
| Qwen3.6-27B | 8 H20, CP8 | Global CP8 max/normalizer/value-sum merge (BF16 numerator) | Global response-attention forward preserved; dQ all-reduced; dK/dV adapter gradients not synchronized |
| GLM-5.2 | 32 H20, TP1/CP32/EP32/ETP1/PP1 | EP32 dispatch/combine; DSA CP-local (no global candidate merge) | All ranks complete 78-layer backwards; custom path skips Megatron gradient finalization for CP-replicated non-expert adapters |

### 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.

| Path | Prompt / response | $G$ | Hardware/layout | Terminal evidence on every worker | Wall (s) | Peak (GB) |
|------|------------------|-----|-----------------|-----------------------------------|----------|-----------|
| Qwen global forward | 2,088,960 / 8,192 input ×2 | 2 | 8 H20, CP8 | Old+ref+backward+local AdamW call | 5,198.780 | 97.503 |
| Qwen global forward | 2,088,960 / 8,192 input ×8 | 8 | 8 H20, CP8 | 8 serial members + local AdamW call | 6,785.225 | 97.711 |
| GLM CP-local DSA | 2,097,152 / 3 input (2 scored) | 2 | 32 H20, TP1/CP32/EP32 | Two old scores + 2×78-layer backward + optimizer call | 2,975.138 | n/r |

### 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.

---

_Markdown view of https://picx.dev/p/745msU, served by PicX — AI-generated visual whiteboard summaries of research papers._