Scalable Construction of Spiking Neural Networks using up to thousands of GPUs

Questo articolo presenta un nuovo metodo basato su MPI per costruire e simulare reti neurali a impulsi su larga scala su cluster multi-GPU e supercomputer exascale, dimostrando un'efficienza di scalabilità per modelli corticali complessi attraverso strategie ottimizzate di connettività locale e scambio di impulsi.

L'essenziale

Immagina di provare a simulare il cervello umano su un computer. Il cervello è una vasta città di circa 86 miliardi di neuroni, dove ogni neurone è una casa che invia minuscoli "messaggi di testo" elettrici (chiamati spike) a migliaia di altre case ogni secondo. Per simulare questo, serve un supercomputer con migliaia di schede grafiche (GPU) che lavorano insieme.

Il problema è che queste GPU sono come isole. Sono veloci, ma non comunicano facilmente tra loro. Se un'isola vuole inviare un messaggio a un'altra, il "postino" (il sistema di comunicazione) deve correre avanti e indietro, il che rallenta tutto.

Questo articolo introduce un nuovo metodo, molto più veloce, per costruire la mappa di queste connessioni prima che inizi la simulazione, in modo che le GPU possano eseguire la simulazione senza rimanere intrappolate nel traffico.

Ecco come l'hanno fatto, spiegato semplicemente:

1. Il Vecchio Metodo: Costruire la Mappa sulla Terraferma

In precedenza, quando gli scienziati volevano simulare una rete neurale, costruivano la "mappa delle connessioni" prima sul computer centrale lento (la CPU). Poi, dovevano copiare questa enorme mappa sulle veloci GPU.

• L'Analogia: Immagina di organizzare una festa enorme. Nel vecchio metodo, scrivevi il nome di ogni singolo ospite e chi conosceva su un foglio di carta in cucina (CPU), poi correvi in ogni singola stanza (GPU) per consegnare loro una copia della lista. Questo richiedeva molto tempo solo per prepararsi.

2. Il Nuovo Metodo: Costruire la Mappa Dentro le Stanze

Gli autori hanno sviluppato un nuovo metodo in cui ogni GPU costruisce la sua propria parte della mappa delle connessioni direttamente nella sua memoria, senza attendere il computer centrale.

• L'Analogia: Ora, invece di scrivere la lista in cucina, ogni stanza ha il suo taccuino. Non appena inizia la festa, gli ospiti in ogni stanza scrivono chi conoscono proprio lì. Non è necessario correre avanti e indietro in cucina.
• Il Risultato: Questa costruzione "a bordo" è più di 10 volte più veloce del vecchio metodo. In un test, ci sono voluti 55 secondi per costruire la rete invece di quasi 12 minuti.

3. Due Modi per Inviare Messaggi

Una volta costruita la mappa, le GPU devono scambiarsi i "messaggi di testo" (spike) durante la simulazione. L'articolo ha testato due strategie diverse per questo, a seconda di come è organizzata la rete:

• Strategia A: La Chiamata Telefonica Diretta (Punto a Punto)

  • Come funziona: Se un neurone nella GPU #1 deve parlare con un neurone specifico nella GPU #2, chiama direttamente quella specifica GPU.
  • Ideale per: Reti in cui le connessioni sono disuguali o specifiche (come un vero cervello, dove alcune aree parlano molto tra loro, ma non con tutti).
  • L'Affermazione dell'Articolo: L'hanno usata per un modello della corteccia visiva della scimmia (32 aree diverse). Ha funzionato perfettamente, dimostrando che il nuovo metodo di costruzione della mappa è compatibile con strutture cerebrali complesse e reali.

• Strategia B: La Chat di Gruppo (Comunicazione Collettiva)

  • Come funziona: Invece di chiamare gli individui, una GPU urla i suoi messaggi a un intero gruppo di GPU contemporaneamente. Tutti nel gruppo sentono l'urlo e controllano se il messaggio è per loro.
  • Ideale per: Reti enormi e casuali dove tutti parlano con tutti (come una folla equilibrata).
  • L'Affermazione dell'Articolo: L'hanno testata su una massiccia "rete bilanciata" che scala fino a 1.024 GPU. Questo è un numero enorme di schede grafiche che lavorano insieme. Hanno dimostrato che anche con così tante schede, il sistema scala fluidamente senza bloccarsi.

