🤖 machine learning

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

LongStraw is an architecture-aware execution stack that enables million-token reinforcement learning post-training under a fixed GPU budget by optimizing memory usage through shared prompt evaluation without autograd and sequential response replay, successfully demonstrating execution capacity on models like Qwen3.6-27B and GLM-5.2 with context lengths exceeding 2 million tokens.

Original authors: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Che

Published 2026-07-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to solve a massive, multi-step mystery. To do this, the robot needs to read a library's worth of clues (the "context") before it can even start guessing the solution. In the world of Artificial Intelligence, this is called "long-context training." Usually, when we train these robots, we feed them a huge book of clues and then ask them to write a short answer. The problem is that the robot's brain (its computer memory) gets so full of the book's details that it runs out of space before it can even finish writing the answer. It's like trying to solve a mystery while holding the entire library in your hands; eventually, your arms get too tired to move.

To fix this, scientists have been building bigger and bigger libraries of computer chips (GPUs) to hold all the information. But what if you don't have a supercomputer? What if you only have a small, fixed number of chips, like a standard set of eight or thirty-two? This is the puzzle the paper "LongStraw" tackles. It asks: Can we teach a robot to read a two-million-word story and write a solution using only a small, fixed set of computer chips, without buying more? The answer isn't "yes, we made the chips bigger," but rather "yes, we changed how we hold the book."

The researchers built a clever system called LongStraw. Think of it like a librarian who is incredibly organized. Instead of trying to keep the entire two-million-word story open on the desk while the robot writes its answer (which would crash the desk), the librarian reads the whole story once, memorizes just the essential "vibe" or "state" of the story, and then puts the heavy book away. Now, the robot only has to write its short answer while looking at that tiny "vibe" note. If the robot needs to check a specific detail, the librarian quickly pulls out just that one page, shows it to the robot, and then immediately puts it back. This way, the desk never gets cluttered, even though the story is huge.

The paper tests this idea on two different types of robot brains: one called Qwen and another called GLM. They tried to make these robots read a prompt (the story) that was 2,097,152 tokens long (which is roughly two million words or characters) and then write a response. They did this on a strict budget: eight H20 GPUs for the Qwen robot and 32 H20 GPUs for the GLM robot.

Here is what they found:

It works for the "reading" part: They successfully managed to process the entire two-million-token story without the computer memory crashing. They proved that you can fit this massive task onto a fixed number of chips by being smart about when you hold onto information and when you let it go.
It saves space, but takes time: By putting the heavy book away and only pulling out pages when needed, they saved a huge amount of memory. For example, when they increased the number of robot "guesses" (the group size) from 2 to 8, the memory usage barely changed (only about 0.21 GB more), but it took longer to finish the job because they had to pull the pages out one by one instead of all at once.
It's a "proof of concept," not a perfect finish: The paper is very honest about what it didn't do. While the system successfully ran the simulation and showed the robot could read the long story and write the answer, it didn't fully prove that the robot learned the perfect way to update its brain. Some of the math that connects the robot's different "brain parts" (gradients) wasn't fully synchronized across all the chips in the way a perfect training session would require. It's like showing that a car engine can run on a tiny amount of fuel, but not yet proving it can drive a whole race track without stalling.

The researchers call their result an "execution receipt." It's a receipt that says, "Yes, we successfully ran this massive task on this specific hardware without crashing." They didn't claim to have won a race or built the fastest car; they just proved that the car can drive on this specific track with this specific amount of gas.

In short, LongStraw shows that you don't need a supercomputer to train AI on massive stories. You just need a smarter way to manage your memory, trading a little bit of extra time for a lot of saved space. It opens the door for smaller teams and researchers with limited budgets to experiment with these giant AI models, proving that you can do big things with a fixed, small toolkit if you're clever enough.

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

1. Problem Statement

A significant gap has emerged between the context lengths supported by inference systems (approaching million-token contexts) and those utilized during Reinforcement Learning (RL) post-training, which often remain capped at 256K tokens or lower. This gap is particularly critical for AI agents whose trajectories accumulate observations, tool outputs, and prior decisions over long sequences.

The core challenge in extending RL post-training (specifically Group Relative Policy Optimization, or GRPO) to million-token contexts under a fixed GPU budget is memory. Unlike inference, where the forward graph can be discarded after state caching, training must score multiple responses and backpropagate through them. Standard approaches require retaining the full prompt and all response graphs in memory simultaneously, leading to quadratic attention computation and massive activation storage. Existing techniques like FlashAttention, LoRA, or Ring Attention address specific bottlenecks but do not inherently allow a fixed-GPU GRPO run to fit a 2M+ token context without scaling out to larger clusters.

2. Methodology: LongStraw Execution Stack

LongStraw is an architecture-aware execution stack designed to decouple the long prompt from the short response computation, trading additional replay time for reduced GPU memory usage. The core design principle is to retain a tensor across the prompt boundary only when a later response token depends on it.

2.1 The Execution Schedule

The system operates in four distinct phases for a grouped update:

