Efficient protein structure prediction fromcompact computers to datacenters withOpenFold-TRT

Este artigo apresenta otimizações de inferência para OpenFold e TensorRT que, combinadas com MMseqs2-GPU em diferentes arquiteturas de hardware, permitem a previsão de estruturas proteicas em alta velocidade e sem perda de precisão, desde computadores compactos até datacenters.

O essencial

Imagine que você é um chef de cozinha tentando recriar uma receita secreta (a estrutura de uma proteína) apenas lendo uma lista de ingredientes (a sequência de aminoácidos). Antigamente, para fazer isso, você precisava de uma biblioteca gigante de receitas antigas para comparar, e o processo de encontrar as receitas certas e depois montar o prato levava dias.

Este artigo é como a história de como um grupo de cientistas e engenheiros da NVIDIA e de universidades transformou essa cozinha lenta em uma fábrica de comida ultrarrápida, capaz de preparar milhares de pratos em segundos, sem perder o sabor (a precisão).

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

1. O Problema: A Cozinha Lenta

Antes, o processo de prever a estrutura de uma proteína tinha duas etapas principais:

• A Pesquisa (MSA): Era como ir a uma biblioteca gigante e procurar manualmente em milhões de livros para encontrar receitas similares. Isso levava muito tempo.
• A Cozinhagem (DL Inference): Depois de achar as receitas, você usava um computador poderoso (uma IA chamada AlphaFold) para montar o prato. Isso também era lento e exigia computadores enormes.

O problema é que a quantidade de "ingredientes" (dados biológicos) está crescendo mais rápido do que a velocidade dos computadores (a famosa "Lei de Moore" está quebrando). Precisávamos de uma maneira mais inteligente de cozinhar.

2. A Solução: O "Super Chef" OpenFold-TRT

Os autores criaram uma nova combinação de ferramentas que funciona como um turbo para a cozinha:

• O Motor (TensorRT): Eles pegaram o software de cozinha (OpenFold) e o "reprogramaram" para falar a língua nativa da placa de vídeo (GPU) da NVIDIA. É como trocar o motor de um carro popular por um de Fórmula 1. O carro é o mesmo, mas agora ele corre muito mais rápido.
• O Filtro Inteligente (MMseqs2-GPU): Para a etapa de pesquisa na biblioteca, eles criaram um filtro super-rápido que roda diretamente na placa de vídeo. Em vez de ler livro por livro, ele "pula" direto para as páginas relevantes.

3. Os Resultados: De Dias para Segundos

Com essa nova combinação, eles conseguiram feitos impressionantes:

• Velocidade: O sistema deles é 131 vezes mais rápido do que a versão original do AlphaFold2.
  • Analogia: Se a versão antiga demorasse 131 dias para prever a estrutura de uma proteína, a nova versão faz isso em apenas 1 dia.
• Precisão: O "prato" final ficou exatamente igual. Não houve perda de qualidade. A comida saiu tão boa quanto a feita no método antigo, só que em segundos.

4. Hardware: De Computadores de Escritório a Supercomputadores

O artigo mostra que essa tecnologia funciona em diferentes tamanhos de "cozinhas":

• O Caminhão de Carga (RTX PRO 6000): Um servidor potente com uma única placa de vídeo nova (Blackwell) consegue fazer o trabalho de uma frota inteira de computadores antigos.
• O Carro Elétrico Eficiente (DGX Spark): Eles também adaptaram o sistema para funcionar em máquinas pequenas e que gastam pouca energia, como o "DGX Spark". É como ter um carro de corrida que cabe no porta-malas de um carro comum.
• O Gigante Inteligente (Grace Hopper): Para quando a biblioteca de receitas é tão grande que não cabe na memória do computador, eles usaram um sistema onde o processador e a memória estão "colados" um no outro (como um chip único), permitindo ler dados gigantes sem travar.

5. Por que isso importa?

Imagine que você quer prever a estrutura de 350 milhões de proteínas (o tamanho de um banco de dados real).

• Com a tecnologia antiga, levaria 500 anos para um único computador fazer isso.
• Com a tecnologia nova, levaria apenas 4 meses e meio.

Isso é crucial para a ciência. Agora, os pesquisadores podem testar milhões de novas ideias para criar novos medicamentos, enzimas para limpar o plástico ou materiais biológicos em tempo recorde. É como ter uma máquina do tempo que acelera a descoberta científica.

Resumo da Ópera:
Eles pegaram um processo biológico complexo e lento e o transformaram em algo rápido e eficiente, permitindo que computadores menores e mais baratos façam o trabalho que antes exigia supercomputadores caríssimos, tudo isso sem perder a precisão. É a democratização da supercomputação para a biologia.

Resumo técnico

1. O Problema

A previsão de estrutura de proteínas, impulsionada por ferramentas como o AlphaFold2, tornou-se fundamental para a biologia e o design de fármacos. No entanto, à medida que as bases de dados de sequências proteicas crescem exponencialmente e a Lei de Moore desacelera, os pipelines atuais enfrentam gargalos de desempenho.

• Gargalos de Hardware e Software: Os pipelines tradicionais dependem de duas etapas principais: geração de Alinhamento de Sequências Múltiplas (MSA) e inferência de Aprendizado Profundo (DL). A geração de MSA é frequentemente limitada pela memória da GPU e pela velocidade de busca homóloga, enquanto a inferência DL é computacionalmente intensiva.
• Escalabilidade: Soluções existentes, como o AlphaFold2 padrão (baseado em JAX) ou o ColabFold, tornam-se inviáveis para previsões em larga escala (ex: bilhões de sequências) devido ao tempo de execução excessivo e à limitação de memória em sistemas convencionais.

2. Metodologia

