TRACE: addestramento consapevole della quantizzazione guidato dai rollout per l’apprendimento per rinforzo in FP4 dei modelli linguistici MoE

Abstract degli autori

L’apprendimento per rinforzo (RL) per il post-addestramento dei modelli linguistici di grandi dimensioni (LLM) comporta un notevole dispendio di risorse di calcolo e memoria durante la generazione dei rollout. Questo rende interessante l’uso di rollout a bassa precisione per addestrare i modelli con RL in modo più efficiente. I metodi esistenti per l’RL in FP4 presentano però un limite fondamentale: ottimizzano principalmente l’accuratezza della quantizzazione nei percorsi di addestramento e di rollout separatamente, anziché ridurre direttamente la discrepanza tra i due percorsi di esecuzione quantizzati. In questo lavoro proponiamo TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), un framework di quantizzazione FP4 per l’addestramento con RL di modelli linguistici Mixture-of-Experts (MoE) che affronta questo limite. TRACE integra un addestramento consapevole della quantizzazione guidato dai rollout: usa i risultati della quantizzazione sul versante dei rollout per orientare le decisioni di arrotondamento FP4 sul versante dell’addestramento, riducendo direttamente la discrepanza tra addestramento e rollout. TRACE adotta inoltre uno schema efficiente di memorizzazione temporanea delle informazioni di quantizzazione, che conserva selettivamente le informazioni sulla mantissa e sulla scala degli strati più profondi per ridurre i costi di archiviazione e comunicazione introdotti dalla guida dei rollout. Valutiamo TRACE su quattro modelli linguistici MoE di grandi dimensioni, in compiti di ragionamento, programmazione e RL a lungo orizzonte. I risultati mostrano che TRACE consente di eseguire rollout con pesi e attivazioni in FP4 e cache KV in FP4, ottenendo prestazioni nell’RL paragonabili a quelle dei rollout in BF16, un’accelerazione dei rollout fino a 5,4 volte e solide prestazioni finali in FP4 rispetto alla quantizzazione FP4 applicata a posteriori alle policy addestrate in BF16.