GPU occupate: il metodo di campionamento e pipelining accelera il deep learning su grafici di grandi dimensioni

I grafici, una rete potenzialmente estesa di nodi collegati da bordi, possono essere utilizzati per esprimere e interrogare le relazioni tra i dati, come le connessioni sociali, le transazioni finanziarie, il traffico, le reti energetiche e le interazioni molecolari. Man mano che i ricercatori raccolgono più dati e costruiscono queste immagini grafiche, i ricercatori avranno bisogno di metodi più veloci ed efficienti, nonché di maggiore potenza computazionale, per condurre il deep learning su di essi, come le reti neurali a grafo (GNN).  

Ora, un nuovo metodo, chiamato SALIENT (SAmpling, sLIcing, and data movemeNT), sviluppato dai ricercatori del MIT e IBM Research, migliora le prestazioni di addestramento e inferenza affrontando tre colli di bottiglia chiave nel calcolo. Ciò riduce drasticamente il tempo di esecuzione dei GNN su set di dati di grandi dimensioni, che, ad esempio, contengono sulla scala di 100 milioni di nodi e 1 miliardo di bordi. Inoltre, il team ha scoperto che la tecnica scala bene quando la potenza di calcolo viene aggiunta da una a 16 unità di elaborazione grafica (GPU). Il lavoro è stato presentato alla Fifth Conference on Machine Learning and Systems.

“Abbiamo iniziato a esaminare le sfide incontrate dai sistemi attuali durante il ridimensionamento delle tecniche di apprendimento automatico all'avanguardia per i grafici a set di dati davvero grandi. Si è scoperto che c'era molto lavoro da fare, perché molti dei sistemi esistenti stavano ottenendo buone prestazioni principalmente su set di dati più piccoli che si adattano alla memoria della GPU", afferma Tim Kaler, l'autore principale e postdoc presso il MIT Computer Science e Laboratorio di Intelligenza Artificiale (CSAIL).

Per vasti set di dati, gli esperti intendono scale come l'intera rete Bitcoin, in cui determinati modelli e relazioni di dati potrebbero enunciare tendenze o giochi scorretti. "Ci sono quasi un miliardo di transazioni Bitcoin sulla blockchain e se vogliamo identificare attività illecite all'interno di una tale rete congiunta, allora ci troviamo di fronte a un grafico di tale portata", afferma il coautore Jie Chen, ricercatore senior e manager. di IBM Research e del MIT-IBM Watson AI Lab. "Vogliamo costruire un sistema in grado di gestire quel tipo di grafico e consentire all'elaborazione di essere il più efficiente possibile, perché ogni giorno vogliamo stare al passo con il ritmo dei nuovi dati che vengono generati."

I coautori di Kaler e Chen includono Nickolas Stathas MEng '21 di Jump Trading, che ha sviluppato SALIENT come parte del suo lavoro di laurea; ex stagista del MIT-IBM Watson AI Lab e studentessa laureata del MIT Anne Ouyang; Alexandros-Stavros Iliopoulos, postdoc al MIT CSAIL; Tao B. Schardl, ricercatore del MIT CSAIL; e Charles E. Leiserson, Edwin Sibley Webster Professor of Electrical Engineering presso il MIT e ricercatore presso il MIT-IBM Watson AI Lab.     

Per questo problema, il team ha adottato un approccio orientato ai sistemi nello sviluppo del proprio metodo: SALIENT, afferma Kaler. Per fare ciò, i ricercatori hanno implementato quelle che consideravano importanti ottimizzazioni di base dei componenti che si adattano ai framework di apprendimento automatico esistenti, come PyTorch Geometric e la libreria di grafi profondi (DGL), che sono interfacce per la creazione di un modello di apprendimento automatico. Stathas afferma che il processo è come sostituire i motori per costruire un'auto più veloce. Il loro metodo è stato progettato per adattarsi alle architetture GNN esistenti, in modo che gli esperti di dominio potessero facilmente applicare questo lavoro ai loro campi specificati per accelerare l'addestramento del modello e ottenere più velocemente informazioni durante l'inferenza. Il trucco, ha stabilito il team, era mantenere tutto l'hardware (CPU, collegamenti dati e GPU) sempre occupato: mentre la CPU campiona il grafico e prepara mini-batch di dati che verranno poi trasferiti attraverso il collegamento dati , la GPU più critica sta lavorando per addestrare il modello di machine learning o condurre l'inferenza. 

