Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

Este trabalho apresenta um fluxo de trabalho exaescala baseado em modelos fundamentais de grafos atômicos (HydraGNN) que, ao ser treinado em 544 milhões de estruturas em supercomputadores como o Frontier, permite a triagem de 1,1 bilhão de estruturas em 50 segundos, acelerando drasticamente a descoberta de materiais em comparação com métodos de primeiros princípios.

O essencial

Imagine que você é um chef de cozinha tentando descobrir a receita perfeita para um novo prato. No mundo da ciência de materiais, os "ingredientes" são átomos e as "receitas" são como esses átomos se organizam para criar coisas como baterias melhores, medicamentos mais eficazes ou materiais superfortes.

Até agora, descobrir essas receitas era como tentar cozinhar um banquete para milhões de pessoas apenas usando uma colher de pau e uma panela pequena. Os cientistas usavam métodos tradicionais (chamados de "primeiros princípios") que eram extremamente precisos, mas lentos e caros. Para testar uma única combinação de átomos, podia levar dias de cálculo em supercomputadores. Testar bilhões de combinações? Isso levaria séculos.

Este artigo apresenta uma revolução: um "Super-Chef de Inteligência Artificial" capaz de cozinhar (ou seja, prever) bilhões de receitas em apenas 50 segundos.

Aqui está a explicação do que eles fizeram, usando analogias do dia a dia:

1. O Problema: A Cozinha Bagunçada

Os cientistas tinham 16 "livros de receitas" diferentes (conjuntos de dados). O problema é que:

• Alguns livros tinham milhões de receitas, outros tinham apenas algumas.
• Alguns livros usavam medidas em xícaras, outros em gramas (diferentes níveis de precisão ou "fidelidade").
• Se você misturasse tudo em uma panela gigante sem cuidado, as receitas grandes dominariam e as pequenas (mas importantes) seriam esquecidas. O resultado seria um prato ruim.

2. A Solução: O "Chef Multi-Tarefa" (HydraGNN)

Os pesquisadores criaram um modelo de IA chamado HydraGNN. Pense nele como um chefe de cozinha com várias mãos especializadas:

• O Corpo do Chef (A Base): Ele aprende as regras básicas da culinária (como os átomos se conectam) usando todos os livros de receitas ao mesmo tempo.
• As Mãos Especiais (Cabeças de Tarefa): Para cada livro de receitas diferente, ele tem uma "mão" especializada que ajusta o sabor para aquele livro específico. Isso permite que ele aprenda com o livro pequeno sem ser sufocado pelo livro gigante.

3. A Máquina: O Supercomputador "Frontier"

Para treinar esse chef, eles usaram o Frontier, um dos supercomputadores mais rápidos do mundo (um computador "Exascale").

• Eles usaram 16.384 chips de processamento (GPUs) trabalhando juntos. É como ter 16.000 cozinheiros ajudando o chefe a ler os livros de receitas simultaneamente.
• Eles treinaram o modelo com 544 milhões de estruturas atômicas (imagens de como os átomos estão organizados).

4. A Grande Proeza: O Teste de Velocidade

Aqui está a parte mais impressionante. Depois de treinado, eles pediram ao Super-Chef para examinar 1,1 bilhão de estruturas atômicas.

• O método antigo: Levaria cerca de 6,7 anos de trabalho contínuo em supercomputadores para fazer isso.
• O método novo: Feito em 50 segundos.
• Analogia: É como se você pudesse ler todos os livros da Biblioteca do Congresso dos EUA em menos de um minuto, entendendo perfeitamente cada palavra.

5. Por que isso é importante?

• Descoberta Rápida: Em vez de esperar anos para encontrar um novo material para uma bateria de carro elétrico, os cientistas podem agora "peneirar" trilhões de possibilidades em segundos e encontrar as joias raras.
• Aprendizado com Poucos Dados: O modelo é tão inteligente que, mesmo quando precisam aprender uma tarefa nova com poucos exemplos (como prever algo sobre um material muito raro), ele se adapta rapidamente, porque já "leu" tantos livros antes.
• Precisão: Eles provaram que o modelo funciona bem mesmo usando diferentes níveis de "precisão" (como usar uma régua de milímetro vs. uma trena), garantindo que as previsões sejam confiáveis para a ciência real.

