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

Dieser Beitrag stellt eine neuartige MPI-basierte Methode zur Konstruktion und Simulation großskaliger spikender neuronaler Netzwerke auf Multi-GPU-Clustern und Exascale-Supercomputern vor, die durch optimierte lokale Konnektivität und Spike-Austauschstrategien eine effiziente Skalierung für komplexe kortikale Modelle demonstriert.

Das Wesentliche

Stellen Sie sich vor, Sie versuchen, das menschliche Gehirn auf einem Computer zu simulieren. Das Gehirn ist eine massive Stadt mit etwa 86 Milliarden Neuronen, wobei jedes Neuron ein Haus ist, das jede Sekunde winzige elektrische „Textnachrichten" (sogenannte Spikes) an Tausende anderer Häuser sendet. Um dies zu simulieren, benötigen Sie einen Supercomputer mit Tausenden von Grafikkarten (GPUs), die zusammenarbeiten.

Das Problem ist, dass diese GPUs wie Inseln sind. Sie sind schnell, sprechen aber nicht leicht miteinander. Wenn eine Insel eine Nachricht an eine andere senden möchte, muss der „Briefträger" (das Kommunikationssystem) hin und her laufen, was alles verlangsamt.

Diese Arbeit stellt eine neue, viel schnellere Methode vor, um die Karte dieser Verbindungen bevor die Simulation beginnt zu erstellen, damit die GPUs die Simulation ausführen können, ohne im Verkehr stecken zu bleiben.

So haben sie es gemacht, einfach erklärt:

1. Der alte Weg: Die Karte auf dem Festland erstellen

Früher, wenn Wissenschaftler ein Gehirnnetzwerk simulieren wollten, erstellten sie die „Verbindungskarte" zuerst auf dem langsamen, zentralen Computer (der CPU). Dann mussten sie diese riesige Karte auf die schnellen GPUs kopieren.

• Die Analogie: Stellen Sie sich vor, Sie organisieren eine massive Party. Bei der alten Methode schrieben Sie den Namen jedes einzelnen Gastes und wen er kennt auf ein Stück Papier in der Küche (CPU) und liefen dann in jeden einzelnen Raum (GPU), um ihnen eine Kopie der Liste zu übergeben. Dies nahm nur für die Vorbereitung sehr lange Zeit in Anspruch.

2. Der neue Weg: Die Karte innerhalb der Räume erstellen

Die Autoren entwickelten eine neue Methode, bei der jede GPU ihren eigenen Teil der Verbindungskarte direkt in ihrem eigenen Speicher erstellt, ohne auf den zentralen Computer zu warten.

• Die Analogie: Jetzt schreibt nicht mehr in der Küche die Liste, sondern jeder Raum hat sein eigenes Notizbuch. Sobald die Party beginnt, schreiben die Gäste in jedem Raum sofort auf, wen sie kennen. Es ist kein Hin- und Herlaufen zur Küche mehr nötig.
• Das Ergebnis: Dieser „onboard"-Aufbau ist mehr als 10-mal schneller als der alte Weg. In einem Test dauerte es 55 Sekunden, um das Netzwerk zu erstellen, statt fast 12 Minuten.

3. Zwei Arten, Nachrichten zu senden

Sobald die Karte erstellt ist, müssen die GPUs während der Simulation die „Textnachrichten" (Spikes) austauschen. Die Arbeit testete zwei verschiedene Strategien hierfür, je nachdem, wie das Netzwerk organisiert ist:

• Strategie A: Der direkte Anruf (Punkt-zu-Punkt)

  • Funktionsweise: Wenn ein Neuron in GPU #1 mit einem spezifischen Neuron in GPU #2 sprechen muss, ruft es diese spezifische GPU direkt an.
  • Am besten geeignet für: Netzwerke, bei denen Verbindungen ungleichmäßig oder spezifisch sind (wie ein reales Gehirn, in dem sich manche Bereiche viel unterhalten, aber nicht mit allen).
  • Die Behauptung der Arbeit: Sie verwendeten dies für ein Modell des visuellen Kortex eines Affen (32 verschiedene Bereiche). Es funktionierte perfekt und bewies, dass die neue Karten-Erstellungsmethode mit komplexen, realen Gehirnstrukturen kompatibel ist.

• Strategie B: Der Gruppenchat (Kollektive Kommunikation)

  • Funktionsweise: Anstatt einzelne Personen anzurufen, schreit eine GPU ihre Nachrichten an eine ganze Gruppe von GPUs gleichzeitig. Jeder in der Gruppe hört den Schrei und prüft, ob die Nachricht für ihn bestimmt ist.
  • Am besten geeignet für: Riesige, zufällige Netzwerke, in denen jeder mit jedem spricht (wie eine ausgeglichene Menge).
  • Die Behauptung der Arbeit: Sie testeten dies an einem massiven „ausgeglichenen Netzwerk", das auf bis zu 1.024 GPUs skalierte. Dies ist eine enorme Anzahl von Grafikkarten, die zusammenarbeiten. Sie zeigten, dass sich das System selbst mit so vielen Karten reibungslos hochskaliert, ohne abzustürzen.

