A Survey of Neural Network Variational Monte Carlo from a Computing Workload Characterization Perspective

Este artigo apresenta uma análise de carga de trabalho e caracterização empírica em GPUs de quatro ansätze representativos de Monte Carlo Variacional com Redes Neurais, identificando gargalos de desempenho relacionados a baixa intensidade computacional e movimentação de dados, e discutindo implicações para o co-design de algoritmos e hardware visando sistemas escaláveis.

O essencial

Imagine que você está tentando prever o tempo perfeito para uma festa, mas em vez de nuvens e vento, você está lidando com elétrons (partículas minúsculas) e átomos. O problema é que esses elétrons são como uma multidão de pessoas muito agitadas: eles se movem rápido, interagem entre si e com os "anfitriões" (os núcleos dos átomos), e prever exatamente onde eles estarão é um pesadelo matemático.

Aqui entra o NNVMC (Monte Carlo Variacional com Redes Neurais). Pense nele como um super-inteligente detetive que usa uma rede neural (uma espécie de cérebro de computador) para adivinhar onde os elétrons estão e como se comportam.

Este artigo é como um relatório de engenharia que olha para dentro desse "cérebro" para entender por que ele às vezes é lento e gasta muita energia, mesmo com computadores modernos superpotentes (chamados GPUs).

Aqui está a explicação simplificada, passo a passo:

1. O Problema: O Detetive está "Engasgado"

Os cientistas criaram quatro versões diferentes desse "detetive" (chamadas FermiNet, PauliNet, Psiformer e Orbformer). A ideia era que eles fossem rápidos e precisos. Mas, na prática, eles estão travando.

• A Analogia: Imagine que você tem um caminhão de entrega (o computador) que pode levar 100 caixas de uma vez. Mas, em vez de encher o caminhão com caixas grandes e pesadas (cálculos complexos), o motorista passa 90% do tempo fazendo curtos trajetos para pegar apenas um grão de areia de cada vez.
• O que o artigo descobriu: O computador está gastando mais tempo movendo dados (pegando informações da memória) do que fazendo cálculos (pensando). É como se o motor do carro fosse potente, mas o pneu estivesse furado e o carro estivesse arrastando o chassi.

2. As Quatro Versões do Detetive (Os Modelos)

Os autores testaram quatro "estilos" de detetive e viram que cada um tem um problema diferente:

• FermiNet e PauliNet (Os Tradicionais): Eles são como chefs que seguem receitas muito rígidas. Eles fazem muitos cálculos repetitivos e pequenos. O problema é que eles ficam presos em "trânsito" (lentidão de memória) porque precisam verificar e re-verificar os dados muitas vezes para garantir que a física está correta.
• Psiformer (O Moderno): Ele usa uma técnica mais inteligente (Transformers, a mesma usada no ChatGPT). Ele consegue fazer cálculos maiores de uma vez só, o que é ótimo. Mas, ele gasta muito tempo na etapa de "amostragem" (tentar adivinhar onde os elétrons estão), o que ainda deixa o computador esperando.
• Orbformer (O Híbrido): Ele tenta ser o melhor dos dois mundos, mas acabou ficando com um pouco de tudo. Ele usa menos cálculos grandes e mais pequenos movimentos de dados, voltando a sofrer com a lentidão de memória.

3. O Grande Segredo: Não é só sobre "Força Bruta"

Muitas pessoas pensam que para acelerar esses sistemas, precisamos apenas de computadores mais rápidos (mais "força bruta").

• A Analogia: Pense em uma cozinha de restaurante. Se você tem um chef que corta legumes muito rápido (cálculo), mas o ajudante que traz os legumes da despensa é lento (memória), o chef vai ficar parado esperando.
• A Conclusão do Artigo: Adicionar mais força bruta não ajuda muito. O que precisamos é melhorar a logística. Precisamos de ajudantes que tragam os legumes mais perto do chef, ou de uma despensa que fique dentro da cozinha.

4. As Soluções Sugeridas (O Futuro)

Os autores sugerem que, para fazer essa tecnologia funcionar de verdade, os engenheiros de hardware e os cientistas precisam trabalhar juntos (co-design):