I ricercatori hanno iniziato analizzando le prestazioni di una libreria di apprendimento automatico comunemente utilizzata per GNN (PyTorch Geometric), che ha mostrato un utilizzo sorprendentemente basso delle risorse GPU disponibili. Applicando semplici ottimizzazioni, i ricercatori hanno migliorato l'utilizzo della GPU dal 10 al 30 percento, ottenendo un miglioramento delle prestazioni da 1.4 a due volte rispetto ai codici di benchmark pubblici. Questo codice di base veloce potrebbe eseguire un passaggio completo su un set di dati di addestramento di grandi dimensioni attraverso l'algoritmo (un'epoca) in 50.4 secondi.                          

Alla ricerca di ulteriori miglioramenti delle prestazioni, i ricercatori hanno deciso di esaminare i colli di bottiglia che si verificano all'inizio della pipeline di dati: gli algoritmi per il campionamento dei grafici e la preparazione dei mini-batch. A differenza di altre reti neurali, i GNN eseguono un'operazione di aggregazione del vicinato, che calcola le informazioni su un nodo utilizzando le informazioni presenti in altri nodi vicini nel grafico, ad esempio, in un grafico di un social network, informazioni da amici di amici di un utente. Con l'aumentare del numero di livelli nella GNN, il numero di nodi a cui la rete deve raggiungere le informazioni può esplodere, superando i limiti di un computer. Gli algoritmi di campionamento di quartiere aiutano selezionando un sottoinsieme casuale più piccolo di nodi da raccogliere; tuttavia, i ricercatori hanno scoperto che le attuali implementazioni di questo erano troppo lente per tenere il passo con la velocità di elaborazione delle moderne GPU. In risposta, hanno identificato un mix di strutture di dati, ottimizzazioni algoritmiche e così via che ha migliorato la velocità di campionamento, migliorando infine la sola operazione di campionamento di circa tre volte, portando il tempo di esecuzione per epoca da 50.4 a 34.6 secondi. Hanno anche scoperto che il campionamento, a una velocità appropriata, può essere eseguito durante l'inferenza, migliorando l'efficienza energetica e le prestazioni complessive, un punto che era stato trascurato in letteratura, osserva il team.      

Nei sistemi precedenti, questa fase di campionamento era un approccio multiprocesso, che creava dati aggiuntivi e movimenti di dati non necessari tra i processi. I ricercatori hanno reso il loro metodo SALIENT più agile creando un singolo processo con thread leggeri che mantenevano i dati sulla CPU nella memoria condivisa. Inoltre, SALIENT sfrutta una cache di processori moderni, afferma Stathas, parallelizzando il feature slicing, che estrae informazioni rilevanti dai nodi di interesse e dai loro vicini e bordi circostanti, all'interno della memoria condivisa della core cache della CPU. Ciò ha nuovamente ridotto il tempo di esecuzione complessivo per epoca da 34.6 a 27.8 secondi.

L'ultimo collo di bottiglia che i ricercatori hanno affrontato è stato quello di convogliare i trasferimenti di dati in mini-batch tra la CPU e la GPU utilizzando una fase di precaricamento, che avrebbe preparato i dati appena prima che fossero necessari. Il team ha calcolato che ciò massimizzerebbe l'utilizzo della larghezza di banda nel collegamento dati e porterebbe il metodo a un utilizzo perfetto; tuttavia, hanno visto solo circa il 90 percento. Hanno identificato e corretto un bug delle prestazioni in una popolare libreria PyTorch che causava comunicazioni di andata e ritorno non necessarie tra CPU e GPU. Con questo bug risolto, il team ha raggiunto un runtime di 16.5 secondi per epoca con SALIENT.

"Il nostro lavoro ha dimostrato, credo, che il diavolo è nei dettagli", afferma Kaler. “Quando presti molta attenzione ai dettagli che influiscono sulle prestazioni durante l'addestramento di una rete neurale a grafo, puoi risolvere un numero enorme di problemi di prestazioni. Con le nostre soluzioni, siamo finiti per essere completamente bloccati dal calcolo della GPU, che è l'obiettivo ideale di un tale sistema".

La velocità di SALIENT è stata valutata su tre set di dati standard ogbn-arxiv, ogbn-products e ogbn-papers100M, nonché in impostazioni multi-macchina, con diversi livelli di fanout (quantità di dati che la CPU preparerebbe per la GPU) e su diverse architetture, inclusa la più recente e all'avanguardia, GraphSAGE-RI. In ogni impostazione, SALIENT ha superato PyTorch Geometric, in particolare sul grande set di dati ogbn-papers100M, contenente 100 milioni di nodi e oltre un miliardo di bordi Qui, era tre volte più veloce, in esecuzione su una GPU, rispetto alla linea di base ottimizzata originariamente creata per questo lavoro; con 16 GPU, SALIENT era otto volte più veloce. 

Mentre altri sistemi avevano hardware e configurazioni sperimentali leggermente diversi, quindi non era sempre un confronto diretto, SALIENT li superava comunque. Tra i sistemi che hanno raggiunto una precisione simile, i numeri rappresentativi delle prestazioni includono 99 secondi utilizzando una GPU e 32 CPU e 13 secondi utilizzando 1,536 CPU. Al contrario, il runtime di SALIENT utilizzando una GPU e 20 CPU era di 16.5 secondi ed era di soli due secondi con 16 GPU e 320 CPU. "Se si guardano i numeri di fondo riportati dal lavoro precedente, il nostro tempo di esecuzione di 16 GPU (due secondi) è di un ordine di grandezza più veloce rispetto ad altri numeri che sono stati riportati in precedenza su questo set di dati", afferma Kaler. I ricercatori hanno attribuito i loro miglioramenti delle prestazioni, in parte, al loro approccio all'ottimizzazione del loro codice per una singola macchina prima di passare all'impostazione distribuita. Stathas afferma che la lezione qui è che per i tuoi soldi "ha più senso utilizzare l'hardware che hai in modo efficiente, e fino all'estremo, prima di iniziare a scalare su più computer", il che può fornire risparmi significativi sui costi e sulle emissioni di carbonio che può venire con la formazione del modello.

Questa nuova capacità consentirà ora ai ricercatori di affrontare e scavare più a fondo in grafici sempre più grandi. Ad esempio, la rete Bitcoin menzionata in precedenza conteneva 100,000 nodi; il sistema SALIENT può gestire abilmente un grafico 1,000 volte (o tre ordini di grandezza) più grande.

"In futuro, cercheremo non solo di eseguire questo sistema di addestramento della rete neurale a grafo sugli algoritmi esistenti che abbiamo implementato per classificare o prevedere le proprietà di ciascun nodo, ma vogliamo anche svolgere attività più approfondite, come identificare modelli comuni in un grafico (modelli di sottografi), [che] potrebbe essere effettivamente interessante per indicare i crimini finanziari", afferma Chen. “Vogliamo anche identificare i nodi in un grafico che sono simili nel senso che potrebbero corrispondere allo stesso cattivo attore in un crimine finanziario. Questi compiti richiederebbero lo sviluppo di algoritmi aggiuntivi e possibilmente anche di architetture di reti neurali».

Questa ricerca è stata supportata dal MIT-IBM Watson AI Lab e in parte dall'US Air Force Research Laboratory e dall'US Air Force Artificial Intelligence Accelerator.