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

Este artigo apresenta um método inovador baseado em MPI para construir e simular redes neurais de spiking em grande escala em clusters multi-GPU e supercomputadores exascale, demonstrando escalabilidade eficiente para modelos corticais complexos por meio de estratégias otimizadas de conectividade local e troca de spikes.

O essencial

Imagine tentar simular o cérebro humano em um computador. O cérebro é uma cidade massiva de cerca de 86 bilhões de neurônios, onde cada neurônio é uma casa enviando pequenas "mensagens de texto" elétricas (chamadas de picos) para milhares de outras casas a cada segundo. Para simular isso, você precisa de um supercomputador com milhares de placas gráficas (GPUs) trabalhando juntas.

O problema é que essas GPUs são como ilhas. Elas são rápidas, mas não conversam entre si facilmente. Se uma ilha quiser enviar uma mensagem para outra, o "carteiro" (o sistema de comunicação) tem que correr de um lado para o outro, o que deixa tudo mais lento.

Este artigo apresenta uma nova maneira, muito mais rápida, de construir o mapa dessas conexões antes da simulação começar, para que as GPUs possam executar a simulação sem ficar presas no trânsito.

Veja como eles fizeram isso, explicado de forma simples:

1. A Maneira Antiga: Construindo o Mapa no Continente

Anteriormente, quando cientistas queriam simular uma rede cerebral, eles construíam o "mapa de conexões" no computador lento e central (a CPU) primeiro. Depois, tinham que copiar esse mapa massivo para as GPUs rápidas.

• A Analogia: Imagine que você está organizando uma festa massiva. No método antigo, você escrevia o nome de cada convidado e quem eles conheciam em um pedaço de papel na cozinha (CPU), depois corria para cada sala (GPU) para entregar uma cópia da lista. Isso levava muito tempo apenas para se preparar.

2. A Maneira Nova: Construindo o Mapa Dentro das Salas

Os autores desenvolveram um novo método onde cada GPU constrói sua própria parte do mapa de conexões diretamente em sua própria memória, sem esperar pelo computador central.

• A Analogia: Agora, em vez de escrever a lista na cozinha, cada sala tem seu próprio bloco de notas. Assim que a festa começa, os convidados em cada sala escrevem quem eles conhecem ali mesmo. Não é necessário correr de um lado para o outro até a cozinha.
• O Resultado: Essa construção "a bordo" é mais de 10 vezes mais rápida do que a maneira antiga. Em um teste, levou 55 segundos para construir a rede em vez de quase 12 minutos.

3. Duas Maneiras de Enviar Mensagens

Uma vez que o mapa é construído, as GPUs precisam trocar as "mensagens de texto" (picos) durante a simulação. O artigo testou duas estratégias diferentes para isso, dependendo de como a rede é organizada:

• Estratégia A: A Ligação Telefônica Direta (Ponto a Ponto)

  • Como funciona: Se um neurônio na GPU #1 precisa falar com um neurônio específico na GPU #2, ele liga diretamente para aquela GPU específica.
  • Melhor para: Redes onde as conexões são desiguais ou específicas (como um cérebro real onde algumas áreas conversam muito entre si, mas não com todos).
  • A Alegação do Artigo: Eles usaram isso para um modelo do córtex visual de macaco (32 áreas diferentes). Funcionou perfeitamente, provando que o novo método de construção de mapas é compatível com estruturas cerebrais complexas e do mundo real.

• Estratégia B: O Chat em Grupo (Comunicação Coletiva)

  • Como funciona: Em vez de ligar para indivíduos, uma GPU grita suas mensagens para um grupo inteiro de GPUs de uma vez. Todos no grupo ouvem o grito e verificam se a mensagem é para eles.
  • Melhor para: Redes enormes e aleatórias onde todos conversam com todos (como uma multidão equilibrada).
  • A Alegação do Artigo: Eles testaram isso em uma enorme "rede balanceada" escalando até 1.024 GPUs. Este é um número enorme de placas gráficas trabalhando juntas. Eles mostraram que, mesmo com tantas placas, o sistema escala suavemente sem travar.

4. O Truque dos "Níveis de Memória"

As GPUs têm muita memória, mas não infinita. Armazenar os mapas de conexão para bilhões de neurônios ocupa muito espaço.

• A Analogia: Imagine que você tem uma pequena mesa (memória da GPU) e um enorme armazém (memória da CPU).
• A Solução: Os autores criaram quatro "níveis" de organização.
  • Nível 0: Mantenha os mapas no armazém (CPU) e traga apenas o que você precisa para a mesa. Isso economiza espaço na mesa, mas é mais lento para buscar.
  • Nível 3: Encha a mesa com tudo. Isso é o mais rápido, mas requer uma mesa maior.
