One-sentence Summary

MindLab, Fudan University, introduces LongStraw, an architecture-aware execution stack that evaluates the shared prompt once without autograd, retains only model-specific state, and replays response branches individually, enabling GRPO-based RL post-training on Qwen3.6-27B and GLM-5.2 with up to 4.46M-token contexts on a fixed GPU budget, thereby lowering the hardware barrier for long-context AI agent training.

Key Contributions

• LongStraw is an execution stack that decouples shared prompt evaluation from response backpropagation by detaching the prompt state and replaying short response branches serially under autograd, converting the full-sequence training graph into a single-response graph to reduce GPU memory usage.
• Implemented on Qwen3.6-27B (8 H20 GPUs) and GLM-5.2 (32 H20 GPUs), LongStraw completes grouped scoring and backward for 2.1M-token prompts, scales to 4.46M positions with only 0.21 GB memory increase from group size 2 to 8, and performs end-to-end backward through all 78 layers of GLM-5.2.
• The experiments show that practical context limits for RL post-training are determined by tensor-state lifetime and physical ownership, and that the execution capacity demonstrated by LongStraw lowers the hardware barrier for long-context training under limited accelerator budgets.

Introduction

AI agents increasingly execute long, multi-step trajectories where reasoning, actions, and observations are carried in the prompt context. Post-training these agents with Group Relative Policy Optimization (GRPO) requires scoring multiple responses that share the same long prompt, but the combined memory footprint of the prompt graph, response graphs, cached attention states, and distributed communication quickly exceeds fixed GPU memory budgets. Existing memory-efficient attention and parameter-efficient fine-tuning methods reduce individual costs, while scale-out systems like Ring Attention extend sequence length by adding more accelerators; neither by itself enables a GRPO update to fit under a strict, fixed device count.

The authors present LongStraw, a shared-prompt execution design that decouples the long prompt from the short response computation. The prompt is evaluated once without automatic differentiation, its architecture-specific state is stored, and then each response is replayed sequentially under autograd, accumulating gradients before the optimizer step. This serial replay avoids keeping the prompt graph and all response graphs live simultaneously. The approach is implemented for two distinct model families: a dense-hybrid Qwen stack with recurrent and full-attention layers, and a GLM compressed-attention/MoE stack, demonstrating that long-context GRPO feasibility under a fixed GPU budget is primarily a problem of tensor lifetime and ownership rather than raw context length.

Method

The authors define the systems problem through the update graph rather than a particular attention kernel. A grouped update consists of five logically different phases: prompt capture, pre-step scoring, advantage construction, policy replay, and optimizer transaction.

Refer to the framework diagram:

In conventional full-sequence autograd, the live graph spans the prompt and response tokens. The authors instead capture a read-only prompt state, compute old and reference scores under unchanged pre-step parameters, and replay one response graph at a time. This serial replay bounds live policy-autograd activations by one response, although group-indexed inputs and frozen scores still grow with the supplied group. The clipped policy term follows PPO, while group-relative advantages and member normalization follow the GRPO objective.

The two models differ along two independent axes: the feed-forward axis (dense versus MoE) and the token-mixing axis (GDN/full attention versus MLA/DSA).

As shown in the figure below:

Qwen3.6-27B combines 48 recurrent GDN and 16 full-attention layers with dense FFNs. GLM-5.2 combines 21 index-computing and 57 IndexShare layers with three dense and 75 MoE FFNs. These architectural choices determine parameter residency, token routing, and the shape of activation buffers, as well as the retained prompt state and response-time collectives.

LongStraw is the execution stack developed to retain a tensor across the prompt boundary only when a later response token depends on it. The runtime performs a transaction where the prompt is run under the current policy with autograd disabled, saving model-specific prompt state and releasing transient tensors. The completed prompt state is treated as read-only. For each group member, the short current-policy response path is rebuilt under autograd, reusing the read-only prompt state, backpropagating the member loss, and immediately releasing that member graph. After all backwards, accumulated local gradients are retained and one optimizer call is issued per worker.

For Qwen, compact GDN and KV state are retained on GPU. For GLM, CP-local MLA and indexer-key pages are retained on CPU. During response replay, GLM stages the pages needed by one layer, executes the short response layer, and releases the staged copy. Whole-layer checkpointing is applied to the short response graph to bound activation lifetime, recomputing attention projection, sparse selection, and MoE paths without retaining the complete layer graph.

The final grouped run for GLM is the endpoint of a budget-constrained progression, holding the 32-H20 allocation fixed and moving one limiting resource at a time.

Refer to the framework diagram:

The progression isolates dependency classes including full-graph memory, prefix-state capacity, differentiable replay, cross-layer architecture state, device placement, and group ordering. Stage I localized the full-graph bottleneck, showing that the long prompt could not remain inside the differentiable graph. Stage II scaled prefix-state capacity by capturing MLA latent and DSA indexer-key pages without autograd. Stage III established single-layer differentiable replay with CPU-resident pages. Stage IV achieved all-layer architecture closure, fixing IndexShare lifetime and activation lifetime issues. Stage V established parallel ownership and CPU pages. Finally, the fresh grouped run executes two serial backwards and calls the distributed optimizer once per worker, though CP-local DSA and missing CP gradient finalization remain limitations.

Experiment

The experiments evaluate fixed-budget execution of long-context GRPO replay for dense Qwen and MoE GLM models, confirming that scoring, shaped backward, and optimizer calls complete on all workers at up to 2.1M and 4.25M context positions. These receipts establish execution capacity and forward fidelity but stop short of distributed update consistency and full gradient parity, as critical cross-rank gradient reductions for replicated adapters are either missing or bypassed. The findings demonstrate that separating prompt capture from differentiable replay and using paged attention allow fitting extreme contexts within limited GPU memory, yet they do not constitute a coherent distributed training step or a model quality claim.

