Il caricatore di dati di IBM ricerca migliora l'addestramento dei modelli AI per la comunità open source

Come si superano i colli di bottiglia quando si addestrano modelli AI su enormi quantità di dati? Alla conferenza PyTorch di quest'anno, IBM Research ha presentato un data loader rivoluzionario per l'addestramento su larga scala di LLM . Lo strumento, ora disponibile per gli utenti PyTorch, mira a semplificare la formazione su larga scala per un destinatario quanto più vasto possibile.

L'idea del data loader ad alta produttività è nata da problemi pratici osservati dagli scienziati durante l'addestramento dei modelli, poiché il loro lavoro richiedeva uno strumento in grado di elaborare grandi quantità di dati su più dispositivi, il tutto tenendo il passo con GPU sempre più efficienti. Come IBM Research nota nel suo blog a proposito del rilascio, "È tutto grazie a un team di ricercatori che stavano semplicemente creando gli strumenti necessari per portare a termine un lavoro".

Davis Wertheimer di IBM Research spiega alcune delle sfide che possono emergere durante la formazione su larga scala: "Quando si tratta di formazione su larga scala, vale più o meno la regola dell'80/20. L'80% di tutta la letteratura pubblicata esamina i compromessi algoritmici tra memoria e comunicazione della GPU e calcolo. Ma quando provi davvero a costruire qualcosa, l'80% delle volte puoi contare su una lunghissima coda di tutte queste altre questioni pratiche perché la pipeline procede alla velocità del collo di bottiglia più stretto."

Man mano che il team IBM sviluppava la propria piattaforma di formazione, continuava a incontrare ostacoli. "Più miglioriamo nell'uso delle nostre GPU, più ci accorgiamo che il collo di bottiglia è il caricatore di dati", osserva Wertheimer.

Questa consapevolezza ha portato a un doppio processo di sviluppo. "C'è stato un percorso parallelo che ha visto, da un lato, l'evoluzione della nostra piattaforma di formazione e, dall'altro, l'evoluzione costante del nostro caricatore di dati, per stare al passo con le richieste di velocità della nostra piattaforma di formazione, per evitare di bloccarla", spiega.

Linsong Chu di IBM Research delinea le caratteristiche essenziali del caricatore di dati:

Stato e checkpoint: "Ogni volta che salvi un modello, viene salvato anche lo stato del tuo data loader, e ogni volta che recuperi da un checkpoint, sia lo stato del modello che quello del data loader devono essere recuperati contemporaneamente", dice Chu.

Ridimensionamento automatico dei checkpoint: il caricatore di dati si adatta automaticamente alle variazioni del workload durante le sessioni di formazione prolungate. "La formazione può facilmente durare settimane o mesi e ci sono tonnellate di ragioni per cui potresti dover ridimensionare il workload a metà", osserva Chu.

Flusso efficiente dei dati: il sistema supporta il data streaming senza alcun overhead di costruzione per lo shuffling.

Operazione distribuita asincrona: "Vogliamo che il caricatore di dati sia non bloccante", spiega Chu. "Quando salviamo lo stato del caricatore dati, vogliamo che il salvataggio venga distribuito in una forma in cui non sia coinvolta alcuna comunicazione."

Miscelazione dinamica dei dati: il caricatore di dati può adattarsi a diversi rapporti di miscelazione dei dati, il che è utile per le esigenze di formazione in evoluzione.

Efficiente shuffling globale: lo strumento risolve i colli di bottiglia della memoria durante la gestione di grandi set di dati, rendendo efficiente lo shuffling anche quando i dati crescono.

PyTorch nativo, modulare ed esteso: progettato per adattabilità e scalabilità, il data loader è preparato per la crescita futura. "Cosa succederebbe se l'anno prossimo dovessimo gestire 30 trilioni, 50 trilioni o 100 trilioni di token?" chiede Chu. "Il mondo sta cambiando rapidamente, quindi dobbiamo sviluppare un data loader che possa non solo sopravvivere oggi, ma anche domani."

Il team di ricerca IBM ha testato rigorosamente il proprio caricatore di dati per diversi mesi, eseguendo centinaia di lavori di piccole e grandi dimensioni. Hanno osservato numeri di codice stabili e uniformi. Inoltre, l'intero data loader opera asincronamente ed è non bloccante.

"Abbiamo utilizzato molte delle funzionalità integrate di PyTorch per realizzare tutto questo", afferma Wertheimer. "Ecco perché stiamo contribuendo, contribuendo."