LongStraw : RL à long contexte au-delà de 2 millions de tokens sous un budget GPU fixe

Résumé

Un écart croissant sépare les longueurs de contexte d'inférence du post-entraînement par RL : les systèmes d'inférence approchent des contextes d'un million de tokens, tandis que les charges de travail de post-entraînement restent souvent à 256K tokens ou moins et comptent sur la généralisation de longueur lors du déploiement. Cet écart est particulièrement important pour les agents d'IA, dont les observations, les sorties d'outils, les documents et les décisions antérieures s'accumulent sur de longues trajectoires. LongStraw est une pile d'exécution architecturale pour le post-entraînement par RL à un million de tokens sous un budget GPU fixe, instanciée avec l'Optimisation de Politique Relative par Groupe (GRPO). Elle évalue le prompt partagé sans autograd, ne conserve que l'état spécifique au modèle nécessaire aux tokens ultérieurs, et rejoue les branches de réponse courtes une par une, réduisant le graphe d'entraînement actif au prix d'un temps de rejeu supplémentaire. Nous l'implémentons pour le modèle hybride récurrent et à attention complète Qwen3.6-27B et le mélange d'experts à attention compressée GLM-5.2. Sur huit GPU H20, LongStraw complète le scoring groupé de Qwen et la rétropropagation des réponses à 2,1 millions de positions pour des groupes de 2 et 8 ; augmenter la taille du groupe n'ajoute que 0,21 Go de mémoire allouée maximale, tandis qu'un test de stress séparé atteint 4,46 millions de positions. Sur 32 GPU H20, nous validons le chemin d'exécution de bout en bout de LongStraw pour un prompt de 2,1 millions de tokens à travers les 78 couches de GLM-5.2. Ces expériences établissent une capacité d'exécution plutôt qu'une correction complète de l'entraînement, car l'état capturé du prompt est détaché et certains chemins de propagation avant et de composition de gradient distribués restent incomplets.

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 : RL à long contexte au-delà de 2 millions de tokens sous un budget GPU fixe

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

Résumé

Support