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

Este artículo presenta un método novedoso basado en MPI para construir y simular redes neuronales de picos a gran escala en clústeres de múltiples GPU y superordenadores exaescala, demostrando una escalabilidad eficiente para modelos corticales complejos mediante estrategias optimizadas de conectividad local e intercambio de picos.

Lo esencial

Imagina intentar simular el cerebro humano en una computadora. El cerebro es una ciudad masiva de aproximadamente 86 mil millones de neuronas, donde cada neurona es una casa que envía "mensajes de texto" eléctricos diminutos (llamados picos) a miles de otras casas cada segundo. Para simular esto, necesitas un supercomputador con miles de tarjetas gráficas (GPUs) trabajando juntas.

El problema es que estas GPUs son como islas. Son rápidas, pero no se comunican entre sí fácilmente. Si una isla quiere enviar un mensaje a otra, el "cartero" (el sistema de comunicación) tiene que correr de ida y vuelta, lo que ralentiza todo.

Este artículo presenta una nueva forma, mucho más rápida, de construir el mapa de estas conexiones antes de que comience la simulación, para que las GPUs puedan ejecutar la simulación sin quedar atrapadas en el tráfico.

Así es como lo hicieron, explicado de forma sencilla:

1. La Vieja Forma: Construir el Mapa en el Continente

Anteriormente, cuando los científicos querían simular una red cerebral, construían el "mapa de conexiones" primero en la computadora central lenta (la CPU). Luego, tenían que copiar este mapa masivo a las GPUs rápidas.

• La Analogía: Imagina que estás organizando una fiesta masiva. En el método antiguo, escribías el nombre de cada invitado y a quién conocía en un papel en la cocina (CPU), y luego corrías a cada habitación (GPU) para entregarles una copia de la lista. Esto tomaba mucho tiempo solo para prepararse.

2. La Nueva Forma: Construir el Mapa Dentro de las Habitaciones

Los autores desarrollaron un nuevo método donde cada GPU construye su propia parte del mapa de conexiones directamente dentro de su propia memoria, sin esperar a la computadora central.

• La Analogía: Ahora, en lugar de escribir la lista en la cocina, cada habitación tiene su propia libreta. Tan pronto como comienza la fiesta, los invitados en cada habitación anotan a quién conocen allí mismo. No es necesario correr de ida y vuelta a la cocina.
• El Resultado: Esta construcción "a bordo" es más de 10 veces más rápida que la vieja forma. En una prueba, tardó 55 segundos en construir la red en lugar de casi 12 minutos.

3. Dos Formas de Enviar Mensajes

Una vez construido el mapa, las GPUs necesitan intercambiar los "mensajes de texto" (picos) durante la simulación. El artículo probó dos estrategias diferentes para esto, dependiendo de cómo esté organizada la red:

• Estrategia A: La Llamada Telefónica Directa (Punto a Punto)

  • Cómo funciona: Si una neurona en la GPU #1 necesita hablar con una neurona específica en la GPU #2, llama directamente a esa GPU específica.
  • Ideal para: Redes donde las conexiones son desiguales o específicas (como un cerebro real donde algunas áreas se comunican mucho entre sí, pero no con todos).
  • La Afirmación del Artículo: Lo utilizaron para un modelo de la corteza visual de un mono (32 áreas diferentes). Funcionó perfectamente, demostrando que el nuevo método de construcción de mapas es compatible con estructuras cerebrales complejas y del mundo real.

• Estrategia B: El Chat de Grupo (Comunicación Colectiva)

  • Cómo funciona: En lugar de llamar a individuos, una GPU grita sus mensajes a todo un grupo de GPUs a la vez. Todos en el grupo escuchan el grito y verifican si el mensaje es para ellos.
  • Ideal para: Redes enormes y aleatorias donde todos hablan con todos (como una multitud equilibrada).
  • La Afirmación del Artículo: Lo probaron en una "red equilibrada" masiva que escala hasta 1,024 GPUs. Este es un número enorme de tarjetas gráficas trabajando juntas. Demostraron que incluso con tantas tarjetas, el sistema escala suavemente sin colapsar.

4. El Truco de los "Niveles de Memoria"

Las GPUs tienen mucha memoria, pero no infinita. Almacenar los mapas de conexiones para miles de millones de neuronas ocupa mucho espacio.

• La Analogía: Imagina que tienes un escritorio pequeño (memoria de la GPU) y un enorme almacén (memoria de la CPU).
• La Solución: Los autores crearon cuatro "niveles" de organización.
  • Nivel 0: Mantén los mapas en el almacén (CPU) y solo trae lo que necesites al escritorio. Esto ahorra espacio en el escritorio pero es más lento de recuperar.
  • Nivel 3: Llena el escritorio con todo. Esto es lo más rápido pero requiere un escritorio más grande.
• La Afirmación del Artículo: Demostraron que, al elegir el nivel correcto, podían ejecutar simulaciones en el supercomputador Leonardo Booster (que tiene 4,096 GPUs) e incluso predecir que el próximo supercomputador JUPITER podría simular una red con 230 millones de neuronas y 2.5 billones de sinapsis. ¡Eso es aproximadamente el tamaño de la corteza humana!

