LongStraw: RL a Contesto Lungo oltre 2 Milioni di Token con un Budget GPU Fisso

Abstract

Un divario crescente separa la lunghezza dei contesti di inferenza dal post-addestramento RL: i sistemi di inferenza si avvicinano a contesti di milioni di token, mentre i carichi di lavoro di post-addestramento rimangono spesso a 256K token o inferiori e si affidano alla generalizzazione della lunghezza in fase di deployment. Il divario è particolarmente importante per gli agenti AI, le cui osservazioni, output degli strumenti, documenti e decisioni precedenti si accumulano lungo traiettorie lunghe. LongStraw è uno stack di esecuzione consapevole dell'architettura per il post-addestramento RL su milioni di token con un budget fisso di GPU, istanziato con Group Relative Policy Optimization (GRPO). Valuta il prompt condiviso senza autograd, conserva solo lo stato specifico del modello necessario ai token successivi e riesegue brevi rami di risposta uno alla volta, riducendo il grafo di addestramento attivo a scapito di un tempo di riesecuzione aggiuntivo. Lo implementiamo per il modello ibrido ricorrente e full-attention Qwen3.6-27B e per il modello mixture-of-experts ad attenzione compressa GLM-5.2. Su otto GPU H20, LongStraw completa lo scoring raggruppato e il backward delle risposte di Qwen su 2,1 milioni di posizioni per gruppi di 2 e 8; aumentare la dimensione del gruppo aggiunge solo 0,21 GB di memoria allocata di picco, mentre un test di stress separato raggiunge 4,46 milioni di posizioni. Su 32 GPU H20, validiamo il percorso di esecuzione end-to-end di LongStraw per un prompt di 2,1 milioni di token attraverso tutti i 78 layer di GLM-5.2. Questi esperimenti stabiliscono la capacità di esecuzione piuttosto che la completa correttezza dell'addestramento, poiché lo stato catturato del prompt è scollegato e alcuni percorsi di composizione forward e del gradiente distribuiti rimangono incompleti.

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 a Contesto Lungo oltre 2 Milioni di Token con un Budget GPU Fisso

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

Abstract

Support