• A Alegação do Artigo: Eles mostraram que, escolhendo o nível certo, podiam executar simulações no supercomputador Leonardo Booster (que possui 4.096 GPUs) e até prever que o próximo supercomputador JUPITER poderia simular uma rede com 230 milhões de neurônios e 2,5 trilhões de sinapses. Isso é aproximadamente o tamanho do córtex humano!

Resumo do Que Eles Conquistaram

• Velocidade: Eles tornaram a fase de "configuração" das simulações cerebrais 10 vezes mais rápida construindo o mapa da rede diretamente nas placas gráficas.
• Escala: Eles provaram que isso funciona em até 1.024 GPUs simultaneamente.
• Flexibilidade: Eles mostraram duas maneiras diferentes de lidar com a comunicação (ligações diretas vs. chats em grupo) para que os cientistas possam escolher o melhor método para seu modelo cerebral específico.
• Preparado para o Futuro: Seus métodos são projetados para funcionar na próxima geração de supercomputadores "Exaescala", que serão poderosos o suficiente para simular um cérebro humano completo com detalhes de sinapses individuais.

Em resumo, eles não apenas fizeram a simulação rodar mais rápido; eles construíram um melhor "sistema de estradas" para os dados, para que o supercomputador não fique preso no trânsito antes mesmo da corrida começar.

Resumo técnico

Resumo Técnico: Construção Escalável de Redes Neurais de Spikes usando até milhares de GPUs

Declaração do Problema

Simular Redes Neurais de Spikes (SNNs) em grande escala, na magnitude do córtex cerebral humano, apresenta dois desafios principais: requisitos substanciais de memória para neurônios e sinapses individuais, e a necessidade de altas velocidades de processamento para resolver dinâmicas com precisão sub-milissegundo. Embora sistemas de Computação de Alto Desempenho (HPC) equipados com milhares de GPUs ofereçam a densidade computacional necessária, o software de simulação baseado em GPUs existente ainda não demonstrou capacidade de escalar para clusters de computação inteiros, atendendo simultaneamente às demandas de infraestrutura e precisão da neurociência computacional.

Um gargalo específico nas simulações distribuídas de redes de neurônios pontuais em grande escala é a comunicação de spikes entre diferentes nós de um cluster de computação. Abordagens anteriores, como Digital Brain ou GeNN, ou omitiam informações de sinapses individuais ou estavam limitadas à execução em uma única GPU. Além disso, simuladores tradicionais baseados em CPU, como o NEST, dependem de distribuição de neurônios em rodízio e comunicação coletiva, o que assume estruturas de rede homogêneas e falha em explorar a heterogeneidade topológica e espacial dos cérebros biológicos. Embora o NEST GPU tenha abordado algumas dessas questões, sua construção inicial de rede dependia da transferência de dados da memória da CPU para a memória da GPU, e métodos de construção dinâmica estavam anteriormente limitados a simulações de GPU única.

Metodologia

Este trabalho apresenta um método inovador e eficiente em termos de memória para construir e simular SNNs em grande escala diretamente em sistemas multi-GPU usando a Interface de Passagem de Mensagens (MPI). A inovação central reside na realização da construção da rede inteiramente dentro da memória da GPU ("onboard"), sem comunicação entre processos durante a fase de construção.

Algoritmo Central

O método distingue entre conexões locais (neurônios dentro do mesmo processo MPI) e conexões remotas (neurônios entre processos diferentes).

1. Construção Independente: Cada processo MPI constrói independentemente sua porção da rede. Ele cria conectividade local e prepara estruturas de dados para conexões remotas sem comunicar-se com outros processos.
2. Representações Proxy: Para conexões remotas, o método utiliza "neurônios imagem" (proxies) nos processos de destino. Estas são representações virtuais de neurônios fonte localizados em outros ranks MPI.
3. Mapas de Comunicação: O algoritmo instancia mapas de comunicação contíguos na memória da GPU para rotear spikes eficientemente. Esses mapas associam o índice de um neurônio fonte em um rank de origem ao índice de seu neurônio imagem em um rank de destino.
4. Esquemas de Comunicação: O framework suporta dois modos de comunicação MPI, selecionáveis pelo usuário com base na arquitetura da rede:
   • Ponto a Ponto: Usa comunicação direta entre dois processos. É otimizado para redes com distribuições desiguais de neurônios ou sinapses (por exemplo, o Modelo de Múltiplas Áreas). Utiliza estruturas de mapeamento específicas $(R_{\tau,\sigma}, L_{\tau,\sigma})$ e sequências $(T, P)$ para rotear spikes.
   • Coletivo: Usa comunicação baseada em grupos (por exemplo, MPI_Allgather). É vantajoso para redes balanceadas com cargas de comunicação homogêneas. Emprega arrays de indexação específicos do grupo e arrays de host para gerenciar o roteamento de spikes entre múltiplos processos simultaneamente.

Otimização de Memória GPU

Para equilibrar o consumo de memória da GPU e a velocidade de simulação, os autores implementaram quatro Níveis de Memória GPU (GMLs):