Resumen de lo que Lograron

• Velocidad: Hicieron que la fase de "configuración" de las simulaciones cerebrales fuera 10 veces más rápida al construir el mapa de la red directamente en las tarjetas gráficas.
• Escala: Demostraron que esto funciona simultáneamente en hasta 1,024 GPUs.
• Flexibilidad: Mostraron dos formas diferentes de manejar la comunicación (llamadas directas vs. chats de grupo) para que los científicos puedan elegir el mejor método para su modelo cerebral específico.
• Preparado para el Futuro: Sus métodos están diseñados para funcionar en la próxima generación de supercomputadores "Exaescala", que serán lo suficientemente potentes como para simular un cerebro humano completo con detalles individuales de sinapsis.

En resumen, no solo hicieron que la simulación se ejecutara más rápido; construyeron un mejor "sistema de carreteras" para los datos para que el supercomputador no quede atrapado en el tráfico antes de que comience la carrera.

Resumen técnico

Resumen Técnico: Construcción Escalable de Redes Neuronales de Spikes utilizando hasta miles de GPUs

Enunciado del Problema

Simular Redes Neuronales de Spikes (SNN) a gran escala, comparable al de la corteza cerebral humana, presenta dos desafíos principales: requisitos sustanciales de memoria para neuronas y sinapsis individuales, y la necesidad de altas velocidades de procesamiento para resolver dinámicas con precisión submilisegundo. Si bien los sistemas de Computación de Alto Rendimiento (HPC) equipados con miles de GPUs ofrecen la densidad computacional necesaria, el software de simulación basado en GPUs existente aún no ha demostrado la capacidad de escalar a clústeres de computación completos mientras cumple con los requisitos de infraestructura y precisión de la neurociencia computacional.

Un cuello de botella específico en las simulaciones distribuidas de redes de neuronas puntuales grandes es la comunicación de picos (spikes) entre diferentes nodos de un clúster de computación. Enfoques anteriores, como Digital Brain o GeNN, omiten información de sinapsis individuales o se limitan a la ejecución en una sola GPU. Además, los simuladores tradicionales basados en CPU como NEST dependen de la distribución de neuronas tipo round-robin y de comunicación colectiva, lo que asume estructuras de red homogéneas y falla al no aprovechar la heterogeneidad topológica y espacial de los cerebros biológicos. Aunque NEST GPU ha abordado algunos de estos problemas, su construcción inicial de redes dependía de transferir datos de la memoria de la CPU a la de la GPU, y los métodos de construcción dinámica estaban previamente limitados a simulaciones de una sola GPU.

Metodología

Este trabajo presenta un método novedoso y eficiente en memoria para construir y simular SNN a gran escala directamente en sistemas multi-GPU utilizando la Interfaz de Paso de Mensajes (MPI). La innovación central reside en realizar la construcción de la red completamente dentro de la memoria de la GPU ("onboard") sin comunicación entre procesos durante la fase de construcción.

Algoritmo Central

El método distingue entre conexiones locales (neuronas dentro del mismo proceso MPI) y conexiones remotas (neuronas a través de diferentes procesos).

1. Construcción Independiente: Cada proceso MPI construye independientemente su porción de la red. Crea conectividad local y prepara estructuras de datos para conexiones remotas sin comunicarse con otros procesos.
2. Representaciones Proxy: Para las conexiones remotas, el método utiliza "neuronas imagen" (proxies) en los procesos destino. Estas son representaciones virtuales de las neuronas fuente ubicadas en otros rangos MPI.
3. Mapas de Comunicación: El algoritmo instancia mapas de comunicación contiguos en la memoria de la GPU para enrutar los picos de manera eficiente. Estos mapas asocian el índice de una neurona fuente en un rango fuente con el índice de su neurona imagen en un rango destino.
4. Esquemas de Comunicación: El marco soporta dos modos de comunicación MPI, seleccionables por el usuario según la arquitectura de la red:
   • Punto a Punto: Utiliza comunicación directa entre dos procesos. Está optimizado para redes con distribuciones desiguales de neuronas o sinapsis (por ejemplo, el Modelo de Múltiples Áreas). Utiliza estructuras de mapeo específicas $(R_{\tau,\sigma}, L_{\tau,\sigma})$ y secuencias $(T, P)$ para enrutar los picos.
   • Colectivo: Utiliza comunicación basada en grupos (por ejemplo, MPI_Allgather). Esto es ventajoso para redes equilibradas con cargas de comunicación homogéneas. Emplea arrays de indexación específicos del grupo y arrays host para gestionar el enrutamiento de picos a través de múltiples procesos simultáneamente.

Optimización de la Memoria GPU

Para equilibrar el consumo de memoria GPU y la velocidad de simulación, los autores implementaron cuatro Niveles de Memoria GPU (GML):