1. Memória Inteligente (PIM): Em vez de mover os dados da memória para o processador, coloque o processador dentro da memória. É como ter o chef cozinhando direto na despensa, sem precisar correr até lá.
2. Equipes Híbridas: Usar diferentes tipos de computadores para tarefas diferentes. O computador forte faz os cálculos grandes, e um sistema mais simples e rápido cuida dos pequenos movimentos de dados.
3. Adaptação Dinâmica: O sistema deve mudar de estratégia dependendo da fase do trabalho. Às vezes ele precisa de velocidade de cálculo, às vezes de velocidade de transporte de dados. Um carro que muda de pneu automaticamente dependendo do terreno.
4. Usar a Memória do Celular/SSD: Para sistemas muito grandes, a memória do computador gráfico não é suficiente. O artigo sugere usar a memória do processador principal (CPU) ou até discos rígidos para guardar dados que não estão sendo usados naquele milissegundo, liberando espaço para o que é importante agora.

Resumo Final

Este artigo diz: "Pare de tentar apenas tornar o motor mais potente. O problema é que o carro está atolado no barro porque estamos movendo os dados de um jeito ineficiente."

Para resolver os problemas da química quântica e descobrir novos materiais ou medicamentos, precisamos redesenhar como os computadores lidam com esses dados específicos, criando sistemas que entendam que, às vezes, mover menos informação é mais rápido do que calcular mais.

Resumo técnico

Resumo Técnico: Caracterização de Carga de Trabalho em NNVMC

1. Problema e Contexto

A equação de Schrödinger eletrônica é fundamental para a simulação de muitos corpos quânticos na química e ciência dos materiais. Métodos tradicionais (como CCSD(T) e FCI) enfrentam um compromisso difícil entre precisão e custo computacional (escalando frequentemente como $O(N^7)$). O Variational Monte Carlo com Redes Neurais (NNVMC) emergiu como uma alternativa promissora, utilizando redes neurais para parametrizar a função de onda, oferecendo melhor escalabilidade assintótica (frequentemente citada como $O(N^4)$).

No entanto, a implementação prática do NNVMC em GPUs modernas enfrenta barreiras significativas:

• Custos Elevados: Alto tempo de execução e consumo de memória limitam a escalabilidade para sistemas maiores.
• Ineficiência de Desempenho: Diferente de cargas de trabalho de linguagem ou visão computacional, o NNVMC possui estágios de execução específicos da física (amostragem MCMC, construção de função de onda, avaliação de derivadas/Laplaciano).
• Falha de Métricas Tradicionais: Contar apenas operações de ponto flutuante (FLOPs) não é um bom preditor de tempo de execução ou comportamento de memória, devido à heterogeneidade dos kernels e gargalos de transferência de dados.

2. Metodologia

Os autores realizaram uma análise de carga de trabalho orientada a dados, combinando uma revisão de literatura com uma caracterização empírica em GPU.

• Modelos Analisados: Quatro ansatzes representativos para problemas de estado fundamental:
  • PauliNet e FermiNet (implementados no codebase DEEPQMC).
  • Psiformer e Orbformer (implementados no codebase ONEQMC).
• Plataforma de Hardware: GPUs NVIDIA (RTX A5000, A100, H200).
• Configuração de Carga de Trabalho: Moléculas variando de pequenas (LiH) a médias (C4H4), com 1024 configurações de elétrons (walkers) em paralelo.
• Ferramentas de Perfilamento:
  • Nsight Systems para rastreamento de execução.
  • Nsight Compute para métricas em nível de kernel (intensidade aritmética, uso de SM, taxa de acerto de cache L2, throughput de memória).
• Abordagem de Análise:
  • Divisão da execução em estágios (A: Embedding, B: Propagação/Message Passing, C: Readout, D: Montagem da Função de Onda, E: Avaliação de Derivadas/Laplaciano).
  • Uso do Modelo Roofline e Intensidade Aritmética (AI) para classificar os kernels como limitados por memória ou por computação.

3. Principais Contribuições

1. Revisão Orientada à Carga de Trabalho: Mapeamento unificado dos fluxos de execução dos quatro ansatzes, destacando como as escolhas arquiteturais (ex: uso de JVP vs. estimadores de Laplaciano estilo Hutchinson) alteram a composição da carga de trabalho.
2. Caracterização em Nível de Kernel e Operador: Demonstração empírica de que kernels elementwise fundidos e de movimentação de dados (baixa intensidade aritmética) frequentemente dominam o tempo de execução, superando até mesmo os kernels de multiplicação de matrizes densas (GEMM).
3. Análise de Comportamento de Hardware: Relato detalhado sobre o uso de Streaming Multiprocessors (SM), atividade Tensor Core, throughput de memória e taxas de acerto de cache L2, conectando essas métricas ao desempenho end-to-end.
4. Diretrizes de Co-design: Proposição de direções para o co-design algoritmo-hardware, focando em otimizações sensíveis à fase de execução e centradas na memória.

