cuNNQS-SCI: A Fully GPU-Accelerated Framework for High-Performance Configuration Interaction Selection withNeural Network QQantum States

Resumo Técnico: cuNNQS-SCI

1. O Problema

O cálculo da estrutura eletrônica de sistemas quânticos de muitos corpos é fundamental para a química e física computacionais, mas enfrenta o desafio da complexidade exponencial do espaço de estados. Métodos exatos, como a Interação de Configuração Completa (FCI), tornam-se intratáveis para sistemas grandes.

A abordagem Interação de Configuração Selecionada com Estados Quânticos de Rede Neural (NNQS-SCI) surgiu como uma técnica de ponta, combinando a precisão da seleção de configurações com a capacidade representacional de redes neurais (como Transformers). No entanto, a implementação atual do NNQS-SCI sofre de gargalos críticos de escalabilidade devido a uma arquitetura híbrida CPU-GPU:

• Desduplicação Centralizada na CPU: O processo de remover configurações redundantes globalmente exige que todos os dados sejam agregados em um único nó mestre (CPU). Isso cria um gargalo de comunicação e de memória, limitando a escalabilidade e subutilizando os GPUs.
• Geração de Configurações Acopladas na CPU: A geração de novas configurações acopladas (estados vizinhos) é realizada predominantemente na CPU, criando um desequilíbrio onde a CPU se torna o fator limitante, enquanto os GPUs ficam ociosos.
• Barreira de Memória: A migração total do fluxo de trabalho para o GPU expõe a limitação de memória de alta largura de banda (HBM) dos dispositivos. Sistemas químicos grandes geram dados intermediários que excedem a capacidade de memória de um único GPU (ex: 40GB em uma A100), tornando a execução "in-core" impossível.

2. Metodologia e Arquitetura

O cuNNQS-SCI propõe um framework totalmente acelerado por GPU, redesenhando o pipeline do NNQS-SCI para eliminar a dependência da CPU e gerenciar restrições de memória. A solução é baseada em três pilares principais:

A. Desduplicação Global Distribuída (Sort-Based Regular Sampling)

• Substitui a abordagem centralizada por um algoritmo distribuído residente em GPU.
• Utiliza uma estratégia de amostragem regular baseada em ordenação (sort-based). Em vez de usar hash (que sofre com desequilíbrio de carga e acesso aleatório à memória), o algoritmo ordena localmente os dados em cada GPU, coleta amostras de pivôs para definir divisores globais e executa uma troca de dados AlltoAll eficiente.
• Isso garante um balanceamento de carga determinístico e minimiza o tráfego de comunicação, eliminando o nó mestre como gargalo.

B. Kernels CUDA de Alta Performance para Geração Acoplada

• Migra a geração de configurações acopladas (baseada nas regras de Slater-Condon) da CPU para kernels CUDA personalizados.
• Layout de Dados Otimizado: Em vez de armazenar a matriz Hamiltoniana completa (que seria exabytes), o sistema usa Tabelas de Excitação Comprimidas (Single e Double Excitation Tables) que reduzem o footprint de memória em mais de 15 ordens de magnitude.
• Paralelismo de Granularidade Fina: O kernel utiliza uma estrutura de dois níveis: um bloco por configuração fonte e um thread por excitação potencial. Emprega operações bit a bit eficientes e compactação de fluxo paralela para filtrar resultados inválidos sem penalidades de serialização.

C. Modelo de Execução Centrado na Memória do GPU (GPU Memory-Centric)

• Para superar a barreira de memória física, o framework adota um modelo de execução "out-of-core" inteligente.
• Processamento em Mini-Lotes (Mini-Batch): Os dados são processados em lotes gerenciáveis, tratando a memória do GPU como um cache de alta velocidade para o conjunto de trabalho ativo, em vez de tentar armazenar tudo.
• Sobreposição de Transferência e Cálculo: Utiliza streams CUDA assíncronos para sobrepor a transferência de dados (Host-to-Device e Device-to-Host) com a computação, escondendo a latência da largura de banda PCIe.
• Pipeline de Três Estágios: O fluxo é dividido em: (1) Geração e Desduplicação, (2) Inferência e Seleção Hierárquica (Top-K), e (3) Cálculo de Energia e Otimização. Dados "frios" são descarregados para a memória da CPU quando não estão em uso imediato.

