Appianare tutti i picchi di memoria nell’addestramento di modelli Mixture-of-Experts con contesti lunghi

Abstract degli autori

L’addestramento di un modello Mixture-of-Experts (MoE) con un contesto lungo o un batch di grandi dimensioni si interrompe non appena il picco di memoria allocata da un solo componente supera la memoria del dispositivo. Occorre quindi contenere tutti i picchi, non soltanto l’occupazione media. I piani di parallelismo comunemente usati ne lasciano quattro senza limiti, e ciascuno cresce in modo diverso: la distribuzione dei token agli esperti con la matrice di instradamento, la proiezione sul vocabolario con il prodotto tra il numero di token e la dimensione del vocabolario, i confini dei checkpoint dei gradienti con il prodotto tra la profondità e la lunghezza della sequenza, e lo stato dell’ottimizzatore con il numero di parametri. Quale esaurisca per primo la memoria dipende dal modello, dalla lunghezza del contesto e dal numero di dispositivi: ridurre il picco maggiore non fa che far emergere quello successivo. Poniamo un limite a tutti e quattro con schemi che fissano all’avvio la memoria di lavoro sulla GPU: PipelinedLLEP estende il parallelismo degli esperti meno carichi limitando il numero di token che ogni sorgente apporta a un blocco di distribuzione; Ring-DTP fa circolare attivazioni o partizioni dei pesi lungo un anello nella proiezione sul vocabolario e incorpora ogni blocco di logit in un calcolo online della log-sum-exp; Selective checkpoint offload (SCO) conserva nella memoria della CPU il tensore di lunga durata presente a ogni confine di checkpoint; e OffloadStreamAdamW trasforma l’aggiornamento seriale di Adam sulla CPU, usato quando lo stato dell’ottimizzatore è trasferito fuori dalla GPU, in una pipeline di blocchi. Tutti e quattro modificano soltanto l’ordine e la granularità dei calcoli e dei trasferimenti di dati, perciò la funzione di perdita e i gradienti restano esatti. Nei test comparativi sui singoli componenti, riducono il picco di memoria della distribuzione agli esperti MoE fino al $59.3\%$ senza perdita di throughput e quello della proiezione sul vocabolario dell’$86.6\%$; rendono inoltre il passo dell’ottimizzatore trasferito fuori dalla GPU $2.05\times$ più veloce. Usati insieme su modelli MoE da 120 a 667 miliardi di parametri, consentono l’addestramento con una lunghezza del contesto di 1M, pari a $8$--$32\times$ quella raggiungibile con una baseline FSDP2 ottimizzata, e un throughput fino a $10.4\times$ quello della stessa baseline.