Prompt Capture (No-Grad): The shared prompt ( $P$ tokens) is evaluated once with automatic differentiation disabled. The system saves only the architecture-specific conditional state required for future tokens (e.g., recurrent states, KV pages, latent representations) and releases all transient activations (FFN intermediates, attention scratch, MoE routing buffers).
Pre-Step Scoring: Old-policy and reference-policy log-probabilities are evaluated for all group members using the frozen prompt state.
Serial Response Replay: Instead of maintaining $G$ response graphs simultaneously, the system replays one response branch at a time under autograd. It reuses the read-only captured prompt state, computes the loss, backpropagates, and immediately frees the response graph.
Gradient Accumulation & Optimization: Local gradients from all $G$ members are accumulated, followed by a single optimizer call per worker.

2.2 Architecture-Specific Implementations

LongStraw adapts this strategy to two distinct model families:

Qwen3.6-27B (Hybrid Recurrent/Full-Attention):
- State Retention: Keeps compact Gated DeltaNet (GDN) recurrent states and context-parallel (CP) sharded Key/Value (KV) pages for the 16 full-attention layers on GPU.
- Mechanism: Uses NF4 QLoRA. The 48 GDN layers provide fixed-size state, while the 16 full-attention layers use physical page compaction. The system performs a global LSE/output merge for the forward pass but acknowledges missing K/V adapter gradient synchronization.
- Hardware: Executed on 8 H20 GPUs (CP8).
GLM-5.2 (Compressed Attention/MoE):
- State Retention: Moves the long prompt state (MLA latent pages and DSA indexer-key pages) to CPU memory. Only one decoder layer is staged to the GPU at a time during response replay.
- Mechanism: Uses a "Zigzag" context page layout across 32 ranks (CP32) and Expert Parallelism (EP32). It employs whole-layer activation checkpointing to bound the response graph. The system reconstructs IndexShare schedules layer-by-layer.
- Hardware: Executed on 32 H20 GPUs (TP1/CP32/EP32).

3. Key Contributions

Shared-Prompt Execution Design: A novel schedule that stores prompt state once (detached from autograd) and serializes response work, reducing the live training graph from $P + \sum R_i$ to $\max(R_i)$ .
Dual-Architecture Implementation: Successful instantiation of this design for two structurally different models: a dense hybrid recurrent/attention stack (Qwen) and a compressed-attention MoE stack (GLM).
Fixed-Budget Execution Receipts: Empirical evidence demonstrating that GRPO execution paths can complete for prompts exceeding 2M tokens (and up to 4.25M for Qwen) without increasing the GPU count, while explicitly documenting the limitations in distributed gradient composition.

4. Experimental Results

4.1 Qwen3.6-27B (8 H20s)

Context Length: Successfully executed a 2,097,152-token context (2,088,960 prompt + 8,192 response).
Memory Efficiency: Increasing the group size from $G=2$ to $G=8$ increased peak allocated memory by only 0.21 GB (0.213%), confirming that memory scales with the largest response branch, not the group size.
Extended Capacity: A stress test reached 4.25M positions (4,448,256 prompt + 8,192 response) on the same 8 H20s, peaking at ~83 GB per rank.
Limitation: The distributed update is not fully coherent; K/V adapter gradients are not synchronized across ranks, and the prompt state gradient is detached.

4.2 GLM-5.2 (32 H20s)

Context Length: Validated the end-to-end execution path for a 2,097,152-token prompt across all 78 layers.
Memory Strategy: CPU-resident prompt state (approx. 5.8 GiB per rank) allowed the GPU to remain within budget during the 78-layer backward pass.
Limitation: The DSA (Dynamic Sparse Attention) operator in the accepted receipt is local to each CP shard (selecting top-2048 from 65k tokens rather than the global 2M context). Additionally, the path bypasses Megatron's finalize_model_grads, leaving CP-replicated non-expert adapter gradients unreduced before the optimizer step.

5. Significance and Claims

The paper positions LongStraw not as a "world record" for context length, but as a fixed-budget feasibility envelope.

State Lifetime vs. Sparsity: The authors argue that fixed-budget capacity is determined by tensor lifetime and physical ownership, not merely by sparsity (MoE) or memory-efficient attention kernels. By allowing prompt intermediates to die immediately after capture, the system avoids the memory explosion of full-sequence autograd.
Physical Ownership: The work emphasizes that logical sharding is insufficient; physical memory allocation must match the sharding strategy (e.g., copying pages to right-sized tensors) to release memory effectively.
Modest Claims: The authors explicitly state that these are execution receipts, not proofs of training correctness. The prompt state is detached (missing the $\partial z_P / \partial \theta$ term), and distributed gradient composition is incomplete for both models. The results establish that the execution path fits within the hardware constraints, enabling researchers to explore long-context RL without massive scale-out, but they do not claim to have solved the distributed gradient parity or policy improvement aspects yet.

In summary, LongStraw demonstrates that with careful management of state lifetime and physical ownership, million-token RL post-training is achievable on fixed, relatively small GPU clusters, shifting the bottleneck from raw memory capacity to scheduling and state management.