LongStraw：固定 GPU 予算で 200 万 Token 超の長文脈 RL を実行

導入

大規模言語モデルの推論コンテキストは百万 token に近づいています。一方で、RL 後学習の多くは 256K token 以下にとどまり、デプロイ時には長さ方向の汎化に頼ることが少なくありません。この差は AI エージェントで特に深刻です。観測、ツール出力、文書、過去の判断が長い軌跡として蓄積されるため、短い訓練文脈だけでは実運用の状況を十分に再現できないからです。

LongStraw はこの問題に対し、モデルそのものではなく実行スタックとして提案されています。論文では GRPO を用いた RL 後学習を例に、固定 GPU 予算のまま百万 token 級の文脈を扱う方法を示しています。

核心ポイント

• 共有プロンプトを応答分岐から分離：長い共有プロンプトは一度だけ、autograd なしで前向き計算します。これにより、巨大なプロンプト全体を訓練グラフに載せる必要がなくなります。
• 必要な状態だけを保持：後続 token が必要とするモデル固有の状態を保存し、全活性を保持し続けることを避けます。
• 短い応答分岐を逐次 replay：勾配が必要な部分では、応答分岐を一つずつ再実行して autograd を有効にします。ピークメモリは下がりますが、再実行時間は増えます。
• 複数アーキテクチャで検証：混合的な recurrent / full-attention 構成を持つ Qwen3.6-27B と、圧縮注意機構を持つ MoE モデル GLM-5.2 に実装されています。

実験では、8 枚の H20 GPU 上で Qwen のグループ化スコアリングと応答の backward を 2.1M positions で完了しています。グループサイズ 2 と 8 の双方で動作し、グループサイズ増加によるピーク割り当てメモリの増加は 0.21GB にとどまりました。別のストレステストでは 4.46M positions に到達しています。GLM-5.2 では、32 枚の H20 GPU 上で 2.1M token プロンプトを 78 層すべてに通す実行経路が検証されました。

意義と影響

LongStraw の意義は、「モデルが百万 token を読めるか」だけでなく、「その長さで行動を後学習できるか」という問題に踏み込んだ点にあります。長時間動作するエージェントでは、文脈は静的な巨大プロンプトではなく、ツール利用や意思決定を通じて積み上がる状態です。

ただし、論文は限界も明確にしています。今回の結果は実行容量の実証であり、完全な訓練正当性を示すものではありません。取得されたプロンプト状態は detach されており、一部の分散 forward や勾配合成経路も未完成です。したがって LongStraw は完成した訓練手法というより、百万 token 級 RL に向けた有力なシステム設計と見るべきでしょう。

出典：Hugging Face Daily Papers