4. Der Trick mit den „Speicherebenen"

GPUs haben viel Speicher, aber nicht unendlich viel. Die Speicherung der Verbindungskarten für Milliarden von Neuronen nimmt viel Platz in Anspruch.

• Die Analogie: Stellen Sie sich vor, Sie haben einen kleinen Schreibtisch (GPU-Speicher) und ein riesiges Lagerhaus (CPU-Speicher).
• Die Lösung: Die Autoren schufen vier „Ebenen" der Organisation.
  • Ebene 0: Halten Sie die Karten im Lagerhaus (CPU) und bringen Sie nur das mit, was Sie auf den Schreibtisch benötigen. Dies spart Schreibtischplatz, ist aber langsamer beim Abrufen.
  • Ebene 3: Füllen Sie den Schreibtisch mit allem. Dies ist am schnellsten, erfordert aber einen größeren Schreibtisch.
• Die Behauptung der Arbeit: Sie zeigten, dass sie durch die Wahl der richtigen Ebene Simulationen auf dem Leonardo Booster-Supercomputer (der 4.096 GPUs besitzt) durchführen konnten, und sagten sogar voraus, dass der kommende JUPITER-Supercomputer ein Netzwerk mit 230 Millionen Neuronen und 2,5 Billionen Synapsen simulieren könnte. Das entspricht ungefähr der Größe des menschlichen Kortex!

Zusammenfassung der erzielten Ergebnisse

• Geschwindigkeit: Sie machten die „Einrichtungs"-Phase von Gehirnsimulationen um das 10-fache schneller, indem sie die Netzwerkkarte direkt auf den Grafikkarten erstellten.
• Skalierbarkeit: Sie bewiesen, dass dies gleichzeitig auf bis zu 1.024 GPUs funktioniert.
• Flexibilität: Sie zeigten zwei verschiedene Möglichkeiten, die Kommunikation zu handhaben (direkte Anrufe vs. Gruppenchats), damit Wissenschaftler die beste Methode für ihr spezifisches Gehirnmodell wählen können.
• Zukunftssicherheit: Ihre Methoden sind so konzipiert, dass sie auf der nächsten Generation von „Exascale"-Supercomputern funktionieren, die stark genug sein werden, um ein komplettes menschliches Gehirn mit Details einzelner Synapsen zu simulieren.

Kurz gesagt: Sie haben nicht nur die Simulation schneller laufen lassen; sie bauten ein besseres „Straßensystem" für die Daten, damit der Supercomputer nicht schon vor dem Start des Rennens im Verkehr stecken bleibt.

Technische Zusammenfassung

Technischer Zusammenfassung: Skalierbare Konstruktion von Spiking Neural Networks mit bis zu tausenden GPUs

Problemstellung

Die Simulation großskaliger Spiking Neural Networks (SNNs) im Maßstab des menschlichen Großhirns stellt zwei Hauptherausforderungen dar: erhebliche Speicheranforderungen für einzelne Neuronen und Synapsen sowie die Notwendigkeit hoher Verarbeitungsgeschwindigkeiten, um Dynamiken mit einer Präzision unterhalb von einer Millisekunde aufzulösen. Obwohl High-Performance-Computing-Systeme (HPC) mit tausenden GPUs die erforderliche Rechenleistung bieten, hat bisherige GPU-basierte Simulationssoftware noch nicht nachgewiesen, dass sie auf gesamte Rechencluster skalieren kann, während sie gleichzeitig die Infrastruktur- und Genauigkeitsanforderungen der computergestützten Neurowissenschaften erfüllt.

Ein spezifischer Flaschenhals bei verteilten Simulationen großer Punkt-Neuronen-Netzwerke ist die Kommunikation von Spikes zwischen verschiedenen Knoten eines Rechenclusters. Frühere Ansätze, wie Digital Brain oder GeNN, lassen entweder Informationen einzelner Synapsen aus oder sind auf die Ausführung auf einer einzigen GPU beschränkt. Darüber hinaus verlassen sich traditionelle CPU-basierte Simulatoren wie NEST auf eine Round-Robin-Verteilung der Neuronen und kollektive Kommunikation, was homogene Netzwerkstrukturen voraussetzt und die topologische sowie räumliche Heterogenität biologischer Gehirne nicht ausnutzt. Zwar hat NEST GPU einige dieser Probleme adressiert, doch beruhte der anfängliche Netzwerkaufbau auf dem Transfer von Daten vom CPU- in den GPU-Speicher, und dynamische Konstruktionsmethoden waren bisher auf Simulationen mit einer einzigen GPU beschränkt.