4. Il Trucco dei "Livelli di Memoria"

Le GPU hanno molta memoria, ma non infinita. Memorizzare le mappe delle connessioni per miliardi di neuroni occupa molto spazio.

• L'Analogia: Immagina di avere una piccola scrivania (memoria GPU) e un enorme magazzino (memoria CPU).
• La Soluzione: Gli autori hanno creato quattro "livelli" di organizzazione.
  • Livello 0: Tieni le mappe nel magazzino (CPU) e porta alla scrivania solo ciò che ti serve. Questo risparmia spazio sulla scrivania ma è più lento da recuperare.
  • Livello 3: Riempie la scrivania con tutto. Questo è il più veloce ma richiede una scrivania più grande.
• L'Affermazione dell'Articolo: Hanno dimostrato che scegliendo il livello giusto, potevano eseguire simulazioni sul supercomputer Leonardo Booster (che ha 4.096 GPU) e persino prevedere che il prossimo supercomputer JUPITER potrebbe simulare una rete con 230 milioni di neuroni e 2,5 trilioni di sinapsi. Questo è circa la dimensione della corteccia umana!

Riepilogo di Cosa Hanno Raggiunto

• Velocità: Hanno reso la fase di "configurazione" delle simulazioni cerebrali 10 volte più veloce costruendo la mappa della rete direttamente sulle schede grafiche.
• Scalabilità: Hanno dimostrato che questo funziona fino a 1.024 GPU simultaneamente.
• Flessibilità: Hanno mostrato due modi diversi per gestire la comunicazione (chiamate dirette vs chat di gruppo) in modo che gli scienziati possano scegliere il metodo migliore per il loro specifico modello cerebrale.
• Preparato per il Futuro: I loro metodi sono progettati per funzionare sulla prossima generazione di supercomputer "Exascale", che saranno abbastanza potenti da simulare un intero cervello umano con dettagli a livello di singola sinapsi.

In breve, non hanno solo fatto funzionare la simulazione più velocemente; hanno costruito un sistema di "strade" migliore per i dati in modo che il supercomputer non rimanga intrappolato nel traffico prima ancora che la gara inizi.

Sommario tecnico

Riepilogo Tecnico: Costruzione Scalabile di Reti Neurali a Spike utilizzando fino a migliaia di GPU

Enunciato del Problema

La simulazione di Reti Neurali a Spike (SNN) su larga scala, alla scala della corteccia cerebrale umana, presenta due sfide principali: notevoli requisiti di memoria per singoli neuroni e sinapsi, e la necessità di elevate velocità di elaborazione per risolvere le dinamiche con una precisione sub-millisecondo. Sebbene i sistemi di Calcolo ad Alte Prestazioni (HPC) dotati di migliaia di GPU offrano la densità computazionale necessaria, il software esistente per la simulazione basato su GPU non ha ancora dimostrato la capacità di scalare a interi cluster di calcolo soddisfacendo contemporaneamente le esigenze di infrastruttura e accuratezza della neuroscienza computazionale.

Un collo di bottiglia specifico nelle simulazioni distribuite di reti di neuroni puntiformi su larga scala è la comunicazione degli spike tra diversi nodi di un cluster di calcolo. Approcci precedenti, come Digital Brain o GeNN, omettono le informazioni sulle singole sinapsi o sono limitati all'esecuzione su singola GPU. Inoltre, i simulatori tradizionali basati su CPU come NEST si affidano alla distribuzione round-robin dei neuroni e alla comunicazione collettiva, che presuppone strutture di rete omogenee e non riesce a sfruttare l'eterogeneità topologica e spaziale dei cervelli biologici. Sebbene NEST GPU abbia affrontato alcuni di questi problemi, la sua costruzione iniziale di rete si basava sul trasferimento di dati dalla memoria della CPU a quella della GPU, e i metodi di costruzione dinamica erano precedentemente limitati a simulazioni su singola GPU.

