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

Cet article présente une nouvelle méthode basée sur MPI pour construire et simuler des réseaux de neurones à spiking à grande échelle sur des clusters multi-GPU et des supercalculateurs exascale, démontrant une mise à l'échelle efficace pour des modèles corticaux complexes grâce à des stratégies optimisées de connectivité locale et d'échange de spiking.

L'essentiel

Imaginez essayer de simuler le cerveau humain sur un ordinateur. Le cerveau est une immense ville d'environ 86 milliards de neurones, où chaque neurone est une maison envoyant de minuscules « messages texte » électriques (appelés potentiels d'action) à des milliers d'autres maisons chaque seconde. Pour simuler cela, vous avez besoin d'un supercalculateur avec des milliers de cartes graphiques (GPU) travaillant ensemble.

Le problème est que ces GPU sont comme des îles. Ils sont rapides, mais ils ne communiquent pas facilement entre eux. Si une île veut envoyer un message à une autre, le « facteur » (le système de communication) doit courir d'avant en arrière, ce qui ralentit tout.

Cet article présente une nouvelle méthode, beaucoup plus rapide, pour construire la carte de ces connexions avant le début de la simulation, afin que les GPU puissent exécuter la simulation sans rester bloqués dans les embouteillages.

Voici comment ils ont procédé, expliqué simplement :

1. L'Ancienne Méthode : Construire la Carte sur le Continent

Auparavant, lorsque les scientifiques voulaient simuler un réseau neuronal, ils construisaient d'abord la « carte de connexion » sur l'ordinateur central lent (le CPU). Ensuite, ils devaient copier cette carte massive vers les GPU rapides.

• L'Analogie : Imaginez que vous organisez une immense fête. Dans l'ancienne méthode, vous écriviez le nom de chaque invité et de ses connaissances sur un papier dans la cuisine (CPU), puis vous couriez dans chaque pièce (GPU) pour leur remettre une copie de la liste. Cela prenait beaucoup de temps juste pour se préparer.

2. La Nouvelle Méthode : Construire la Carte à l'Intérieur des Pièces

Les auteurs ont développé une nouvelle méthode où chaque GPU construit sa propre partie de la carte de connexion directement dans sa propre mémoire, sans attendre l'ordinateur central.

• L'Analogie : Maintenant, au lieu d'écrire la liste dans la cuisine, chaque pièce a son propre bloc-notes. Dès que la fête commence, les invités de chaque pièce notent qui ils connaissent sur place. Plus besoin de courir d'avant en arrière vers la cuisine.
• Le Résultat : Cette construction « embarquée » est plus de 10 fois plus rapide que l'ancienne méthode. Dans un test, il a fallu 55 secondes pour construire le réseau au lieu de près de 12 minutes.

3. Deux Façons d'Envoyer des Messages

