🤖 machine learning

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

LongStraw è uno stack di esecuzione consapevole dell'architettura che consente l'apprendimento per rinforzo post-addestramento da un milione di token sotto un budget GPU fisso ottimizzando l'uso della memoria attraverso la valutazione del prompt condivisa senza autograd e il replay sequenziale della risposta, dimostrando con successo la capacità di esecuzione su modelli come Qwen3.6-27B e GLM-5.2 con lunghezze di contesto superiori a 2 milioni di token.

Autori originali: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Che

Pubblicato 2026-07-17

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di cercare di insegnare a un robot super intelligente come risolvere un mistero enorme e composto da molti passaggi. Per farlo, il robot deve leggere una biblioteca intera di indizi (il "contesto") prima ancora di poter iniziare a ipotizzare la soluzione. Nel mondo dell'Intelligenza Artificiale, questo viene chiamato "addestramento a lungo contesto" (long-context training). Di solito, quando addestriamo questi robot, diamo loro un enorme libro di indizi e poi chiediamo loro di scrivere una risposta breve. Il problema è che il cervello del robot (la sua memoria informatica) si riempie così tanto di dettagli del libro che finisce lo spazio prima ancora di poter finire di scrivere la risposta. È come cercare di risolvere un mistero mentre tieni l'intera biblioteca tra le mani; alla fine, le tue braccia si stancheranno troppo per muoversi.

Per risolvere questo problema, gli scienziati hanno costruito librerie di chip per computer (GPU) sempre più grandi per contenere tutte le informazioni. Ma cosa succede se non hai un supercomputer? Cosa succede se hai solo un numero piccolo e fisso di chip, come un set standard di otto o trentadue? Questo è il puzzle che il documento "LongStraw" affronta. Chiede: Possiamo insegnare a un robot a leggere una storia di due milioni di parole e scrivere una soluzione usando solo un piccolo, fisso set di chip per computer, senza comprarne altri? La risposta non è "sì, abbiamo reso i chip più grandi", ma piuttosto "sì, abbiamo cambiato il modo in cui teniamo il libro in mano".

I ricercatori hanno costruito un sistema ingegnoso chiamato LongStraw. Immaginalo come un bibliotecario incredibilmente organizzato. Invece di cercare di tenere l'intera storia di due milioni di parole aperta sulla scrivania mentre il robot scrive la sua risposta (il che farebbe crollare la scrivania), il bibliotecario legge tutta la storia una volta, memorizza solo l'essenziale "vibe" o stato della storia, e poi mette via il pesante libro. Ora, il robot deve solo scrivere la sua breve risposta guardando quella minuscola nota sul "vibe". Se il robot ha bisogno di controllare un dettaglio specifico, il bibliotecario tira fuori rapidamente proprio quella singola pagina, la mostra al robot e poi la rimette immediatamente a posto. In questo modo, la scrivania non si intasa mai, anche se la storia è enorme.

Il documento testa questa idea su due diversi tipi di cervelli robotici: uno chiamato Qwen e un altro chiamato GLM. Hanno cercato di far leggere a questi robot un prompt (la storia) lungo 2.097.152 token (che sono circa due milioni di parole o caratteri) e poi scrivere una risposta. Lo hanno fatto con un budget rigoroso: otto GPU H20 per il robot Qwen e 32 GPU H20 per il robot GLM.

Ecco cosa hanno scoperto:

Funziona per la parte di "lettura": Sono riusciti con successo a elaborare l'intera storia di due milioni di token senza che la memoria del computer andasse in crash. Hanno dimostrato che puoi far entrare questo compito massiccio su un numero fisso di chip essendo intelligenti su quando trattenere le informazioni e quando lasciarle andare.
Risparmia spazio, ma richiede tempo: Mettendo via il pesante libro e tirando fuori le pagine solo quando necessario, hanno risparmiato una enorme quantità di memoria. Per esempio, quando hanno aumentato il numero di "ipotesi" del robot (la dimensione del gruppo) da 2 a 8, l'uso della memoria è cambiato appena (solo circa 0,21 GB in più), ma ci è voluto più tempo per finire il lavoro perché dovevano tirare fuori le pagine una alla volta invece che tutte insieme.
È una "prova di concetto", non una fine perfetta: Il documento è molto onesto su ciò che non ha fatto. Sebbene il sistema abbia eseguito con successo la simulazione e dimostrato che il robot poteva leggere la lunga storia e scrivere la risposta, non ha provato completamente che il robot ha imparato il modo perfetto di aggiornare il proprio cervello. Parte della matematica che collega le diverse "parti del cervello" del robot (i gradienti) non era completamente sincronizzata tra tutti i chip nel modo in cui richiederebbe una sessione di addestramento perfetta. È come mostrare che il motore di un'auto può funzionare con una piccola quantità di carburante, ma non ancora provare che possa percorrere un intero circuito di gara senza stallare.