• Nivel 0: Los mapas de conexión remota y los conteos de conexiones se almacenan en memoria de CPU.
• Nivel 1: Similar al Nivel 0, pero asume que todas las neuronas fuente tienen imágenes en los procesos destino, evitando comprobaciones de uso real (construcción más rápida, potencialmente mayor desperdicio de memoria).
• Nivel 2: Los mapas y los índices de conexión se almacenan en memoria GPU; los conteos de conexiones se calculan sobre la marcha. Este es el nivel predeterminado.
• Nivel 3: Todas las estructuras de datos, incluidos los conteos de conexiones, se almacenan en memoria GPU, minimizando la transferencia de datos CPU-GPU a costa de un mayor uso de memoria GPU.

Modelos Evaluados

1. Modelo de Múltiples Áreas (MAM): Un modelo biológicamente detallado de 32 áreas relacionadas con la visión de la corteza del macaco ($4.13 \times 10^6$ neuronas, $24.2 \times 10^9$ sinapsis). Este modelo presenta conectividad compleja y jerárquica y fue simulado utilizando comunicación punto a punto.
2. Red Equilibrada Escalable: Una red aleatoria de neuronas excitatorias e inhibitorias con conectividad de grado de entrada fija, diseñada para evaluar el rendimiento de escalado débil. Este modelo fue simulado utilizando comunicación colectiva en hasta 1.024 GPUs.

Resultados Clave

Rendimiento de Construcción de Red

El método de construcción "onboard" en GPU demostró aceleraciones significativas en comparación con el enfoque anterior "offboard" (basado en CPU):

• Simulación MAM: El tiempo de construcción de la red disminuyó de 686,0 s (offboard) a 55,5 s (onboard), una aceleración de 12,4x.
  • La creación de conexiones locales vio una aceleración de 20x.
  • La creación de conexiones remotas vio una aceleración de 9x.
  • La creación de neuronas/dispositivos y la preparación de la simulación vieron aceleraciones de 350x y 50x, respectivamente.
• Red Equilibrada Escalable: El método construyó con éxito redes de hasta 230,4 millones de neuronas y 2,59 billones de sinapsis a través de 1.024 GPUs (256 nodos).

Propagación de Estado y Escalado

• MAM: El tiempo de propagación de estado (medido como Factor de Tiempo Real) se mantuvo comparable entre las versiones offboard y onboard (aproximadamente 15–16), lo que indica que la optimización de la construcción no impactó negativamente en la dinámica de la simulación.
• Red Equilibrada: El sistema demostró escalado débil hasta 1.024 GPUs.
  • Eficiencia de Memoria: El Nivel de Memoria GPU 0 permitió que las simulaciones alcanzaran 4.096 nodos sin exceder los límites de memoria de las GPUs NVIDIA A100 (64 GB). Los niveles de memoria más altos (2 y 3) ofrecieron velocidades de construcción y simulación más rápidas, pero alcanzaron el límite de memoria con conteos de nodos más bajos (aproximadamente 3.072 nodos para el Nivel 3).
  • Rendimiento: Desactivar la grabación de picos en la red equilibrada redujo el tiempo de propagación de estado en aproximadamente un 20%.

Validación

El nuevo método de construcción fue validado frente a la versión offboard anterior y el simulador basado en CPU NEST. A pesar de los cambios en las secuencias de generación de números aleatorios debido al nuevo algoritmo, las propiedades estadísticas de la actividad de picos (tasas de disparo, coeficiente de variación de los intervalos entre picos y correlaciones de Pearson por pares) se preservaron, confirmando la validez biológica de la simulación.

Significado y Afirmaciones

El artículo afirma que este trabajo proporciona el primer software de simulación de SNN basado en GPU capaz de escalar a clústeres de computación completos (hasta miles de GPUs) mientras almacena información de sinapsis individuales. Las contribuciones principales son:

1. Construcción Escalable: Un algoritmo novedoso que construye la conectividad de la red directamente en la memoria GPU, eliminando el cuello de botella de transferencia CPU-GPU y evitando la comunicación MPI durante la fase de construcción.
2. Flexibilidad: Soporte para comunicación MPI tanto punto a punto como colectiva, permitiendo la adaptación a diferentes topologías de red (jerárquica frente a aleatoria/equilibrada).
3. Preparación para la Escala Exa: Los autores extrapolan que su enfoque podría simular redes de $2 \times 10^{10}$ neuronas y $10^{14}$ sinapsis en el próximo superordenador exaescalar JUPITER. Esta escala se acerca a la conectividad de la corteza humana manteniendo la resolución de sinapsis individuales.
4. Eficiencia: Al optimizar el uso de memoria a través del sistema GML, el método permite la simulación de redes más grandes en hardware existente (por ejemplo, ajustar el MAM en 8 GPUs en lugar de 32) y proporciona una vía para utilizar la capacidad completa de los futuros sistemas exaescalares.

Los autores concluyen que este enfoque aborda el cuello de botella crítico de la comunicación de picos en simulaciones distribuidas y establece a NEST GPU como una plataforma de referencia para simulaciones neuronales a gran escala y biológicamente detalladas en arquitecturas modernas de HPC.