AIAny
Icon for item

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

Enables RL post-training with million-token prompts under a fixed GPU budget by evaluating shared prompt state without autograd, retaining only minimal model state, and replaying short response branches; instantiated as GRPO and demonstrated on Qwen3.6-27B and GLM-5.2 up to multi-million token execution.

Introduction

Most RL post-training pipelines still assume prompt lengths ≤256K, while modern inference systems and agent trajectories routinely approach or exceed millions of tokens. That gap matters for agent-style workloads where observations, tool outputs and prior decisions accumulate over very long trajectories. LongStraw targets that execution gap: instead of redesigning model internals, it trades extra replay time for a much smaller live training graph so million-token RL post-training can run on a fixed GPU budget.

Key Findings
  • Architecture-aware execution stack: LongStraw evaluates the shared prompt without autograd, keeps only the per-model state required for later tokens, and replays response branches one at a time to reduce live memory and graph size. This shifts cost from peak memory to additional replay compute.
  • Empirical capacity limits: On eight H20 GPUs, LongStraw completes grouped Qwen scoring and response backward at 2.1M positions for groups of 2 and 8; increasing group size added only ~0.21 GB peak memory. A separate stress test reached 4.46M positions. On 32 H20 GPUs, the end-to-end execution path was validated for a 2.1M-token prompt across all 78 layers of GLM-5.2.
  • Instantiation and scope: The paper instantiates LongStraw with Group Relative Policy Optimization (GRPO) and implements it for hybrid recurrent/full-attention Qwen3.6-27B and compressed-attention mixture-of-experts GLM-5.2. The evaluations emphasize execution capacity (what can be run) rather than full training correctness.
Who It's For and Tradeoffs

Great fit if you need to validate whether million-token RL post-training execution is feasible on a fixed GPU fleet, or to prototype long-context agent training pipelines without redesigning model internals. LongStraw is useful for stress-testing execution, measuring memory scaling vs. replay cost, and integrating with large LLMs that mix recurrent and attention components. Look elsewhere if you require fully correct end-to-end distributed gradient composition with no detached prompt state, minimal replay overhead, or production-ready wall-clock training speed—some distributed forward/gradient composition paths remain incomplete and the prompt state is detached, so the work demonstrates capacity more than finished training correctness.

How it works (brief)

The core mechanism detaches the shared prompt from autograd, reconstructs only model-specific state needed by future tokens, and replays short response branches sequentially so each branch builds gradients locally while keeping the global live graph small. The approach reduces peak memory growth per group at the cost of extra compute for replay; measured memory overheads and multi-million-token checkpoints quantify that tradeoff.

Information

  • Websitearxiv.org
  • AuthorsChanghai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei
  • Published date2026/07/16

More Items

GitHub
AI Infra2026

Defines OpenTelemetry semantic conventions for generative AI telemetry — spans, metrics, and events for GenAI clients, the Model Context Protocol (MCP), and provider-specific integrations. Includes YAML models, human-readable docs, and reference implementations to standardize observability across GenAI deployments.

GitHub
AI Infra2024

Provides a lightweight build platform for HIP and ROCm that supports building ROCm, PyTorch, and JAX from source, multi-architecture nightly releases, and integrated CI/CD and developer tooling for Linux and Windows.

Analyzes adversarial weaknesses of World-Action Models (WAMs) via BadWAM, a framework that crafts visual perturbations to decouple a model’s imagined future from its executed actions. Introduces two attack modes—action-only (disruptive) and imagination-preserving (stealthy)—and shows large drops in closed-loop task success (e.g., 96.5%→43.1%).