Une fois la carte construite, les GPU doivent échanger les « messages texte » (potentiels d'action) pendant la simulation. L'article a testé deux stratégies différentes pour cela, selon la manière dont le réseau est organisé :

• Stratégie A : L'Appel Téléphonique Direct (Point à Point)

  • Fonctionnement : Si un neurone dans le GPU #1 doit parler à un neurone spécifique dans le GPU #2, il appelle directement ce GPU précis.
  • Idéal pour : Des réseaux où les connexions sont inégales ou spécifiques (comme un vrai cerveau où certaines zones parlent beaucoup entre elles, mais pas à tout le monde).
  • L'Affirmation de l'Article : Ils l'ont utilisée pour un modèle du cortex visuel de singe (32 zones différentes). Cela a fonctionné parfaitement, prouvant que la nouvelle méthode de construction de carte est compatible avec des structures cérébrales complexes et réelles.

• Stratégie B : Le Chat de Groupe (Communication Collective)

  • Fonctionnement : Au lieu d'appeler des individus, un GPU crie ses messages à tout un groupe de GPU à la fois. Tout le monde dans le groupe entend le cri et vérifie si le message lui est destiné.
  • Idéal pour : De gigantesques réseaux aléatoires où tout le monde parle à tout le monde (comme une foule équilibrée).
  • L'Affirmation de l'Article : Ils l'ont testée sur un immense « réseau équilibré » s'étendant jusqu'à 1 024 GPU. C'est un nombre énorme de cartes graphiques travaillant ensemble. Ils ont montré que même avec autant de cartes, le système monte en puissance de manière fluide sans planter.

4. L'Astuce des « Niveaux de Mémoire »

Les GPU ont beaucoup de mémoire, mais pas infinie. Stocker les cartes de connexion pour des milliards de neurones prend beaucoup de place.

• L'Analogie : Imaginez que vous avez un petit bureau (mémoire GPU) et un immense entrepôt (mémoire CPU).
• La Solution : Les auteurs ont créé quatre « niveaux » d'organisation.
  • Niveau 0 : Gardez les cartes dans l'entrepôt (CPU) et n'apportez au bureau que ce dont vous avez besoin. Cela économise de l'espace sur le bureau mais est plus lent à récupérer.
  • Niveau 3 : Remplissez le bureau avec tout. C'est le plus rapide mais nécessite un plus grand bureau.
• L'Affirmation de l'Article : Ils ont montré qu'en choisissant le bon niveau, ils pouvaient exécuter des simulations sur le supercalculateur Leonardo Booster (qui possède 4 096 GPU) et même prédire que le futur supercalculateur JUPITER pourrait simuler un réseau de 230 millions de neurones et 2,5 billions de synapses. Cela correspond à peu près à la taille du cortex humain !

Résumé de Ce Qu'ils Ont Réussi

• Vitesse : Ils ont rendu la phase de « configuration » des simulations cérébrales 10 fois plus rapide en construisant la carte du réseau directement sur les cartes graphiques.
• Échelle : Ils ont prouvé que cela fonctionne simultanément sur jusqu'à 1 024 GPU.
• Flexibilité : Ils ont montré deux façons différentes de gérer la communication (appels directs vs chats de groupe) afin que les scientifiques puissent choisir la meilleure méthode pour leur modèle cérébral spécifique.
• Préparation pour l'Avenir : Leurs méthodes sont conçues pour fonctionner sur la prochaine génération de supercalculateurs « Exascale », qui seront assez puissants pour simuler un cerveau humain complet avec des détails de synapses individuels.

En bref, ils n'ont pas seulement fait en sorte que la simulation tourne plus vite ; ils ont construit un meilleur « système routier » pour les données afin que le supercalculateur ne reste pas bloqué dans les embouteillages avant même que la course ne commence.

Résumé technique

Résumé technique : Construction évolutive de réseaux de neurones à spiking utilisant jusqu'à des milliers de GPU

Énoncé du problème

La simulation de réseaux de neurones à spiking (SNN) à grande échelle, à l'échelle du cortex cérébral humain, présente deux défis principaux : des exigences mémoire substantielles pour les neurones et les synapses individuels, et la nécessité de vitesses de traitement élevées pour résoudre les dynamiques avec une précision inférieure à la milliseconde. Bien que les systèmes de calcul haute performance (HPC) équipés de milliers de GPU offrent la densité de calcul nécessaire, les logiciels de simulation existants basés sur les GPU n'ont pas encore démontré leur capacité à s'étendre à l'ensemble des clusters de calcul tout en répondant aux exigences d'infrastructure et de précision de la neuroscience computationnelle.

Un goulot d'étranglement spécifique dans les simulations distribuées de grands réseaux de neurones ponctuels est la communication des spikes entre les différents nœuds d'un cluster de calcul. Les approches précédentes, telles que Digital Brain ou GeNN, omettent soit les informations sur les synapses individuelles, soit sont limitées à une exécution sur un seul GPU. De plus, les simulateurs traditionnels basés sur le CPU comme NEST reposent sur une distribution des neurones en round-robin et une communication collective, ce qui suppose des structures de réseau homogènes et ne parvient pas à exploiter l'hétérogénéité topologique et spatiale des cerveaux biologiques. Bien que NEST GPU ait résolu certains de ces problèmes, sa construction initiale de réseau reposait sur le transfert de données de la mémoire CPU vers la mémoire GPU, et les méthodes de construction dynamique étaient auparavant limitées aux simulations sur un seul GPU.

Méthodologie

Ce travail présente une méthode novatrice et économe en mémoire pour construire et simuler des SNN à grande échelle directement sur des systèmes multi-GPU en utilisant l'interface de passage de messages (MPI). L'innovation centrale réside dans l'exécution de la construction du réseau entièrement dans la mémoire GPU (« onboard »), sans communication inter-processus pendant la phase de construction.

Algorithme principal

La méthode distingue les connexions locales (neurones au sein du même processus MPI) des connexions distantes (neurones entre différents processus).

1. Construction indépendante : Chaque processus MPI construit indépendamment sa portion du réseau. Il crée la connectivité locale et prépare les structures de données pour les connexions distantes sans communiquer avec d'autres processus.
2. Représentations proxy : Pour les connexions distantes, la méthode utilise des « neurones image » (proxys) dans les processus cibles. Il s'agit de représentations virtuelles des neurones sources situés dans d'autres rangs MPI.
3. Cartes de communication : L'algorithme instancie des cartes de communication contiguës dans la mémoire GPU pour acheminer efficacement les spikes. Ces cartes associent l'indice d'un neurone source dans un rang source à l'indice de son neurone image dans un rang cible.
4. Schémas de communication : Le cadre prend en charge deux modes de communication MPI, sélectionnables par l'utilisateur en fonction de l'architecture du réseau :
   • Point à point : Utilise une communication directe entre deux processus. Il est optimisé pour les réseaux avec des distributions inégales de neurones ou de synapses (par exemple, le modèle multi-zone). Il utilise des structures de mappage spécifiques $(R_{\tau,\sigma}, L_{\tau,\sigma})$ et des séquences $(T, P)$ pour acheminer les spikes.
   • Collectif : Utilise une communication basée sur des groupes (par exemple, MPI_Allgather). Cela est avantageux pour les réseaux équilibrés avec des charges de communication homogènes. Il emploie des tableaux d'indexation spécifiques au groupe et des tableaux hôtes pour gérer l'acheminement des spikes à travers plusieurs processus simultanément.

Optimisation de la mémoire GPU

Pour équilibrer la consommation de mémoire GPU et la vitesse de simulation, les auteurs ont mis en œuvre quatre niveaux de mémoire GPU (GML) :

• Niveau 0 : Les cartes de connexions distantes et les comptes de connexions sont stockés dans la mémoire CPU.
• Niveau 1 : Similaire au niveau 0, mais suppose que tous les neurones sources ont des images dans les processus cibles, évitant ainsi les vérifications d'utilisation réelle (construction plus rapide, gaspillage mémoire potentiellement plus élevé).
• Niveau 2 : Les cartes et les indices de connexion sont stockés dans la mémoire GPU ; les comptes de connexion sont calculés à la volée. C'est le niveau par défaut.
• Niveau 3 : Toutes les structures de données, y compris les comptes de connexion, sont stockées dans la mémoire GPU, minimisant le transfert de données CPU-GPU au prix d'une utilisation accrue de la mémoire GPU.

Modèles évalués

1. Modèle multi-zone (MAM) : Un modèle biologiquement détaillé de 32 zones visuelles du cortex du macaque ($4,13 \times 10^6$ neurones, $24,2 \times 10^9$ synapses). Ce modèle présente une connectivité complexe et hiérarchique et a été simulé en utilisant une communication point à point.
2. Réseau équilibré évolutif : Un réseau aléatoire de neurones excitateurs et inhibiteurs avec une connectivité d'entrée fixe, conçu pour évaluer les performances de mise à l'échelle faible. Ce modèle a été simulé en utilisant une communication collective sur jusqu'à 1 024 GPU.

Résultats clés

Performance de construction du réseau

La méthode de construction « onboard » sur GPU a démontré des accélérations significatives par rapport à l'approche précédente « offboard » (basée sur le CPU) :

• Simulation MAM : Le temps de construction du réseau est passé de 686,0 s (offboard) à 55,5 s (onboard), soit une accélération de 12,4x.
  • La création de connexions locales a connu une accélération de 20x.
  • La création de connexions distantes a connu une accélération de 9x.
  • La création de neurones/périphériques et la préparation de la simulation ont connu des accélérations de 350x et 50x, respectivement.
• Réseau équilibré évolutif : La méthode a construit avec succès des réseaux allant jusqu'à 230,4 millions de neurones et 2,59 billions de synapses sur 1 024 GPU (256 nœuds).

Propagation de l'état et mise à l'échelle

• MAM : Le temps de propagation de l'état (mesuré comme facteur temps réel) est resté comparable entre les versions offboard et onboard (environ 15–16), indiquant que l'optimisation de la construction n'a pas eu d'impact négatif sur les dynamiques de simulation.
• Réseau équilibré : Le système a démontré une mise à l'échelle faible jusqu'à 1 024 GPU.
  • Efficacité mémoire : Le niveau de mémoire GPU 0 a permis aux simulations d'atteindre 4 096 nœuds sans dépasser les limites de mémoire des GPU NVIDIA A100 (64 Go). Des niveaux de mémoire plus élevés (2 et 3) ont offert des vitesses de construction et de simulation plus rapides mais ont atteint la limite de mémoire à des nombres de nœuds inférieurs (environ 3 072 nœuds pour le niveau 3).
  • Performance : La désactivation de l'enregistrement des spikes dans le réseau équilibré a réduit le temps de propagation de l'état d'environ 20 %.

Validation

La nouvelle méthode de construction a été validée par rapport à la version offboard précédente et au simulateur NEST basé sur le CPU. Malgré les changements dans les séquences de génération de nombres aléatoires dus au nouvel algorithme, les propriétés statistiques de l'activité de spiking (taux de décharge, coefficient de variation des intervalles inter-spike et corrélations de Pearson par paires) ont été préservées, confirmant la validité biologique de la simulation.

Importance et revendications

L'article revendique que ce travail fournit le premier logiciel de simulation de SNN basé sur les GPU capable de s'étendre à l'ensemble des clusters de calcul (jusqu'à des milliers de GPU) tout en stockant les informations sur les synapses individuelles. Les contributions principales sont :