4. Resultados Chave

• Escalonamento de Tempo de Execução:

  • O tempo de execução escala fortemente com o tamanho da molécula, mas a taxa de crescimento depende do ansatz.
  • PauliNet/FermiNet: Escalam de forma íngreme (30-42x) devido à estratégia de Laplaciano baseada em JVP (Jacobian-Vector Product), que exige reexecução dos estágios A-D para cada direção cartesiana, gerando muitos kernels finos e sensíveis à largura de banda.
  • Psiformer/Orbformer: Escalam de forma mais suave (8-9x) devido ao uso de estimadores de Laplaciano estilo Hutchinson (menos reexecuções), deslocando o tempo para a fase de amostragem e kernels GEMM maiores.

• Composição de Kernels e Gargalos:

  • Dominância de Baixa Intensidade: Em todos os modelos, kernels elementwise fundidos e de layout (transpose, reshape) ocupam uma grande fatia do tempo de execução (ex: 52% no PauliNet) e operam na região limitada por memória do modelo Roofline.
  • Fase E (Derivadas): A avaliação do Laplaciano é um gargalo crítico. No PauliNet/FermiNet, ela amplifica a granularidade dos kernels, mantendo o sistema limitado por memória mesmo com a presença de grandes multiplicações de matrizes.
  • Diferenças entre Modelos:
    • FermiNet desloca mais tempo para GEMM em comparação ao PauliNet, mas ainda é sensível à memória.
    • Psiformer tem uma carga mais pesada em GEMM e na fase de amostragem.
    • Orbformer (com FlashAttention) reduz a dominância do GEMM, mas aumenta a proporção de kernels de movimentação de dados, tornando-se novamente fortemente limitado por memória.

• Uso de Hardware:

  • A utilização de pico de instruções e FMA (Fused Multiply-Add) é frequentemente baixa (ex: ~26% no PauliNet), indicando subutilização da capacidade de computação devido ao gargalo de transferência de dados.
  • As taxas de acerto de cache L2 variam, mas a transferência de dados off-chip permanece o fator limitante principal.

• Memória:

  • O uso de memória de treinamento é consistentemente maior que o de inferência.
  • Modelos baseados em ONEQMC (Psiformer/Orbformer) mostram saltos de memória em tamanhos específicos de moléculas devido a efeitos de alocação dinâmica e granularidade de pools de memória, além do crescimento não linear com o tamanho do sistema.

5. Significado e Implicações para Co-design

O estudo conclui que a otimização puramente focada em computação (aumentar FLOPs) ou apenas em kernels específicos (como atenção) é insuficiente para o NNVMC. As implicações para o futuro de sistemas escaláveis incluem:

1. Processamento em Memória (PIM): Devido à dominância de kernels de baixa intensidade aritmética e alta movimentação de dados, mover a computação para perto da memória (PIM) pode ser mais eficaz do que aumentar o throughput de computação bruto.
2. Sistemas Heterogêneos Colaborativos (GPU-PIM): Uma abordagem fixa (apenas GPU) não é robusta. O ideal é particionar a carga: manter os blocos pesados em GEMM na GPU e descarregar clusters limitados por memória (elementwise, layout) para engines próximas à memória.
3. Aceleração Reconfigurável: Como a mistura de operadores muda drasticamente entre estágios (ex: da propagação para a avaliação de derivadas), aceleradores reconfiguráveis que adaptam o fabric de computação e caminhos de dados para cada fase podem melhorar a eficiência.
4. Suporte Arquitetural Além da Atenção: A otimização de kernels de atenção (como FlashAttention) tem retornos decrescentes em modelos como o Orbformer, onde kernels de movimentação de dados e layout são igualmente críticos. Novas arquiteturas devem focar em padrões de acesso à memória eficientes.
5. Offload de Memória: Para sistemas grandes que excedem a memória da GPU, estratégias de offload para CPU/SSD com prefetching assíncrono são necessárias, mas exigem validação específica para o NNVMC para não destabilizar o treinamento.

Em suma, o NNVMC é uma carga de trabalho híbrida e dependente de fase, onde o gargalo principal não é a capacidade de ponto flutuante, mas sim a eficiência da movimentação de dados e a granularidade dos kernels. O co-design futuro deve priorizar estratégias sensíveis à memória e à fase de execução.