Metodologia

Questo lavoro presenta un metodo innovativo ed efficiente in termini di memoria per costruire e simulare SNN su larga scala direttamente su sistemi multi-GPU utilizzando l'interfaccia di passaggio messaggi (MPI). L'innovazione fondamentale risiede nell'esecuzione della costruzione della rete interamente all'interno della memoria GPU ("onboard") senza comunicazione inter-processo durante la fase di costruzione.

Algoritmo Core

Il metodo distingue tra connessioni locali (neuroni all'interno dello stesso processo MPI) e connessioni remote (neuroni tra processi diversi).

1. Costruzione Indipendente: Ogni processo MPI costruisce indipendentemente la propria porzione di rete. Crea la connettività locale e prepara le strutture dati per le connessioni remote senza comunicare con altri processi.
2. Rappresentazioni Proxy: Per le connessioni remote, il metodo utilizza "neuroni immagine" (proxy) nei processi target. Si tratta di rappresentazioni virtuali dei neuroni sorgente situati in altri rank MPI.
3. Mappe di Comunicazione: L'algoritmo istanzia mappe di comunicazione contigue nella memoria GPU per instradare gli spike in modo efficiente. Queste mappe associano l'indice di un neurone sorgente in un rank sorgente all'indice del suo neurone immagine in un rank target.
4. Schemi di Comunicazione: Il framework supporta due modalità di comunicazione MPI, selezionabili dall'utente in base all'architettura di rete:
   • Point-to-Point: Utilizza la comunicazione diretta tra due processi. È ottimizzato per reti con distribuzioni disuguali di neuroni o sinapsi (ad esempio, il Multi-Area Model). Utilizza strutture di mappatura specifiche $(R_{\tau,\sigma}, L_{\tau,\sigma})$ e sequenze $(T, P)$ per instradare gli spike.
   • Collettiva: Utilizza la comunicazione basata su gruppi (ad esempio, MPI_Allgather). Questo è vantaggioso per reti bilanciate con carichi di comunicazione omogenei. Impiega array di indicizzazione specifici per gruppo e array host per gestire l'instradamento degli spike attraverso più processi simultaneamente.

Ottimizzazione della Memoria GPU

Per bilanciare il consumo di memoria GPU e la velocità di simulazione, gli autori hanno implementato quattro Livelli di Memoria GPU (GML):

• Livello 0: Le mappe di connessione remota e i conteggi delle connessioni sono memorizzati nella memoria della CPU.
• Livello 1: Simile al Livello 0, ma assume che tutti i neuroni sorgente abbiano immagini nei processi target, evitando controlli sull'effettivo utilizzo (costruzione più veloce, potenziale spreco di memoria più elevato).
• Livello 2: Le mappe e gli indici di connessione sono memorizzati nella memoria GPU; i conteggi delle connessioni sono calcolati al volo. Questo è il livello predefinito.
• Livello 3: Tutte le strutture dati, inclusi i conteggi delle connessioni, sono memorizzate nella memoria GPU, minimizzando il trasferimento di dati CPU-GPU a scapito di un maggiore utilizzo della memoria GPU.

Modelli Valutati

1. Multi-Area Model (MAM): Un modello biologicamente dettagliato di 32 aree visive correlate della corteccia di macaco ($4.13 \times 10^6$ neuroni, $24.2 \times 10^9$ sinapsi). Questo modello presenta una connettività complessa e gerarchica ed è stato simulato utilizzando comunicazione point-to-point.
2. Scalable Balanced Network: Una rete casuale di neuroni eccitatori e inibitori con connettività a grado in ingresso fisso, progettata per valutare le prestazioni di weak scaling. Questo modello è stato simulato utilizzando comunicazione collettiva su fino a 1.024 GPU.

Risultati Chiave

Prestazioni di Costruzione della Rete

Il metodo di costruzione GPU "onboard" ha dimostrato accelerazioni significative rispetto al precedente approccio "offboard" (basato su CPU):

• Simulazione MAM: Il tempo di costruzione della rete è diminuito da 686,0 s (offboard) a 55,5 s (onboard), un accelerazione di 12,4x.
  • La creazione di connessioni locali ha visto un accelerazione di 20x.
  • La creazione di connessioni remote ha visto un accelerazione di 9x.
  • La creazione di neuroni/dispositivi e la preparazione della simulazione hanno visto accelerazioni di 350x e 50x, rispettivamente.
• Scalable Balanced Network: Il metodo ha costruito con successo reti fino a 230,4 milioni di neuroni e 2,59 trilioni di sinapsi su 1.024 GPU (256 nodi).

Propagazione dello Stato e Scaling

• MAM: Il tempo di propagazione dello stato (misurato come Fattore di Tempo Reale) è rimasto comparabile tra le versioni offboard e onboard (circa 15–16), indicando che l'ottimizzazione della costruzione non ha impattato negativamente le dinamiche di simulazione.
• Balanced Network: Il sistema ha dimostrato weak scaling fino a 1.024 GPU.
  • Efficienza di Memoria: Il Livello di Memoria GPU 0 ha permesso alle simulazioni di raggiungere 4.096 nodi senza superare i limiti di memoria delle GPU NVIDIA A100 (64 GB). Livelli di memoria più alti (2 e 3) hanno offerto velocità di costruzione e simulazione più elevate, ma hanno raggiunto il limite di memoria a conteggi di nodi inferiori (circa 3.072 nodi per il Livello 3).
  • Prestazioni: Disabilitare la registrazione degli spike nella rete bilanciata ha ridotto il tempo di propagazione dello stato di circa il 20%.

Validazione

Il nuovo metodo di costruzione è stato validato contro la precedente versione offboard e il simulatore basato su CPU NEST. Nonostante i cambiamenti nelle sequenze di generazione dei numeri casuali dovuti al nuovo algoritmo, le proprietà statistiche dell'attività a spike (tassi di scarica, coefficiente di variazione degli intervalli inter-spike e correlazioni di Pearson a coppie) sono state preservate, confermando la validità biologica della simulazione.

Significato e Affermazioni

Il lavoro afferma che questa ricerca fornisce il primo software di simulazione SNN basato su GPU capace di scalare a interi cluster di calcolo (fino a migliaia di GPU) memorizzando le informazioni sulle singole sinapsi. I contributi principali sono:

1. Costruzione Scalabile: Un algoritmo innovativo che costruisce la connettività di rete direttamente nella memoria GPU, eliminando il collo di bottiglia del trasferimento CPU-GPU ed evitando la comunicazione MPI durante la fase di costruzione.
2. Flessibilità: Supporto per entrambe le comunicazioni MPI point-to-point e collettive, permettendo l'adattamento a diverse topologie di rete (gerarchiche vs casuali/bilanciate).
3. Prontezza per l'Exascale: Gli autori estrapolano che il loro approccio potrebbe simulare reti di $2 \times 10^{10}$ neuroni e $10^{14}$ sinapsi sul futuro supercomputer exascale JUPITER. Questa scala si avvicina alla connettività della corteccia umana mantenendo la risoluzione delle singole sinapsi.
4. Efficienza: Ottimizzando l'uso della memoria attraverso il sistema GML, il metodo consente la simulazione di reti più grandi sull'hardware esistente (ad esempio, adattando il MAM su 8 GPU invece di 32) e fornisce una via per sfruttare la piena capacità dei futuri sistemi exascale.

Gli autori concludono che questo approccio affronta il collo di bottiglia critico della comunicazione degli spike nelle simulazioni distribuite e stabilisce NEST GPU come piattaforma di riferimento per simulazioni neurali su larga scala e biologicamente dettagliate su moderne architetture HPC.