Os autores propuseram uma reengenharia completa do pipeline de inferência, focando na otimização de hardware e software para acelerar ambas as etapas (MSA e DL) sem sacrificar a precisão.

• Aceleração de Inferência DL (OpenFold-TRT):

  • Substituição do backend PyTorch padrão do OpenFold por TensorRT (NVIDIA).
  • Uso de precisão mista: TF32 para o módulo ExtraMSA e BF16 (Bfloat16) para o núcleo Evoformer.
  • Implementação de suporte a formas dinâmicas (Dynamic Shapes) via TorchDynamo e exportação ONNX, permitindo o processamento de sequências de proteínas de comprimentos variáveis sem recompilação.
  • Fusão de kernels (Kernel Fusion) para reduzir o tráfego de memória e aumentar a intensidade aritmética.

• Aceleração de Busca Homóloga (MMseqs2-GPU):

  • Otimização dos kernels de alinhamento do MMseqs2-GPU para arquiteturas NVIDIA Blackwell (RTX PRO 6000), aproveitando novas instruções de programação dinâmica (DPX) e aumentando o throughput de operações inteiras.
  • Otimizações específicas para ARM (NVIDIA Grace-Hopper e DGX Spark), incluindo a reimplementação de operações SIMD (256-bit) usando instruções NEON nativas para melhorar a eficiência de ciclos de CPU e permitir buscas em bases de dados maiores que a memória da GPU, utilizando memória compartilhada CPU-GPU.

• Configurações de Hardware Testadas:

  • Servidores x86 com GPUs NVIDIA L40S, H100 e RTX PRO 6000 Blackwell.
  • Sistemas ARM: DGX GH200 (Grace-Hopper Superchip) e DGX Spark (GB10).

3. Principais Contribuições

1. OpenFold-TRT: Uma nova implementação de inferência que é a mais rápida disponível para modelos estilo AlphaFold2, oferecendo uma aceleração significativa em relação ao OpenFold padrão e ao AlphaFold2 original.
2. Otimizações de MMseqs2-GPU: Melhorias de desempenho para as arquiteturas Blackwell e ARM, permitindo que sistemas de menor porte (como DGX Spark) e sistemas de grande porte (GH200) realizem buscas homólogas em escala metagenômica com eficiência.
3. Pipeline End-to-End Otimizado: Demonstração de que a combinação de MMseqs2-GPU (para MSA) e OpenFold-TRT (para DL) em um servidor x86 com uma única RTX PRO 6000 constitui o pipeline mais rápido para previsão de estrutura de proteínas.
4. Resiliência à Escala de Memória: Solução para o problema de limitação de memória da GPU em sistemas ARM (GH200), onde a memória unificada permite buscar em bases de dados maiores que a capacidade da VRAM da GPU sem perda significativa de velocidade.

4. Resultados

Os testes foram realizados em 20 alvos difíceis do CASP14 e em benchmarks de escalabilidade de banco de dados.

• Velocidade de Inferência DL:

  • O OpenFold-TRT foi 2,54× mais rápido que o OpenFold-PyTorch padrão.
  • Comparado ao AlphaFold2 (JAX), foi 20,69× mais rápido.
  • Comparado ao ColabFold-batch, foi 6,13× mais rápido.
  • Em um sistema RTX PRO 6000, a inferência DL média foi de 5,6 segundos por proteína.

• Velocidade de Geração de MSA:

  • O MMseqs2-GPU no RTX PRO 6000 foi 1,4× mais rápido que no L40S anterior.
  • Foi 191,4× mais rápido que o pipeline clássico (JackHMMER + HHblits) usado no AlphaFold2 original.
  • Foi 5,8× mais rápido que o ColabFold-search usando CPU.

• Desempenho End-to-End (Total):

  • O pipeline completo (MSA + DL) no RTX PRO 6000 levou em média 15,93 segundos por alvo.
  • Isso representa uma aceleração de 131,4× em relação ao pipeline AlphaFold2 original e 5,94× em relação ao ColabFold padrão.
  • O sistema DGX GH200 (ARM) alcançou desempenho comparável, com a vantagem de manter a velocidade mesmo quando o banco de dados excede a memória da GPU (96GB), graças à interconexão de alta largura de banda CPU-GPU.

• Precisão:

  • As otimizações não comprometeram a precisão. Os escores TM-score (medida de similaridade estrutural) permaneceram consistentes com os métodos de base (AlphaFold2 e ColabFold), variando em torno de 0,67 a 0,71.

• Impacto em Escala:

  • O artigo estima que prever 350 milhões de sequências (como no banco de dados AlphaFold) levaria cerca de 500 anos com a solução anterior mais rápida em um único servidor. Com o OpenFold-TRT, esse tempo cai para quatro meses e meio, tornando viável a geração de dados em larga escala para novos métodos de design de proteínas.

5. Significado e Impacto

Este trabalho demonstra que a co-otimização de hardware e software é essencial para o futuro da biologia computacional. Ao reduzir drasticamente o tempo de inferência e superar as limitações de memória, as técnicas apresentadas permitem:

• Democratização do Acesso: Previsão de estruturas em computadores compactos e eficientes energeticamente (como DGX Spark), além de datacenters.
• Aceleração da Descoberta: Possibilita a análise de bancos de dados de proteínas em escala petabyte em tempos viáveis, acelerando a descoberta de fármacos e o design de novas enzimas.
• Sustentabilidade: A eficiência energética e de tempo reduz o custo computacional e ambiental da pesquisa em biologia estrutural.

O código e as otimizações foram disponibilizados como software de código aberto (MMseqs2, OpenFold e TensorRT), permitindo que a comunidade científica reproduza e utilize esses ganhos de desempenho imediatamente.