IT
Un sistema di addestramento adatta la memoria ai rari picchi di carico degli esperti
BaGuaLu-2: Taming Dynamic Memory Fluctuation for Efficient MoE Model Training over 10,000 GPUs · arXiv:2610.23064
In parole semplici
Alcuni modelli linguistici di grandi dimensioni inviano ogni token a reti esperte selezionate. Il numero di token inviati a ciascun esperto può diventare improvvisamente disuguale, aumentando la memoria temporanea necessaria durante l’addestramento.
Una strategia fissa per la memoria può sprecare risorse nei batch normali o rischiare arresti durante un picco. BaGuaLu-2 è un sistema di addestramento che cambia le misure per risparmiare memoria a ogni piccolo batch e sovrappone parte di quel lavoro ad altre attività.
In un’esecuzione su 512 GPU, la sua misura dichiarata dell’uso della capacità di calcolo è rimasta stabile durante picchi che hanno abbassato la misura nei due sistemi di confronto.
La ricerca
Nell’addestramento su larga scala, i modelli mixture-of-experts instradano i token tra molte reti esperte. Rari picchi nell’instradamento possono aumentare il fabbisogno di memoria delle attivazioni; una strategia fissa per la memoria deve quindi rallentare il lavoro ordinario oppure esporre l’esecuzione ad arresti per memoria esaurita.
Come può l’addestramento reagire a queste variazioni brevi senza sostenere lo stesso costo a ogni batch?
BaGuaLu-2 sceglie misure per risparmiare memoria per ogni piccolo batch di addestramento e ne sovrappone il lavoro tra batch mentre cambia la pressione sulla memoria.
BaGuaLu-2 stima il fabbisogno di memoria delle attivazioni e sceglie, separatamente per ogni micro-batch, se ricalcolare i valori, spostarli fuori dalla GPU, ridurne la precisione o non applicare misure.
Sovrappone questo lavoro alle operazioni dei batch vicini per ridurre il tempo dedicato alla gestione della memoria.
Nelle configurazioni di modello e cluster valutate, da 128 a 12,288 GPU, gli autori riportano un throughput medio sull’intera esecuzione pari a 1.14 volte quello di Megatron-Kuai e a 1.13 volte quello di AdaPipe. Qui il throughput è la media dell’MFU (Model FLOPs Utilization), una misura della quota di capacità di calcolo disponibile delle GPU effettivamente utilizzata, per ogni iterazione dell’intera esecuzione, incluso il tempo di riscaldamento delle macchine dopo gli arresti. I risultati supportano un throughput complessivo più elevato in questi test.
Il confronto medio riguarda le configurazioni valutate, che all’inizio usavano la stessa configurazione di parallelismo scelta offline; il paper non specifica se condividessero gli stessi checkpoint iniziali.
Il calcolo dei riavvii include il tempo di riscaldamento delle macchine, ma esclude il tempo dedicato dalle persone alla diagnosi e all’intervento: perciò sottostima questa parte del costo dei riavvii.
Cosa mostrano le prove
- Rari picchi nell’instradamento
I modelli mixture-of-experts su larga scala instradano i token tra reti esperte. Il carico su questi esperti può variare da un’iterazione di addestramento all’altra, modificando la quantità di memoria temporanea necessaria per le attivazioni.
Gli autori hanno misurato uno sbilanciamento superiore a 1.2 volte il carico medio nel 2.9% delle iterazioni, con un tasso massimo di 2.6 volte. La domanda centrale è quindi come gestire picchi occasionali di memoria senza usare sempre una strategia fissa e costosa.
- Memoria adattata a ogni batch
BaGuaLu-2 è un sistema di addestramento che stima il fabbisogno di memoria di ogni micro-batch e poi decide quali operazioni ricalcolare, spostare fuori dalla GPU, eseguire con precisione ridotta o lasciare invariate. Il suo budget di memoria durante l’esecuzione si adatta alle condizioni correnti della pipeline.
Inoltre alterna il lavoro di micro-batch vicini per nascondere parte del tempo aggiuntivo dovuto alla gestione della memoria.
- Confronto degli addestramenti completi
Gli autori hanno confrontato BaGuaLu-2 con Megatron-Kuai e AdaPipe in configurazioni da 128 a 12,288 GPU. I sistemi sono partiti dalla stessa configurazione di parallelismo scelta offline; i sistemi di confronto usavano strategie di memoria statiche e potevano essere riavviati manualmente dopo gli arresti.
Il throughput complessivo è la media dell’MFU (Model FLOPs Utilization), una misura della quota di capacità di calcolo disponibile delle GPU effettivamente utilizzata, per ogni iterazione dell’intera esecuzione, incluso il riscaldamento delle macchine dopo i riavvii.
La media di BaGuaLu-2 è stata pari a 1.14 volte quella di Megatron-Kuai e a 1.13 volte quella di AdaPipe.
- Evidenze su scala e componenti
L’esecuzione valutata alla scala più grande ha addestrato un modello da 1,040 miliardi di parametri su 12,288 GPU e ha riportato un throughput complessivo di 2.012 exaFLOPS. Un confronto progressivo tra componenti ha inoltre associato l’aggiunta successiva di misure statiche granulari, gestione dinamica e pianificazione che sovrappone attività durante l’esecuzione a miglioramenti del throughput rispetto alla configurazione precedente.
Questi risultati descrivono l’esecuzione su larga scala e lo studio di ablazione riportati; non isolano l’effetto di ogni componente su tutti i carichi di lavoro.
- Prestazioni durante i picchi
In un’esecuzione rappresentativa con un modello da 423 miliardi di parametri su 512 GPU, la misura MFU dei sistemi di confronto è diminuita durante i picchi di sbilanciamento, mentre quella di BaGuaLu-2 è rimasta stabile.
Durante questa fase di picchi, il throughput medio di BaGuaLu-2 è stato pari a 1.28 volte quello dei sistemi di confronto. Il risultato riguarda la condizione specifica a cui il sistema è rivolto, ma proviene da una sola configurazione di modello e cluster.
Dettagli
- Autori del paper originale
- Mingshu Zhai, Minkun Guo, Hongrui Lian, Shuangyu Li, Qi Zhang
- Inviato ad arXiv
- Licenza
- CC BY 4.0