The Pulse
TRACE registra rollout FP4 5,4 volte più veloci per l’addestramento MoE
Un nuovo studio su TRACE riporta rollout FP4 fino a 5,4 volte più veloci rispetto a BF16 in un test con output lungo su un modello Qwen mixture-of-experts. Il metodo usa i risultati dell’arrotondamento ottenuti durante il rollout per guidar
AI.info Team · Scritto con l’IA ·
Il 6 ottobre, alcuni ricercatori hanno pubblicato TRACE, un metodo per eseguire i rollout dell’apprendimento per rinforzo con numeri in virgola mobile a quattro bit, addestrando al contempo il modello a seguire le stesse decisioni di arrotondamento. Lo studio riporta una capacità di elaborazione in decodifica fino a 5,4 volte superiore rispetto a un rollout BF16 con una lunghezza dell’output di 128K. Il risultato proviene da test su Qwen3.5-35B-A3B e non garantisce le stesse prestazioni su altri modelli o hardware.
TRACE affronta le divergenze tra rollout e addestramento
I rollout sono le sequenze generate da un modello linguistico durante l’addestramento con apprendimento per rinforzo. Queste generazioni forniscono l’esperienza usata per aggiornare la policy, ma l’arrotondamento a bassa precisione può far sì che il rollout e l’addestramento calcolino valori diversi. In un modello mixture-of-experts, le differenze numeriche possono anche influire su quale esperto gestisce un token, creando un’ulteriore fonte di divergenza.
TRACE registra i risultati della quantizzazione durante il rollout e li usa per guidare l’arrotondamento FP4 durante l’addestramento. Gli autori associano le informazioni registrate alle attivazioni dell’addestramento in base alla traiettoria, alla posizione del token e al punto in cui avviene la quantizzazione. Per limitare il carico di archiviazione e trasferimento, il metodo conserva le informazioni su mantissa e scala provenienti da alcuni strati più profondi, anziché salvare ogni valore quantizzato.
«I nostri risultati dimostrano che TRACE consente rollout con pesi e attivazioni FP4 e KV-cache FP4, ottenendo nell’apprendimento per rinforzo prestazioni paragonabili a quelle dei rollout BF16, con rollout fino a 5,4 volte più veloci e solide prestazioni finali in FP4 rispetto alla quantizzazione FP4 a posteriori di policy addestrate in BF16.»
Xin Wang et al., «TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models»
Un bit fornisce l’indicazione per l’arrotondamento
Nella configurazione predefinita dello studio per Qwen3.5-35B-A3B, TRACE conserva un’indicazione di un bit sulla mantissa e informazioni sulla scala per gli ultimi 20 dei 40 strati MoE del modello. Gli autori riferiscono che oltre il 99% dei valori quantizzati che non coincidono differisce di una voce adiacente nel codebook FP4: è questo il motivo per cui usano una piccola indicazione invece di trasferire le informazioni complete sulla quantizzazione. Stimano che questa cache occupi circa 7,5 KB per token generato, all’incirca sei volte meno rispetto alla memorizzazione nella cache dell’intera attivazione quantizzata.
Su quattro benchmark di ragionamento, TRACE ottiene un punteggio medio di 75,3, contro 74,9 del riferimento con rollout BF16 e 68,8 di QUADS, un altro approccio FP4. Il confronto comprende LiveCodeBench, AIME24, AIME25 e HMMT25. Anche i risultati su tre ulteriori modelli MoE sono vicini ai punteggi BF16 riportati: TRACE ottiene 33,0 contro 33,4 su DeepSWE, 70,6 contro 68,8 su Terminal-Bench e 90,2 contro 90,3 su GDPval.
Il dato sulla velocità riguarda il rollout
Il picco di 5,4 volte riguarda la decodifica con una lunghezza dell’output di 128K su Qwen3.5-35B-A3B. Gli autori hanno misurato la capacità di elaborazione del rollout usando SGLang su quattro GPU GB200; il risultato confronta la velocità di generazione, non l’intero processo di addestramento con apprendimento per rinforzo. Nell’intero processo, nel test riportato TRACE porta il tempo di un passo di RL da 664 a 713, un aumento del 7,4% rispetto al rollout FP4 congiunto standard in una configurazione separata che usa 72 GPU GB200.
Il metodo lascia inoltre irrisolta parte della divergenza. Secondo lo studio, gli aggiornamenti asincroni della policy possono far sì che le attivazioni durante l’addestramento differiscano da quelle registrate durante il rollout; TRACE riduce la discrepanza aggiuntiva dovuta a un arrotondamento FP4 incoerente, ma non elimina questa differenza di fondo. Le prove presentate si limitano a quattro modelli MoE, ad alcuni compiti selezionati e a specifiche configurazioni GPU: il dato sulla velocità va quindi inteso come un risultato ottenuto in quelle condizioni.