I ricercatori chiamano il loro risultato una "ricevuta di esecuzione". È una ricevuta che dice: "Sì, abbiamo eseguito con successo questo compito massiccio su questo hardware specifico senza crashare". Non hanno sostenuto di aver vinto una gara o di aver costruito l'auto più veloce; hanno solo dimostrato che l'auto può guidare su questa specifica pista con questa specifica quantità di benzina.

In breve, LongStraw mostra che non serve un supercomputer per addestrare l'IA su storie enormi. Serve solo un modo più intelligente di gestire la propria memoria, scambiando un po' di tempo extra per molto spazio risparmiato. Apre la porta a team più piccoli e ricercatori con budget limitati per sperimentare con questi giganti modelli di IA, dimostrando che si possono fare grandi cose con un kit di strumenti piccolo e fisso se si è abbastanza ingegnosi.

Riepilogo Tecnico: LongStraw: RL a Lungo Contesto Oltre i 2M di Token sotto un Budget Fisso di GPU

1. Definizione del Problema

Si è delineato un divario significativo tra le lunghezze di contesto supportate dai sistemi di inferenza (che si avvicinano ai contesti da milioni di token) e quelle utilizzate durante l'apprendimento per rinforzo (RL) post-training, che spesso rimangono limitate a 256K token o meno. Questo divario è particolarmente critico per gli agenti IA le cui traiettorie accumulano osservazioni, output di strumenti e decisioni pregresse su sequenze molto lunghe.

La sfida principale nell'estendere il post-training RL (specificamente la Group Relative Policy Optimization, o GRPO) a contesti da un milione di token sotto un budget fisso di GPU è la memoria. A differenza dell'inferenza, dove il grafo forward può essere scartato dopo il caching dello stato, l'addestramento deve calcolare il punteggio di più risposte e backpropagarle attraverso di esse. Gli approcci standard richiedono la ritenzione dell'intero prompt e di tutti i grafi delle risposte simultaneamente in memoria, portando a una computazione dell'attenzione quadratica e a uno stoccaggio massiccio delle attivazioni. Tecniche esistenti come FlashAttention, LoRA o Ring Attention affrontano colli di bottiglia specifici, ma non permettono intrinsecamente a un run GRPO con budget GPU fisso di ospitare un contesto di oltre 2M di token senza scalare verso cluster più grandi.

2. Metodologia: Execution Stack di LongStraw

LongStraw è uno stack di esecuzione consapevole dell'architettura, progettato per disaccoppiare il lungo prompt dal calcolo della risposta breve, scambiando tempo di replay aggiuntivo con una riduzione dell'uso della memoria GPU. Il principio di design fondamentale è trattenere un tensore attraverso il confine del prompt solo quando un token della risposta successiva dipende da esso.

2.1 Lo Schema di Esecuzione

Il sistema opera in quattro fasi distinte per un aggiornamento raggruppato:

Cattura del Prompt (No-Grad): Il prompt condiviso ( $P$ token) viene valutato una sola volta con differenziazione automatica disabilitata. Il sistema salva solo lo stato condizionale specifico dell'architettura richiesto per i token futuri (ad esempio, stati ricorrenti, pagine KV, rappresentazioni latenti) e rilascia tutte le attivazioni transitorie (intermedi FFN, scratch dell'attenzione, buffer di routing MoE).
Scoring Pre-Step: Le log-probabilità della vecchia policy e della reference-policy vengono valutate per tutti i membri del gruppo utilizzando lo stato del prompt congelato.
Replay Seriale della Risposta: Invece di mantenere $G$ grafi di risposta simultaneamente, il sistema riproduce un ramo di risposta alla volta sotto autograd. Riutilizza lo stato del prompt catturato (read-only), calcola la perdita, esegue il backpropagation e libera immediatamente il grafo della risposta.
Accumulo dei Gradienti e Ottimizzazione: I gradienti locali di tutti i $G$ membri vengono accumulati, seguiti da una singola chiamata all'ottimizzatore per ogni worker.

2.2 Implementazioni Specifiche dell'Architettura

LongStraw adatta questa strategia a due diverse famiglie di modelli:

Qwen3.6-27B (Ibrido Ricorrente/Full-Attention):
- Ritenzione dello Stato: Mantiene stati ricorrenti compatti Gated DeltaNet (GDN) e pagine KV (Key/Value) sharded per il contesto-parallelismo (CP) per i 16 strati di full-attention su GPU.
- Meccanismo: Utilizza NF4 QLoRA. I 48 strati GDN forniscono uno stato di dimensione fissa, mentre i 16 strati di full-attention utilizzano la compattazione fisica delle pagine. Il sistema esegue un merge globale LSE/output per il forward pass, ma riconosce la mancanza di sincronizzazione del gradiente dell'adapter K/V.
- Hardware: Eseguito su 8 H20 GPU (CP8).
GLM-5.2 (Compressed Attention/MoE):
- Ritenzione dello Stato: Sposta lo stato del lungo prompt (pagine latenti MLA e pagine indexer-key DSA) nella memoria CPU. Solo uno strato di decoder alla volta viene caricato sulla GPU durante il replay della risposta.
- Meccanismo: Utilizza un layout delle pagine di contesto "Zigzag" attraverso 32 rank (CP32) ed Expert Parallelism (EP32). Impiega il checkpointing delle attivazioni per intero strato per limitare il grafo della risposta. Il sistema ricostruisce gli schedule IndexShare strato per strato.
- Hardware: Eseguito su 32 H20 GPU (TP1/CP32/EP32).

