LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Immagina di cercare di insegnare a un robot a comprendere il mondo mostrandogli un enorme gomitolo di lana aggrovigliato. Ogni nodo nella lana è una persona, un luogo o una cosa, e i fili che li collegano sono le loro relazioni. Questo è il modo in cui i computer imparano a conoscere reti complesse come gli amici sui social media, le citazioni scientifiche o i sistemi di raccomandazione. Questo campo è chiamato Graph Neural Networks (GNN). Il robot deve guardare un nodo, vedere con chi è connesso, guardare le connessioni di quei nodi, e così via, per comprendere l'immagine d'insieme.

Il problema è che man mano che il gomitolo di lana diventa più grande, diventa impossibile far stare l'intero oggetto sulla scrivania del robot. Se il robot cercasse di tenere l'intero gomitolo tra le mani (la sua memoria) per studiarlo, finirebbe lo spazio. Così, gli scienziati di solito tagliano la lana in pezzi più piccoli e gestibili chiamati "mini-batch" per studiarli uno alla volta. Ma c'è un problema: se la lana è conservata in un enorme magazzino lontano, e il robot si trova in un piccolo ufficio, il robot deve correre avanti e indietro verso il magazzino per prendere questi pezzi. Il tempo trascorso a correre avanti e indietro (inviare dati sulla rete) spesso richiede più tempo rispetto al tempo effettivamente impiegato per studiare i nodi. Questo è il "collo di bottiglia della comunicazione" che rallenta tutto.

È qui che entra in gioco una nuova idea chiamata LGNNIC. I ricercatori si sono posti una domanda semplice: e se non ci limitassimo a conservare la lana in un magazzino, ma dessimo al magazzino anche un paio di forbici intelligenti e un piccolo assistente super veloce proprio accanto alla lana? Invece di far correre il robot fino al magazzino per prendere un enorme pezzo di lana e poi tagliarlo per ridurne le dimensioni, l'assistente al magazzino potrebbe fare il taglio e la riduzione prima ancora che il robot lo chieda.

Il documento, intitolato "LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs," esplora esattamente questo. Il team ha costruito un sistema in cui il "magazzino" (un nodo di memoria remoto) ha una scheda di rete speciale e intelligente chiamata SmartNIC (nello specifico, una NVIDIA BlueField-2). Questa SmartNIC agisce come quel piccolo assistente intelligente. Si trova proprio accanto ai dati e compie due trucchi magici prima di inviare qualsiasi cosa al computer principale (il "nodo di addestramento" dotato di GPU potente):

1. Neighbor Sampling (Campionamento dei vicini): Invece di inviare al robot un pezzo enorme e disordinato di grafo, la SmartNIC recide via le connessioni non necessarie, mantenendo solo i vicini più rilevanti affinché il robot possa studiarli. Questo trasforma un enorme e pesante fascio di lana in un pacchetto piccolo e ordinato.
2. Quantization (Quantizzazione): La SmartNIC riduce anche la dimensione delle informazioni all'interno del pacchetto. Converte i dati da un formato pesante ad alta precisione (floating-point a 32 bit) a uno più leggero e leggermente meno preciso (floating-point a 16 bit). Immaginatelo come la compressione di un video in alta definizione in un file più piccolo che appare ancora ottimo, ma occupa la metà dello spazio.

I ricercatori hanno testato questa configurazione utilizzando dataset reali come Reddit (un enorme forum) e reti di articoli accademici. Hanno scoperto che, lasciando alla SmartNIC il lavoro pesante di tagliare e rimpicciolire i dati, potevano ridurre drasticamente la quantità di informazioni che dovevano viaggiare sulla rete.

I risultati sono stati impressionanti. Quando hanno utilizzato un metodo standard e più lento per spostare i dati (come l'invio di email tramite una normale connessione internet, che chiamano "Sockets"), il pre-taglio e la riduzione della SmartNIC hanno reso il processo di addestramento fino a 62,4 volte più veloce per alcuni compiti. Anche utilizzando un metodo di connessione più veloce e diretto (chiamato "DOCA-DMA"), hanno comunque osservato accelerazioni fino a 17,5 volte. Il trucco della "riduzione" (quantizzazione) ha aggiunto ulteriore velocità, rendendo i trasferimenti fino a 3,6 volte più veloci con il metodo standard e 1,3 volte più veloci con il metodo diretto.

Fondamentalmente, il documento sottolinea che, sebbene la SmartNIC sia più lenta nel compiere il taglio rispetto a un supercomputer estremamente potente, il tempo risparmiato dal non dover trascinare un enorme e non tagliato fascio attraverso la rete vale la pena. L' "assistente" al magazzino è più lento nel lavoro, ma evita al "corritore" (la rete) di trasportare un carico pesante. I ricercatori hanno anche verificato che questa riduzione non rendesse errate le risposte del robot; l'accuratezza dei risultati è rimasta quasi esattamente la stessa, con cambiamenti minimi e trascurabili.

In breve, il documento suggerisce che spostando parte del lavoro verso il limite della rete (edge), proprio dove risiedono i dati, possiamo evitare che il robot sprechi tempo correndo avanti e indietro. È un modo intelligente per rendere l'addestramento di IA gigantesche e complesse molto più veloce senza dover costruire computer più grandi e costosi, semplicemente essendo più intelligenti su come spostiamo i dati.