Quanto mismatch tra addestramento e inferenza può tollerare la distillazione on-policy?

Abstract degli autori

La distillazione on-policy (OPD) addestra un modello linguistico studente sulle proprie risposte, guidato da un insegnante fisso. Generare queste risposte è costoso, quindi i sistemi ricorrono all’inferenza quantizzata e ai rollout asincroni per ridurre i costi. Entrambi, tuttavia, introducono una discrepanza tra addestramento e inferenza tra il motore dei rollout e il modello in fase di addestramento. Questa discrepanza può destabilizzare il reinforcement learning (RL), ma i suoi effetti sull’OPD restano poco chiari. Questo articolo analizza quanta discrepanza l’OPD sia in grado di tollerare, da dove derivino i suoi effetti e come i sistemi di addestramento possano sfruttare questa tolleranza. Misurando tutte le fonti di discrepanza sulla stessa scala, esaminiamo la quantizzazione su tre coppie studente-insegnante, con test aggiuntivi sull’obsolescenza dei pesi e sul rumore iniettato in una coppia. L’OPD tollera discrepanze considerevoli senza perdite significative nella qualità dell’addestramento, ma peggiora quando i livelli aumentano, rivelando un budget empirico di discrepanza. Sia il costo in termini di qualità della discrepanza sia quella introdotta dalla quantizzazione dipendono dal modello; alcuni strati sono particolarmente sensibili alla riduzione della precisione. Nei confronti controllati, l’OPD resta stabile con configurazioni di rollout che destabilizzano i metodi di RL. In presenza di rumore sui pesi, la discrepanza sposta anche i prefissi raggiunti durante la generazione, introducendo errori che le correzioni per token non possono affrontare. Per applicare questi risultati, introduciamo un controllo preliminare leggero che verifica le configurazioni di rollout rispetto a un budget calibrato, senza richiedere un addestramento completo per ogni configurazione candidata. Ricette FP8 specifiche per il modello preservano la qualità dell’addestramento e consentono di gestire almeno $1.93\times$ il numero di token della cache KV rispetto a BF16. La combinazione di FP8 e rollout asincroni raggiunge la stessa qualità di addestramento $1.62\times$ più rapidamente rispetto all’addestramento sincrono in BF16. Il codice è disponibile all’indirizzo https://github.com/benyucong/train-inference-mismatch-opd