Methodik

Diese Arbeit stellt eine neuartige, speichereffiziente Methode zur Konstruktion und Simulation großskaliger SNNs direkt auf Multi-GPU-Systemen unter Verwendung des Message Passing Interface (MPI) vor. Die Kerninnovation besteht darin, den Netzwerkaufbau vollständig im GPU-Speicher („onboard") durchzuführen, ohne während der Konstruktionsphase eine Kommunikation zwischen Prozessen zu benötigen.

Kernalgorithmus

Die Methode unterscheidet zwischen lokalen Verbindungen (Neuronen innerhalb desselben MPI-Prozesses) und fernen Verbindungen (Neuronen über verschiedene Prozesse hinweg).

1. Unabhängiger Aufbau: Jeder MPI-Prozess baut seinen Teil des Netzwerks unabhängig auf. Er erstellt lokale Konnektivität und bereitet Datenstrukturen für ferne Verbindungen vor, ohne mit anderen Prozessen zu kommunizieren.
2. Proxy-Repräsentationen: Für ferne Verbindungen verwendet die Methode „Bildneuronen" (Proxys) in Zielprozessen. Dies sind virtuelle Repräsentationen von Quellneuronen, die sich in anderen MPI-Rängen befinden.
3. Kommunikationskarten: Der Algorithmus instanziiert zusammenhängende Kommunikationskarten im GPU-Speicher, um Spikes effizient zu routen. Diese Karten verknüpfen den Index eines Quellneurons in einem Quellrang mit dem Index seines Bildneurons in einem Zielrang.
4. Kommunikationsschemata: Das Framework unterstützt zwei MPI-Kommunikationsmodi, die vom Benutzer je nach Netzwerkarchitektur auswählbar sind:
   • Punkt-zu-Punkt: Verwendet direkte Kommunikation zwischen zwei Prozessen. Es ist für Netzwerke mit ungleichmäßigen Verteilungen von Neuronen oder Synapsen optimiert (z. B. das Multi-Area-Modell). Es nutzt spezifische Abbildungsstrukturen $(R_{\tau,\sigma}, L_{\tau,\sigma})$ und Sequenzen $(T, P)$ zum Routen von Spikes.
   • Kollektiv: Verwendet gruppengestützte Kommunikation (z. B. MPI_Allgather). Dies ist vorteilhaft für ausgeglichene Netzwerke mit homogenen Kommunikationslasten. Es verwendet gruppenspezifische Indexierungsarrays und Host-Arrays, um das Spike-Routing über mehrere Prozesse gleichzeitig zu verwalten.

GPU-Speicher-Optimierung

Um den GPU-Speicherbedarf und die Simulationsgeschwindigkeit in Einklang zu bringen, implementierten die Autoren vier GPU-Speicherlevel (GMLs):

• Level 0: Karten für ferne Verbindungen und Verbindungszahlen werden im CPU-Speicher gespeichert.
• Level 1: Ähnlich wie Level 0, geht jedoch davon aus, dass alle Quellneuronen Bilder in Zielprozessen haben, wodurch Prüfungen auf tatsächliche Nutzung entfallen (schnellerer Aufbau, potenziell höherer Speicherverschwendung).
• Level 2: Karten und Verbindungsindizes werden im GPU-Speicher gespeichert; Verbindungszahlen werden on-the-fly berechnet. Dies ist das Standardlevel.
• Level 3: Alle Datenstrukturen, einschließlich Verbindungszahlen, werden im GPU-Speicher gespeichert, was den CPU-GPU-Datentransfer minimiert, jedoch einen höheren GPU-Speicherbedarf mit sich bringt.

Bewertete Modelle

1. Multi-Area-Modell (MAM): Ein biologisch detailliertes Modell von 32 visuell relevanten Arealen des Makaken-Kortex ($4,13 \times 10^6$ Neuronen, $24,2 \times 10^9$ Synapsen). Dieses Modell weist komplexe, hierarchische Konnektivität auf und wurde unter Verwendung von Punkt-zu-Punkt-Kommunikation simuliert.
2. Skalierbares Ausgeglichenes Netzwerk: Ein zufälliges Netzwerk aus erregenden und hemmenden Neuronen mit fester Eingangsgrad-Konnektivität, entwickelt zur Bewertung der schwachen Skalierbarkeit. Dieses Modell wurde unter Verwendung kollektiver Kommunikation auf bis zu 1.024 GPUs simuliert.

Hauptergebnisse

Leistung beim Netzwerkaufbau

Die „onboard"-GPU-Konstruktionsmethode zeigte signifikante Beschleunigungen im Vergleich zum vorherigen „offboard"-Ansatz (CPU-basiert):

• MAM-Simulation: Die Zeit für den Netzwerkaufbau verringerte sich von 686,0 s (offboard) auf 55,5 s (onboard), was einer 12,4-fachen Beschleunigung entspricht.
  • Die Erstellung lokaler Verbindungen erzielte eine 20-fache Beschleunigung.
  • Die Erstellung ferner Verbindungen erzielte eine 9-fache Beschleunigung.
  • Die Erstellung von Neuronen/Geräten und die Simulationsvorbereitung erzielten Beschleunigungen von 350-fach bzw. 50-fach.
• Skalierbares Ausgeglichenes Netzwerk: Die Methode konstruierte erfolgreich Netzwerke mit bis zu 230,4 Millionen Neuronen und 2,59 Billionen Synapsen über 1.024 GPUs (256 Knoten).

Zustandsausbreitung und Skalierung

• MAM: Die Zeit für die Zustandsausbreitung (gemessen als Real-Time-Faktor) blieb zwischen offboard- und onboard-Versionen vergleichbar (ca. 15–16), was darauf hindeutet, dass die Optimierung des Aufbaus die Simulationsdynamik nicht negativ beeinflusste.
• Ausgeglichenes Netzwerk: Das System zeigte eine schwache Skalierung bis zu 1.024 GPUs.
  • Speichereffizienz: GPU-Speicherlevel 0 ermöglichte Simulationen bis zu 4.096 Knoten, ohne die Speichergrenzen von NVIDIA A100 GPUs (64 GB) zu überschreiten. Höhere Speicherlevel (2 und 3) boten schnellere Aufbau- und Simulationsgeschwindigkeiten, erreichten jedoch die Speichergrenze bei niedrigeren Knotenanzahlen (ca. 3.072 Knoten für Level 3).
  • Leistung: Das Deaktivieren der Spike-Aufzeichnung im ausgeglichenen Netzwerk reduzierte die Zeit für die Zustandsausbreitung um etwa 20 %.

Validierung

Die neue Konstruktionsmethode wurde gegen die vorherige offboard-Version und den CPU-basierten NEST-Simulator validiert. Trotz Änderungen in den Sequenzen der Zufallszahlengenerierung durch den neuen Algorithmus blieben die statistischen Eigenschaften der Spike-Aktivität (Feuerraten, Variationskoeffizient der Inter-Spike-Intervalle und paarweise Pearson-Korrelationen) erhalten, was die biologische Validität der Simulation bestätigt.

Bedeutung und Behauptungen

Die Autoren behaupten, dass diese Arbeit die erste GPU-basierte SNN-Simulationssoftware bereitstellt, die in der Lage ist, auf gesamte Rechencluster (bis zu tausende GPUs) zu skalieren, während sie Informationen einzelner Synapsen speichert. Die Hauptbeiträge sind:

1. Skalierbarer Aufbau: Ein neuartiger Algorithmus, der Netzwerkkonnektivität direkt im GPU-Speicher aufbaut, wodurch der Engpass des CPU-GPU-Transfers eliminiert und die MPI-Kommunikation während der Konstruktionsphase vermieden wird.
2. Flexibilität: Unterstützung sowohl für Punkt-zu-Punkt- als auch für kollektive MPI-Kommunikation, was eine Anpassung an verschiedene Netzwerktopologien (hierarchisch vs. zufällig/ausgeglichen) ermöglicht.
3. Exascale-Bereitschaft: Die Autoren extrapolieren, dass ihr Ansatz Netzwerke mit $2 \times 10^{10}$ Neuronen und $10^{14}$ Synapsen auf dem kommenden JUPITER-Exascale-Supercomputer simulieren könnte. Dieses Maß nähert sich der Konnektivität des menschlichen Kortex an, während die Auflösung einzelner Synapsen beibehalten wird.
4. Effizienz: Durch die Optimierung der Speichernutzung mittels des GML-Systems ermöglicht die Methode die Simulation größerer Netzwerke auf vorhandener Hardware (z. B. passt das MAM auf 8 GPUs statt auf 32) und bietet einen Weg, die volle Kapazität zukünftiger Exascale-Systeme zu nutzen.

Die Autoren schließen, dass dieser Ansatz den kritischen Flaschenhals der Spike-Kommunikation in verteilten Simulationen adressiert und NEST GPU als Referenzplattform für großskalige, biologisch detaillierte neuronale Simulationen auf modernen HPC-Architekturen etabliert.