Prompt-boundary state is divided into durable tensors that persist across no-grad prefix capture and transient work discarded after capture. Durable components like Qwen full-attention KV pages and GLM MLA latent pages are stored compactly across ranks and scale sub-linearly with prompt length, while transient activations are recomputed blockwise during response replay to keep peak memory within a fixed budget. This separation allows long-context execution without adding devices by trading CPU offload and recomputation for GPU residency. Qwen full-attention KV pages are stored in compact GPU pages across 8 ranks, scaling as O(P/8) per rank, and are read by every response query during global LSE/output merge. GLM durable states are offloaded to CPU, with MLA latent pages and DSA indexer-key pages each scaling as O(P/32) per rank, and are restored layer by layer during response replay while transient attention and MoE work is recomputed under whole-layer checkpointing.

Context parallelism and expert parallelism distribute prompt state and routed experts, but neither layout alone yields a synchronized distributed update. In the examined configurations, missing gradient reductions for adapters or non-expert weights cause each rank to apply its own optimizer step, so the global update is not coherent. The experiments demonstrate training feasibility envelopes of up to 4.25M context on eight GPUs and 2.1M on 32 GPUs, without convergence or learning guarantees. CP8 runs a global response-attention forward but only reduces dQ in the backward, leaving dK/dV and adapter gradients unreduced; AdamW is called per rank, so a consistent global adapter update is not established. EP32 dispatches routed experts across 32 ranks, but DSA scoring and sparse attention remain local to each CP shard, and a custom code path skips finalize_model_grads, leaving CP-replicated non-expert adapter gradients unreduced. In both layouts, the prompt is detached and the response objective is local to each context shard, so the distributed optimizer step does not produce a globally consistent model update. The Qwen envelope reaches 4,456,448 positions on eight H20 GPUs, while the GLM envelope reaches 2,097,152 positions on 32 H20 GPUs; both are fixed-budget feasibility envelopes, not evidence of convergent training or a context-length record.

Two audited fixed-budget execution receipts confirm full training steps (forward, backward, and optimizer calls) near 2 million context positions: Qwen on 8 H20s with a 2.09M-context prompt and 8,192 response inputs, and GLM on 32 H20s with a 2.1M-context prompt and only 3 response inputs. Both record only worker-local events without verified distributed consistency, and GLM’s whole-run peak memory was not captured. The same 8 H20 budget later extends Qwen to a complete 4.25M context receipt, showing that 2M is an achieved operating point rather than a capacity ceiling. Qwen’s G=8 receipt on 8 H20s completes eight serial members and a local AdamW call in 6,785 seconds, peaking at 97.711 GB per rank. GLM’s receipt on 32 H20s finishes in 2,975 seconds (plus model creation) but lacks a whole-run peak memory measurement; a capture-window diagnostic shows 112.6–145.1 GB per rank. All receipts use detached prompt states and local optimizer steps, so distributed parameter updates are not guaranteed to be coherent.

The GLM rerun protocol pins exact source identities for MinT Runtime, Megatron-LM, Megatron-Bridge, verl, and the GLM-5.2 model snapshot, all declared as clean checkouts or with two known patches. These identities are prospective, binding future reruns to an auditable stack but not retroactively verifying the earlier 2M execution receipt. The results further show that on the same eight-H20 budget used for the GLM 2M receipts, Qwen reaches a 4.25M context, underscoring that the work is an accelerator-bounded feasibility envelope rather than a context-length record. The declared source stack includes MinT Runtime, Megatron-LM, Megatron-Bridge, verl, and GLM-5.2 at exact commits, with clean checkouts and two documented patches. These pinned identities are prospective: they bind future reruns to an auditable stack but do not retroactively attest to the historical 2M receipt. The GLM-5.2 model snapshot has a native maximum of 1,048,576 positions, and the 2M prompt is twice that limit; on the same eight-H20 budget, Qwen reaches 4.25M, confirming the fixed-budget envelope framing.

The GLM final-run trace inventory records 128 trace files split across old and policy phases, where old phases capture only forward layer-end events and policy phases capture both forward and backward events for all 78 layers on every rank. Every policy rank checkpointed all 78 layers, totaling 4,992 forward and 4,992 backward events, though the trace is an execution receipt that does not guarantee numerical correctness. Policy phases emit 4,992 forward and 4,992 backward layer-end events across all ranks, with full checkpoint coverage of 78 of 78 layers per rank. Old phases contain only 78 forward events per rank and no backward events, providing no policy backward coverage. The trace serves as an execution-level receipt, confirming terminal path completion on every worker but not verifying gradient correctness or cross-rank reductions.

The experiments assess a memory-saving design that splits prompt states into durable, compactly stored tensors (offloaded or recomputed) and transient, recomputed activations, enabling full training steps on up to 4.25M-token contexts on 8–32 H20 GPUs. Both context-parallel and expert-parallel layouts are explored, but missing gradient reductions mean each rank applies its own optimizer step, so no globally coherent model update is produced and convergence is not guaranteed. Audited execution receipts confirm complete forward, backward, and optimizer calls at 2M and 4.25M context sizes, though these are fixed-budget feasibility envelopes rather than evidence of effective long-context training. The GLM rerun protocol pins source identities for prospective reproducibility but does not retroactively verify the earlier runs' correctness.