LongStraw: 固定GPU予算下での2Mトークンを超える長文脈RL

要旨

推論コンテキスト長とRL後トレーニングの間には、拡大する乖離が存在する。推論システムは100万トークン規模のコンテキストに近づいている一方、後トレーニングのワークロードは多くの場合256Kトークン以下に留まり、デプロイ時の長さの一般化に依存している。この乖離は、エージェントが長期間にわたる軌跡の中で観測、ツール出力、ドキュメント、過去の判断を蓄積するAIエージェントにとって特に重要である。LongStrawは、固定GPU予算の下で100万トークン規模のRL後トレーニングを実現するアーキテクチャ認識型の実行スタックであり、Group Relative Policy Optimization（GRPO）を用いて実現されている。本システムは、共有プロンプトをautogradなしで評価し、後続のトークンに必要なモデル固有の状態のみを保持し、短い応答ブランチを一度に一つずつリプレイすることで、追加のリプレイ時間と引き換えに訓練中の生グラフを削減する。我々はこれを、ハイブリッドリカレント・全注意機構を備えたQwen3.6-27Bと、圧縮注意機構を備えた混合エキスパートモデルGLM-5.2に実装した。8台のH20 GPU上で、LongStrawはグループ化されたQwenのスコアリングおよび応答逆伝播を、グループサイズ2および8において210万ポジションで完了する。グループサイズを増加させても、ピーク時割り当てメモリは0.21 GBしか増加せず、別途のストレステストでは446万ポジションに達する。32台のH20 GPU上では、GLM-5.2の全78層にわたる210万トークンのプロンプトに対して、LongStrawのエンドツーエンド実行経路を検証する。これらの実験は、完全なトレーニングの正確性ではなく実行容量を確立するものである。なぜなら、キャプチャされたプロンプト状態は切り離されており、分散フォワードおよび勾配合成経路の一部が未完了のままであるためである。

English

A growing gap separates inference context lengths from RL post-training: inference systems are approaching million-token contexts, while post-training workloads often remain at 256K tokens or below and rely on length generalization at deployment. The gap is especially important for AI agents, whose observations, tool outputs, documents, and prior decisions accumulate over long trajectories. LongStraw is an architecture-aware execution stack for million-token RL post-training under a fixed GPU budget, instantiated with Group Relative Policy Optimization (GRPO). It evaluates the shared prompt without autograd, retains only model-specific state needed by later tokens, and replays short response branches one at a time, reducing the live training graph at the cost of additional replay time. We implement it for the hybrid recurrent and full-attention Qwen3.6-27B and the compressed-attention mixture-of-experts GLM-5.2. On eight H20 GPUs, LongStraw completes grouped Qwen scoring and response backward at 2.1M positions for groups of 2 and 8; increasing the group size adds only 0.21 GB of peak allocated memory, while a separate stress test reaches 4.46M positions. On 32 H20 GPUs, we validate the end-to-end LongStraw execution path for a 2.1M-token prompt across all 78 layers of GLM-5.2. These experiments establish execution capacity rather than complete training correctness because the captured prompt state is detached and some distributed forward and gradient composition paths remain incomplete.

LongStraw: 固定GPU予算下での2Mトークンを超える長文脈RL

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

要旨

Support