3. Contributi Chiave

Design di Esecuzione Shared-Prompt: Uno schema di esecuzione innovativo che memorizza lo stato del prompt una sola volta (distaccato dall'autograd) e serializza il lavoro della risposta, riducendo il grafo di addestramento live da $P + \sum R_i$ a $\max(R_i)$ .
Implementazione Dual-Architecture: Istanziazione riuscita di questo design per due modelli strutturalmente diversi: uno stack denso ibrido ricorrente/attenzione (Qwen) e uno stack MoE a attenzione compressa (GLM).
Ricevute di Esecuzione a Budget Fisso: Evidenza empirica che dimostra come i percorsi di esecuzione GRPO possano completare per prompt che superano i 2M di token (e fino a 4.25M per Qwen) senza aumentare il numero di GPU, documentando esplicitamente i limiti nella composizione dei gradienti distribuiti.

4. Risultati Sperimentali

4.1 Qwen3.6-27B (8 H20s)

Lunghezza del Contesto: Ha eseguito con successo un contesto di 2.097.152 token (2.088.960 prompt + 8.192 risposta).
Efficienza della Memoria: Aumentare la dimensione del gruppo da $G=2$ a $G=8$ ha aumentato la memoria allocata di picco di soli 0,21 GB (0,213%), confermando che la memoria scala con il ramo di risposta più grande, non con la dimensione del gruppo.
Capacità Estesa: Un test di stress ha raggiunto 4,25M di posizioni (4.448.256 prompt + 8.192 risposta) sugli stessi 8 H20, raggiungendo un picco di circa 83 GB per rank.
Limitazione: L'aggiornamento distribuito non è completamente coerente; i gradienti dell'adapter K/V non sono sincronizzati tra i rank e il gradiente dello stato del prompt è staccato (detached).

4.2 GLM-5.2 (32 H20s)

Lunghezza del Contesto: Ha validato il percorso di esecuzione end-to-end per un prompt di 2.097.152 token attraverso tutti i 78 strati.
Strategia di Memoria: Lo stato del prompt residente in CPU (circa 5,8 GiB per rank) ha permesso alla GPU di rimanere entro il budget durante il backward pass di 78 strati.
Limitazione: L'operatore DSA (Dynamic Sparse Attention) nella ricevuta accettata è locale a ciascun shard CP (seleziona i top-2048 da 65k token invece del contesto globale di 2M). Inoltre, il percorso bypassa finalize_model_grads di Megatron, lasciando i gradienti degli adapter non-expert replicati in CP non ridotti prima dello step dell'ottimizzatore.

5. Significato e Rivendicazioni

Il paper posiziona LongStraw non come un "record mondiale" di lunghezza del contesto, ma come un inviluppo di fattibilità a budget fisso.

Durata dello Stato vs Sparsità: Gli autori sostengono che la capacità a budget fisso sia determinata dalla durata del tensore e dalla proprietà fisica, non semplicemente dalla sparsità (MoE) o dai kernel di attenzione efficienti in termini di memoria. Permettendo agli intermedi del prompt di "morire" immediatamente dopo la cattura, il sistema evita l'esplosione di memoria dell'autograd a sequenza completa.
Proprietà Fisica: Il lavoro sottolinea che lo sharding logico non è sufficiente; l'allocazione della memoria fisica deve corrispondere alla strategia di sharding (ad esempio, copiando le pagine in tensori di dimensioni corrette) per rilasciare la memoria efficacemente.
Rivendicazioni Modeste: Gli autori dichiarano esplicitamente che queste sono ricevute di esecuzione, non prove di correttezza dell'addestramento. Lo stato del prompt è staccato (manca il termine $\partial z_P / \partial \theta$ ) e la composizione dei gradienti distribuiti è incompleta per entrambi i modelli. I risultati stabiliscono che il percorso di esecuzione rientra nei vincoli hardware, consentendo ai ricercatori di esplorare l'RL a lungo contesto senza un massiccio scale-out, ma non pretendono di aver risolto ancora la parità del gradiente distribuito o gli aspetti di miglioramento della policy.

In sintesi, LongStraw dimostra che con una gestione attenta della durata dello stato e della proprietà fisica, l'addestramento post-training a milioni di token è realizzabile su cluster GPU relativamente piccoli e a budget fisso, spostando il collo di bottiglia dalla pura capacità di memoria alla gestione dello stato e della pianificazione (scheduling).