• Nível 0: Mapas de conexão remota e contagens de conexão são armazenados na memória da CPU.
• Nível 1: Similar ao Nível 0, mas assume que todos os neurônios fonte possuem imagens nos processos de destino, evitando verificações de uso real (construção mais rápida, potencialmente maior desperdício de memória).
• Nível 2: Mapas e índices de conexão são armazenados na memória da GPU; contagens de conexão são computadas sob demanda. Este é o nível padrão.
• Nível 3: Todas as estruturas de dados, incluindo contagens de conexão, são armazenadas na memória da GPU, minimizando a transferência de dados CPU-GPU ao custo de maior uso de memória da GPU.

Modelos Avaliados

1. Modelo de Múltiplas Áreas (MAM): Um modelo biologicamente detalhado de 32 áreas relacionadas à visão do córtex de macaco-prego ($4,13 \times 10^6$ neurônios, $24,2 \times 10^9$ sinapses). Este modelo apresenta conectividade complexa e hierárquica e foi simulado usando comunicação ponto a ponto.
2. Rede Balanceada Escalável: Uma rede aleatória de neurônios excitatórios e inibitórios com conectividade de grau de entrada fixo, projetada para avaliar o desempenho de escalabilidade fraca. Este modelo foi simulado usando comunicação coletiva em até 1.024 GPUs.

Principais Resultados

Desempenho de Construção de Rede

O método de construção "onboard" na GPU demonstrou acelerações significativas em comparação com a abordagem anterior "offboard" (baseada em CPU):

• Simulação MAM: O tempo de construção da rede diminuiu de 686,0 s (offboard) para 55,5 s (onboard), uma aceleração de 12,4x.
  • A criação de conexões locais teve uma aceleração de 20x.
  • A criação de conexões remotas teve uma aceleração de 9x.
  • A criação de neurônios/dispositivo e preparação de simulação tiveram acelerações de 350x e 50x, respectivamente.
• Rede Balanceada Escalável: O método construiu com sucesso redes de até 230,4 milhões de neurônios e 2,59 trilhões de sinapses em 1.024 GPUs (256 nós).

Propagação de Estado e Escalabilidade

• MAM: O tempo de propagação de estado (medido como Fator de Tempo Real) permaneceu comparável entre as versões offboard e onboard (aproximadamente 15–16), indicando que a otimização de construção não impactou negativamente as dinâmicas de simulação.
• Rede Balanceada: O sistema demonstrou escalabilidade fraca até 1.024 GPUs.
  • Eficiência de Memória: O Nível de Memória GPU 0 permitiu que simulações alcançassem 4.096 nós sem exceder os limites de memória das GPUs NVIDIA A100 (64 GB). Níveis de memória mais altos (2 e 3) ofereceram velocidades de construção e simulação mais rápidas, mas atingiram o limite de memória em contagens de nós mais baixas (aproximadamente 3.072 nós para o Nível 3).
  • Desempenho: Desativar a gravação de spikes na rede balanceada reduziu o tempo de propagação de estado em aproximadamente 20%.

Validação

O novo método de construção foi validado contra a versão offboard anterior e o simulador baseado em CPU NEST. Apesar das mudanças nas sequências de geração de números aleatórios devido ao novo algoritmo, as propriedades estatísticas da atividade de spikes (taxas de disparo, coeficiente de variação dos intervalos entre spikes e correlações de Pearson pareadas) foram preservadas, confirmando a validade biológica da simulação.

Significado e Alegações

O artigo afirma que este trabalho fornece o primeiro software de simulação de SNN baseado em GPU capaz de escalar para clusters de computação inteiros (até milhares de GPUs) enquanto armazena informações de sinapses individuais. As principais contribuições são:

1. Construção Escalável: Um algoritmo inovador que constrói conectividade de rede diretamente na memória da GPU, eliminando o gargalo de transferência CPU-GPU e evitando comunicação MPI durante a fase de construção.
2. Flexibilidade: Suporte para comunicação MPI ponto a ponto e coletiva, permitindo adaptação a diferentes topologias de rede (hierárquica vs. aleatória/balanceada).
3. Preparação para Exaescala: Os autores extrapolam que sua abordagem poderia simular redes de $2 \times 10^{10}$ neurônios e $10^{14}$ sinapses no futuro supercomputador exaescala JUPITER. Esta escala aproxima-se da conectividade do córtex humano, mantendo resolução de sinapses individuais.
4. Eficiência: Ao otimizar o uso de memória através do sistema GML, o método permite a simulação de redes maiores em hardware existente (por exemplo, ajustando o MAM em 8 GPUs em vez de 32) e fornece um caminho para utilizar a capacidade total de futuros sistemas exaescala.

Os autores concluem que esta abordagem aborda o gargalo crítico da comunicação de spikes em simulações distribuídas e estabelece o NEST GPU como uma plataforma de referência para simulações neurais em grande escala e biologicamente detalhadas em arquiteturas modernas de HPC.