3. Principais Contribuições

1. Framework Fully GPU-Accelerated: Primeiro sistema a migrar todo o fluxo de trabalho do NNQS-SCI (geração, desduplicação, inferência e seleção) para o GPU, eliminando gargalos de CPU.
2. Algoritmo de Desduplicação Distribuído: Um novo algoritmo baseado em ordenação que garante balanceamento de carga perfeito e alta eficiência de comunicação em clusters de GPUs.
3. Kernels Otimizados para Memória: Implementação de kernels CUDA que equilibram alta vazão com footprint de memória mínimo através de tabelas de excitação comprimidas e compactação de fluxo.
4. Modelo de Execução Escalável: Uma arquitetura que permite simular sistemas quânticos cujos dados intermediários excedem a capacidade física de um único GPU, mantendo o desempenho através de streaming e sobreposição de I/O.

4. Resultados Experimentais

Os testes foram realizados em um cluster com 64 GPUs NVIDIA A100 (40GB cada), comparando o cuNNQS-SCI com a implementação híbrida de ponta (NNQS-SCI).

• Aceleração de Ponta a Ponta:
  • O sistema alcançou uma aceleração de 2,32x no sistema N2 e 1,88x no sistema Cr2 (84 qubits) em comparação com a linha de base.
  • A aceleração específica nos componentes não-AI foi massiva: 432x na geração de configurações acopladas e 40x na desduplicação global.
• Mudança de Gargalo: O framework deslocou o gargalo de desempenho de limitações de CPU/memória host para a inferência da rede neural no dispositivo, permitindo o uso total da capacidade computacional dos GPUs.
• Escalabilidade:
  • Strong Scaling: Mantém 91% de eficiência paralela em 64 GPUs para um problema fixo.
  • Weak Scaling: Demonstra aceleração superlinear (até 163% de eficiência) em 64 GPUs. Isso ocorre porque a taxa de redundância global aumenta com o número de GPUs, e o algoritmo de desduplicação eficiente reduz o trabalho efetivo por nó.
• Precisão: O método mantém a precisão química (erro < 1.6 × 10⁻³ Hartree), com resultados quase idênticos ao método exato NNQS-SCI (erro médio absoluto de ~3.91 × 10⁻⁴ Hartree no Cr2), validando que a execução em GPU não compromete a fidelidade numérica.
• Gerenciamento de Memória: O modelo permitiu simular sistemas (como Cr2 e N2 com bases grandes) que teoricamente exigiriam >65GB de VRAM, mantendo o uso real abaixo de 40GB através do streaming de dados.

5. Significado e Impacto

O cuNNQS-SCI representa um avanço significativo na interseção entre HPC (High-Performance Computing) e IA para Ciência (AI4S).

• Viabilização de Sistemas Maiores: Ao quebrar a barreira de memória e de escalabilidade de CPU, o framework permite a aplicação de métodos de alta precisão (SCI) a sistemas quânticos complexos e fortemente correlacionados que antes eram intratáveis.
• Paradigma de Arquitetura: Estabelece um novo padrão para algoritmos científicos que combinam inferência de IA com lógica de domínio complexa, demonstrando que a migração total para o acelerador (GPU) é possível e necessária para a próxima geração de simulações.
• Eficiência de Recursos: Maximiza o retorno sobre o investimento em clusters de GPU, eliminando a ociosidade causada por gargalos de CPU e comunicação ineficiente.

Em resumo, o cuNNQS-SCI não é apenas uma otimização incremental, mas uma reestruturação sistêmica que expande drasticamente o horizonte de problemas solúveis em química quântica computacional.