Sistema Cerebras CS-1 Integrado ao Supercomputador Lassen

Um novo estudo de caso realizado pela Cerebras em parceria com Laboratório Nacional Lawrence Livermore (LLNL) detalha como o sistema Cerebras CS-1 foi integrado ao supercomputador Lassen do LLNL para permitir avanços em simulações de fusão nuclear. 

O LLNL é uma instalação de pesquisa federal em Livermore, Califórnia, e é financiado principalmente pelo Departamento de Energia dos EUA, Administração Nacional de Segurança Nuclear (NNSA). De acordo com o LLNL, sua missão é fortalecer a segurança dos EUA desenvolvendo e aplicando ciência, tecnologia e engenharia de classe mundial. 

O laboratório abriga a Instalação Nacional de Ignição (NIF), que realiza pesquisas de fusão nuclear com o laser mais poderoso do mundo. Com isso, alguns dos principais obstáculos incluem experimentos de confinamento inercial caros e demorados, então o laboratório executa experimentos simulados com um pacote de software de multi-física chamado HYDRA no supercomputador Lassen. Os modelos HYDRA são validados por meio de dados do mundo real do NIF, o que permite que os modelos sejam mais precisos na previsão do resultado de experimentos do mundo real. 

Parte dos modelos HYDRA inclui cinética atômica e radiação, e essa parte é chamada de CRETIN. Ele prevê como um átomo se comportará sob certas condições, e o CRETIN pode representar dezenas de por cento da carga de computação total para o HYDRA.

Substituindo o CRETIN por um modelo de rede neural profunda (DNN), ou o CRETIN-surrogate, os pesquisadores do LLNL podem reduzir a intensidade computacional. 

Sistema Cerebras CS-1

O sistema Cerebras CS-1 foi escolhido pelo LLNL para realizar a inferência do CRETIN-surrogate. O sistema foi integrado ao supercomputador Lassen, e a instalação levou menos de 20 horas. Técnicos da Cerebras também instalaram uma “casca de refrigeração” e os trilhos e hardware de suporte mecânico. 

Engenheiros de software de aprendizado de máquina trabalharam com colegas do LLNL para escrever uma API em C++ que permite que o código HYDRA chame o modelo CRETIN-surrogate. O modelo depende de um autoencoder para comprimir os dados de entrada em representações de dimensão mais baixa, e esses são processados por um modelo preditivo construído com DJINN, que é um algoritmo de rede neural profunda novo. Esse algoritmo escolhe automaticamente uma arquitetura de rede neural apropriada para os dados dados, e não requer que o usuário ajuste manualmente as configurações.

Resultados do Estudo de Caso

Os resultados iniciais demonstraram que a combinação do sistema Lassen com o acelerador Cerebras é extremamente eficiente. Ao conectar o sistema CS-1 à rede InfiniBand do Lassen, foi possível alcançar uma largura de banda de 1,2 terabits por segundo para o sistema CS-1.

Devido aos 19 GB de memória SRAM acoplados a 400.000 núcleos de computação de IA, o sistema CS-1 foi capaz de executar muitas instâncias do modelo DNN relativamente compacto em paralelo. Por meio da combinação de largura de banda e potência, o HYDRA foi capaz de realizar inferência em 18 milhões de amostras por segundo. 

Isso significa que o LLNL agora pode executar experimentos que eram anteriormente inacessíveis computacionalmente com o sistema Cerebras, e isso envolve apenas uma integração simples e uma fração do custo. 

A pesquisa agora se concentrará em direcionar a simulação e fornecer insights sobre a simulação enquanto ela está em execução, o que permite que os pesquisadores monitorem e interrompam a execução se a simulação não estiver funcionando bem. Os resultados de cada execução se tornam parte do conjunto de treinamento do modelo, então ele pode ser treinado continuamente. Um modelo de “aprendizado ativo” pode ser criado, e ele pode otimizar as execuções futuras escolhendo os parâmetros e a condicionamento de limite inicial para o próximo experimento.