BioIR da Nvidia acelera previsão de proteínas em 2,9x com PyTorch puro

A previsão de estrutura de biomoléculas mudou de corridas de alvo único para listas de trabalho em escala de proteoma. O gargalo deixou de ser “o modelo consegue dobrar uma proteína?” e virou “quão rápido uma fila inteira de alvos independentes passa por parsing, extração de features, inferência em GPU e escrita de saída?”. A resposta da Nvidia é o BioNeMo Inference Runtime (BioIR), uma biblioteca Python que acelera modelos de previsão de estrutura em GPUs da empresa mantendo o fluxo padrão do PyTorch.

Já testado em escala de produção

O BioIR não é um experimento. Ele foi a engrenagem por trás da recente expansão do AlphaFold Database, gerando estruturas de complexos proteicos em 4.777 proteomas — cerca de 31 milhões de complexos candidatos, dos quais 1,81 milhão foram liberados como predições de alta confiança.

A biblioteca está disponível em um repositório aberto no GitHub, com um wheel contendo CUBINs pré-compilados. Para uso em tempo de execução, bastam Python 3.12+, uma GPU Nvidia compatível com driver, o checkpoint do modelo e MSAs A3M por cadeia. Não exige nvcc, código-fonte CUDA, CMake ou o toolkit CUDA.

Três camadas de aceleração

O BioIR ataca os pontos que stacks de inferência genéricas não otimizam por completo — como as pilhas Pairformer e Evoformer, operações de triângulo, atenção pairwise, transformers de difusão e módulos atômicos. Os modelos continuam sendo objetos torch.nn.Module comuns, sem passo de export ou artefato separado. A aceleração acontece em três níveis:

• Seleção de kernel: operações suportadas escolhem entre kernels customizados do BioIR, cuEquivariance ou fallback PyTorch, conforme o modelo, GPU, tipo de dado e formato do tensor.
• Otimização de módulo: o mecanismo optimize() habilita captura de CUDA Graph em módulos compatíveis, reduzindo overhead de lançamento.
• Escala de pipeline: um executor Ray coloca uma réplica completa do modelo em cada GPU visível e distribui entradas independentes, sobrepondo etapas de CPU (parsing, featurização, escrita) à dobra em GPU.

Um detalhe importante: o Ray não divide uma única passada forward entre GPUs. O modo réplica escala listas de trabalho, não alvos individuais.

Os números do benchmark

Em testes pareados com 1.000 dímeros humanos em 8 GPUs H100 de 80 GB, comparando Boltz-2 acelerado por BioIR contra uma implementação open source compilada com torch.compile, os resultados foram claros:

• 58,5 mil resíduos dobrados por GPU-hora com BioIR, contra 20,2 mil da implementação pública — um ganho de 2,90x em throughput normalizado por resíduo.
• A implementação pública estourou memória em 29 dos 1.000 alvos; o BioIR completou todos.

No nível de forward de modelo, os speedups geométricos médios sobre o baseline torch.compile foram de 1,55x (OpenFold3), 1,78x (Boltz2) e 2,56x (OpenFold2 monomer) em H100. A Nvidia ressalta que são números da etapa de dobra, específicos daquele conjunto e hardware, e que não devem ser generalizados para todos os modelos suportados.

Energia também conta

Extrapolando linearmente para 1 milhão de alvos comparáveis, o BioIR consumiria cerca de 11 MWh contra 35 MWh da implementação pública em equivalentes de TDP de 8 GPUs. São estimativas de potência nominal, sem PUE de data center — mas uma economia de 23 a 43 MWh por milhão de alvos é um número material para campanhas em escala de proteoma.