1. Construction évolutive : Un algorithme novateur qui construit la connectivité du réseau directement dans la mémoire GPU, éliminant le goulot d'étranglement du transfert CPU-GPU et évitant la communication MPI pendant la phase de construction.
2. Flexibilité : Prise en charge des communications MPI point à point et collectives, permettant l'adaptation à différentes topologies de réseau (hiérarchique vs aléatoire/équilibrée).
3. Préparation à l'exaéchelle : Les auteurs extrapolent que leur approche pourrait simuler des réseaux de $2 \times 10^{10}$ neurones et $10^{14}$ synapses sur le futur supercalculateur exaéchelle JUPITER. Cette échelle approche la connectivité du cortex humain tout en maintenant une résolution synaptique individuelle.
4. Efficacité : En optimisant l'utilisation de la mémoire grâce au système GML, la méthode permet la simulation de réseaux plus grands sur le matériel existant (par exemple, faire tenir le MAM sur 8 GPU au lieu de 32) et offre une voie pour exploiter la pleine capacité des futurs systèmes exaéchelle.

Les auteurs concluent que cette approche résout le goulot d'étranglement critique de la communication des spikes dans les simulations distribuées et établit NEST GPU comme une plateforme de référence pour les simulations neuronales à grande échelle et biologiquement détaillées sur les architectures HPC modernes.