Resumo da Ópera

Esta pesquisa criou um sistema de "super-leitura" para o mundo atômico. Eles transformaram um processo que era lento, caro e limitado em algo rápido, barato e capaz de explorar o universo de possibilidades químicas que antes era invisível para a ciência.

É como se a humanidade tivesse acabado de ganhar um telescópio que não apenas vê estrelas distantes, mas que consegue mapear a atmosfera de cada planeta em um segundo, permitindo-nos encontrar novos mundos habitáveis instantaneamente.

Resumo técnico

Título: Modelos Fundamentais de Grafos Atômicos Multi-Tarefa em Escala Exascale para Dados Atômicos Desequilibrados e Multi-Fidelidade

1. O Problema

A descoberta de materiais e a modelagem atômica precisas são fundamentais para avanços em química e ciência dos materiais. No entanto, os métodos tradicionais de primeiros princípios, como a Teoria do Funcional da Densidade (DFT), são computacionalmente proibitivos para a triagem em larga escala de vastos espaços químicos.

• Desafios Específicos:
  • Dados Heterogêneos: Não existe um único conjunto de dados ou configuração eletrônica que seja confiável uniformemente para sistemas orgânicos, inorgânicos e híbridos.
  • Desequilíbrio e Multi-Fidelidade: Os conjuntos de dados disponíveis variam drasticamente em tamanho (de milhares a centenas de milhões de estruturas) e em fidelidade (diferentes teorias de aproximação, funcionais e bases). Treinar modelos em dados desbalanceados leva a instabilidade na otimização e viés para os domínios dominantes.
  • Erro Numérico: Modelos de potencial interatômico de aprendizado de máquina (MLIPs) dependem de gradientes de energia (forças), onde erros numéricos são amplificados, especialmente em precisões reduzidas (BF16/FP32) e em ambientes distribuídos.
  • Escalabilidade: A necessidade de treinar modelos fundamentais (Foundation Models) em supercomputadores líderes (Exascale) exige arquiteturas que suportem bilhões de parâmetros e dados sem perder eficiência.

2. Metodologia

Os autores apresentaram um workflow completo em escala exascale baseado no HydraGNN, integrando infraestrutura de dados, otimização de hiperparâmetros e aprendizado multi-tarefa.

• Coleta e Pré-processamento de Dados:

  • Treinamento conjunto em 16 conjuntos de dados abertos de primeiros princípios, totalizando 544+ milhões de estruturas cobrindo mais de 85 elementos.
  • Os dados incluem moléculas orgânicas, cristais inorgânicos, ligas e polímeros, gerados com diferentes teorias de aproximação (DFT com diversos funcionais).
  • Uso de ADIOS2 e DDStore para criar um pipeline de dados escalável, utilizando armazenamento NVMe local para staging e reduzindo a latência de leitura em sistemas de arquivos compartilhados.

• Arquitetura Multi-Tarefa (MTL):

  • Utilização de uma arquitetura com camadas de passagem de mensagens (message-passing) compartilhadas para aprender características atômicas fundamentais transferíveis.
  • Heads (Cabeças) Específicas por Conjunto de Dados: Cada fonte de dados possui sua própria camada de saída (head) para absorver características dependentes do dataset e offsets de fidelidade, reduzindo a interferência de gradientes entre tarefas desbalanceadas.
  • Ponderação Condicional à Composição: Um MLP compacto aprende a ponderar as saídas das 16 branches (ramificações) em tempo de inferência, com base na composição química da estrutura, permitindo que o modelo generalize para composições não vistas durante o treinamento.

• Seleção de Modelo e Otimização (HPO):

  • Execução de 6 campanhas massivas de Otimização de Hiperparâmetros (HPO) no supercomputador Frontier (16.384 GPUs), testando 6 arquiteturas de redes neurais de passagem de mensagens (MPNN): EGNN, SchNet, DimeNet, MACE, PaiNN e PNAEq.
  • Uso do DeepHyper para explorar o espaço de hiperparâmetros, selecionando modelos que equilibram precisão, tempo de treinamento e custo computacional.
  • O modelo vencedor (baseado em PaiNN) foi promovido para treinamento contínuo em 2.048 nós (16.384 GPUs).

