The Pulse
Ai2 rende accessibile Olmo-core 3 per addestrare MoE da mille miliardi di parametri
Ai2 ha rilasciato Olmo-core 3, un framework di addestramento aperto progettato per grandi modelli mixture-of-experts. Nei test è arrivato a 1.200 miliardi di parametri su 512 GPU NVIDIA B300, ma ha usato un instradamento casuale anziché mis
AI.info Team · Scritto con l’IA ·
Ai2 afferma che il suo nuovo sistema di addestramento Olmo-core 3 è stato sottoposto a benchmark su un modello mixture-of-experts da 1,2 bilioni di parametri, con 58,36 miliardi di parametri attivi per ciascun token e il lavoro distribuito su 512 GPU NVIDIA B300. L’istituto ha reso disponibile il framework il 1° ottobre, presentandolo come un’infrastruttura aperta per addestrare modelli sparsi molto più grandi. Il benchmark ha misurato le prestazioni del sistema, non se un modello addestrato con questo sistema produca buone risposte.
Un bilione di parametri, con la riserva di un test
Nel test da 1,2 bilioni di parametri, Ai2 ha registrato un picco osservato di 858 TFLOP/s per GPU, una misura della capacità di calcolo del modello. Ma l’azienda precisa che il test ha usato un instradamento casuale per valutare il sistema. Il test non dimostra la qualità di un modello addestrato, né che quella configurazione abbia completato un intero ciclo di addestramento.
Ai2 descrive anche una configurazione da 2,38 bilioni di parametri che usa DeepEP v2, un metodo alternativo per gestire la comunicazione tra esperti su GPU diverse. L’istituto la definisce un test di capacità di breve durata, non un addestramento prolungato. La distinzione è importante: questi dati indicano che l’infrastruttura può essere configurata a quelle scale, non che Ai2 abbia addestrato un modello capace con bilioni di parametri.
Mantenere gli esperti sulle GPU
I modelli mixture-of-experts suddividono alcune parti di una rete in componenti specializzati chiamati esperti. Il modello indirizza ogni token solo a un sottoinsieme di essi, così il calcolo attivo può rimanere molto inferiore al numero totale di parametri. L’intero modello deve comunque essere memorizzato e aggiornato, mentre l’instradamento degli input tra gli esperti in un cluster comporta un ulteriore carico di comunicazione.
Olmo-core 3 cambia il modo in cui Ai2 gestisce questo lavoro. La precedente implementazione MoE usava il parallelismo dei dati con partizionamento completo, raccogliendo e ripartizionando i pesi per ogni piccolo batch. Il nuovo sistema usa il parallelismo dei dati distribuito e mantiene gli esperti residenti sulle GPU, inviando loro invece i dati pertinenti. Ai2 combina questa impostazione con il parallelismo tra esperti, il parallelismo a pipeline e un ottimizzatore distribuito, ripartendo il modello e lo stato dell’addestramento tra i dispositivi.
Un confronto preliminare su otto GPU B300 ha misurato, per un MoE da 47 miliardi di parametri, 52.000 token al secondo per GPU con Olmo-core 3, contro 19.400 con la precedente implementazione di Ai2: un throughput circa 2,7 volte superiore. In un test separato, l’aumento del numero di esperti da otto a 128, selezionandone quattro per token, ha portato la capacità totale del modello da 4,6 miliardi a 47 miliardi di parametri. I parametri attivi sono rimasti vicini a 3,2 miliardi per token e il throughput è diminuito di meno del 5%.
Precisione inferiore, throughput superiore
Ai2 ha testato anche MXFP8, un formato numerico a precisione inferiore, su quattro GPU B300. In questo benchmark controllato, il throughput dell’addestramento è aumentato di circa il 21% rispetto al valore di riferimento con BF16, mentre il picco di memoria attiva è sceso da 103 GiB a 95 GiB. Secondo il post, la maggior parte del miglioramento in velocità è derivata dal calcolo feed-forward e dal trasferimento di dati tra esperti, più che dal solo meccanismo di attention.
L’istituto segnala i compromessi, anziché considerare ogni ottimizzazione un vantaggio automatico. Il suo rapporto descrive casi in cui la sovrapposizione tra comunicazione e calcolo ha rallentato l’esecuzione complessiva e in cui una metrica pensata per favorire un instradamento bilanciato è migliorata mentre il carico di lavoro effettivo diventava meno bilanciato. Ai2 chiama questo problema «token gerrymandering».
Il codice è pubblico; il prossimo Olmo deve ancora arrivare
Ai2 afferma che ricercatori e sviluppatori possono usare il framework per addestrare modelli MoE e sperimentare con l’instradamento e il parallelismo. Il repository pubblico di Olmo-core include istruzioni per l’installazione, tra cui un pacchetto PyPI, ed è distribuito con licenza Apache-2.0. Rendere disponibile il sistema di addestramento offre ai gruppi esterni accesso a qualcosa di più dei pesi del modello, anche se riprodurre test a questa scala richiede comunque una notevole capacità di calcolo e hardware compatibile.
Questa versione costituisce anche la base per l’Olmo di nuova generazione di Ai2, che secondo l’istituto userà un’architettura MoE. L’annuncio non presenta risultati sulle prestazioni di quel futuro modello. Per ora, il traguardo concreto è più circoscritto: un sistema aperto che Ai2 ha testato su scala di bilioni di parametri, insieme a risultati che distinguono la capacità dell’hardware dall’addestramento completato e dalla qualità del modello.
«Questi test hanno usato un instradamento casuale per misurare le prestazioni del sistema, non la qualità di un modello addestrato».
Kyle Wiggers, Ai2Comms, nell’annuncio di Ai2