NVIDIA avalia métodos Tensor Core para acelerar simulações moleculares
A NVIDIA publicou testes de métodos Tensor Core no PET, modelo de simulação molecular desenvolvido pelo Lab COSMO da EPFL. Em uma GPU B300, o BF16x9 acelerou em até 1,7x o passo completo da dinâmica molecular.
A NVIDIA publicou uma avaliação de métodos Tensor Core para acelerar o PET, modelo de aprendizado de máquina para simulações atomísticas desenvolvido pelo Lab COSMO da EPFL. Os testes compararam BF16x9 em GPUs Blackwell e TF32x3 em GPUs Ampere com o caminho FP32, usado como referência de precisão simples.
No sistema maior, com 17.280 átomos, o BF16x9 acelerou grandes operações de matriz entre 2,2 e 2,5 vezes. Na GPU B300, o ganho chegou a 1,7x no passo completo da simulação. Com TF32x3, as principais operações de matriz na A100 aceleraram entre 1,7 e 1,9x, e o passo completo, em até 1,5x.
A avaliação também verificou estabilidade numérica. Nos testes descritos, o TF32x3 manteve oscilações de energia comparáveis às do FP32 e atingiu o critério de relaxamento estrutural em 218 a 222 etapas, contra 218 no FP32. O TF32 de uma passagem foi mais rápido em alguns testes, mas apresentou maior deriva de energia e não convergiu dentro do limite de etapas da relaxação.
Segundo a NVIDIA, o ganho depende do tamanho do sistema e da parcela do tempo gasta em operações de matriz. A empresa não mediu o efeito combinado com outras otimizações de software e recomenda validar desempenho e precisão para cada modelo e aplicação.
O que muda para você
A avaliação interessa a pesquisadores e equipes que executam modelos de potenciais interatômicos em GPUs NVIDIA. Em GPUs Blackwell compatíveis, o BF16x9 pode ser habilitado no cuBLAS sem alterar os pesos ou a arquitetura do PET. O TF32x3 testado em Ampere exigiu substituir camadas lineares por uma implementação baseada em CUTLASS. Para outros modelos, ainda é preciso medir os ganhos e verificar a estabilidade.
Minha leitura
Em simulações, ganhar velocidade só serve se a trajetória continuar confiável. A avaliação mostra que isso pode ser possível com TF32x3 nos testes realizados, mas não autoriza generalizar o resultado para qualquer modelo ou sistema. Eu começaria verificando quanto do tempo de execução está nas operações de matriz; depois, compararia a estabilidade e a convergência com o caminho FP32.
Thiago FelizolaConsultor de IA, curador do Diário da IA