• Análise de Precisão:

  • Caracterização sistemática do impacto das precisões FP64, FP32 e BF16 tanto no treinamento quanto na inferência, quantificando o compromisso entre velocidade e precisão científica.

3. Principais Contribuições

• Primeiro Workflow Multi-Tarefa em Escala Exascale: Demonstração prática de treinamento conjunto em 544 milhões de estruturas heterogêneas, superando o desafio do desbalanceamento de dados através de cabeças de tarefa desacopladas.
• Pipeline de Dados Co-Design: Integração de ADIOS2/DDStore com NVMe local para sustentar o treinamento contínuo em 2.048 nós sem gargalos de I/O.
• Seleção de Arquitetura Guiada por HPO em Escala: Elevação da seleção de arquitetura a um problema de otimização em escala exascale, identificando que o PaiNN era o mais eficiente para o orçamento fixo de tempo/nós.
• Inferência em Bilhões de Estruturas: Desenvolvimento de um pipeline de inferência otimizado (com fused gradient, reutilização de codificador e torch.compile) capaz de avaliar 1,1 bilhão de estruturas em 50 segundos.
• Portabilidade Cross-Platform: Validação do workflow em três supercomputadores líderes (Frontier, Aurora e Perlmutter), demonstrando escalabilidade forte e fraca consistente.

4. Resultados

• Desempenho de Triagem (Inferência):

  • O modelo final avaliou 1,1 bilhão de estruturas atômicas em apenas 50 segundos no Frontier.
  • Isso representa uma aceleração de 33x em relação a abordagens anteriores e transforma uma tarefa que levaria 6,7 anos em DFT (em escala similar) em uma tarefa de segundos.
  • Throughput sustentado: ~21,8 milhões de estruturas/segundo no sistema completo.

• Ajuste Fino (Fine-Tuning) e Transferência:

  • O modelo pré-treinado foi ajustado para 12 tarefas downstream diversas (incluindo QM9, MD17, OQMD, etc.).
  • Para tarefas alinhadas à Superfície de Energia Potencial (PES), o fine-tuning com o backbone desbloqueado reduziu o erro de validação em até uma ordem de magnitude comparado ao treinamento do zero ou com backbone congelado.
  • O modelo demonstrou robustez em regimes de dados escassos (ex: 150 estruturas), provando a eficácia da representação fundamental aprendida.

• Escalabilidade:

  • Escalabilidade forte quase linear observada até 2.048 GPUs no Perlmutter, 6.144 no Aurora e 1.024 no Frontier.
  • O modelo final possui apenas 12,1 milhões de parâmetros e ocupa ~92 MB de memória, permitindo implantação em edge e integração em fluxos de trabalho de física complexa.

• Precisão vs. Desempenho:

  • A precisão FP64 foi mantida para garantir reprodutibilidade científica.
  • O uso de FP32 introduziu um desvio de energia constante de ~0,021 eV, enquanto BF16 resultou em ~0,080 eV, fornecendo métricas claras para decisões de implantação.

5. Significado e Impacto

Este trabalho representa um marco na interseção entre aprendizado de máquina científico e computação de alto desempenho (HPC).

• Aceleração da Descoberta de Materiais: Transforma a triagem de espaços químicos vastos de um gargalo combinatório proibitivo em um fluxo de trabalho viável e rápido.
• Validação de Modelos Fundamentais Científicos: Demonstra que modelos fundamentais treinados em escala exascale não são apenas grandes, mas são ferramentas práticas que suportam fine-tuning eficiente em dados e inferência de alto rendimento.
• Reprodutibilidade e Portabilidade: Estabelece um padrão para workflows portáveis entre diferentes arquiteturas de aceleradores (AMD MI250/MI350, NVIDIA H100/A100), essencial para a comunidade científica global.
• Abordagem para Dados Desbalanceados: A estratégia de Multi-Task Learning com cabeças específicas oferece uma solução robusta para a integração de dados científicos heterogêneos, um problema comum em ciências computacionais.

Em resumo, os autores criaram um ecossistema completo que permite explorar espaços de design químico que eram anteriormente inacessíveis aos métodos de primeiros princípios, democratizando o acesso a simulações atômicas de alta fidelidade em escala industrial.