Efficient protein structure prediction fromcompact computers to datacenters withOpenFold-TRT

Ce papier présente des accélérations combinant OpenFold, TensorRT et MMseqs2-GPU qui permettent une prédiction de structure protéique à haut débit, jusqu'à 131 fois plus rapide qu'AlphaFold2, sur des architectures allant des serveurs x86 aux superpuces ARM, sans compromettre la précision.

L'essentiel

Imaginez que prédire la forme d'une protéine (la "brique" de la vie) est comme essayer de reconstruire un château de cartes géant et complexe, juste en regardant la liste des cartes utilisées. Pendant des années, les scientifiques ont dû utiliser des supercalculateurs gigantesques et des mois de temps pour faire cela. C'était lent, cher et réservé à une élite.

Ce papier, intitulé "OpenFold-TRT", raconte l'histoire d'une équipe (principalement chez NVIDIA) qui a réussi à rendre ce processus extrêmement rapide, même sur des machines plus petites, sans perdre en précision.

Voici l'explication simple, avec quelques images pour mieux comprendre :

1. Le Problème : La course de relais trop lente

Pour prédire la forme d'une protéine, il faut faire deux étapes principales, comme dans une course de relais :

• L'étape 1 (La recherche) : Il faut chercher dans une bibliothèque immense (des milliards de séquences) des "cousins" de la protéine pour avoir des indices. C'est comme chercher une aiguille dans une botte de foin, mais la botte de foin grossit chaque jour.
• L'étape 2 (La construction) : Une fois les indices trouvés, une intelligence artificielle (Deep Learning) assemble les pièces pour dessiner la forme finale.

Avant, l'étape 1 prenait des heures et l'étape 2 prenait aussi beaucoup de temps. C'était comme essayer de construire un château de cartes avec des mains engourdies.

2. La Solution : Des outils de nouvelle génération

Les auteurs ont créé une nouvelle combinaison d'outils pour accélérer tout ça. Ils ont utilisé deux "super-pouvoirs" :

A. Le Super-Scanner (MMseqs2-GPU)

Pour l'étape de recherche, ils ont amélioré un outil appelé MMseqs2.

• L'analogie : Imaginez que vous cherchiez un mot dans un dictionnaire de 1 milliard de pages.
  • L'ancienne méthode : Vous lisiez chaque page, une par une, avec un crayon (CPU).
  • La nouvelle méthode (GPU) : Vous avez 1000 yeux qui lisent en même temps sur des pages différentes.
• L'innovation : Ils ont optimisé cet outil pour qu'il fonctionne encore mieux sur les nouvelles puces graphiques (comme la RTX PRO 6000). Résultat : la recherche est 190 fois plus rapide qu'avant ! C'est passer de la marche à pied à un avion supersonique.

B. Le Moteur de Formule 1 (OpenFold-TRT)

Pour l'étape de construction (l'intelligence artificielle), ils ont pris un logiciel open-source appelé OpenFold (une copie de la célèbre AlphaFold) et l'ont optimisé avec TensorRT (un moteur de NVIDIA).

• L'analogie : Imaginez que votre voiture (le logiciel) roule bien, mais qu'elle a des pneus d'hiver et un moteur un peu bridé.
  • Les chercheurs ont changé les pneus pour des pneus de course (optimisation de la précision mathématique) et ont débridé le moteur (fusion des opérations).
  • Résultat : La voiture fait le même trajet, mais en 20 fois moins de temps que l'ancienne version d'AlphaFold.

3. Les Résultats : Plus rapide, partout, et moins cher

Le papier montre que cette nouvelle combinaison permet des choses incroyables :

• Sur un serveur puissant (Datacenter) : Avec une seule carte graphique nouvelle génération (RTX PRO 6000), ils peuvent prédire la structure d'une protéine 131 fois plus vite que la méthode originale d'AlphaFold2.
  • Concrètement : Ce qui prenait 40 minutes, prend maintenant moins de 15 secondes.
• Sur de petites machines (Compact) : Grâce à des optimisations pour les puces ARM (comme celles des ordinateurs portables ou des serveurs économes en énergie), ils peuvent faire ces calculs sur des machines très compactes, comme le DGX Spark.
  • L'image : C'est comme réussir à faire tourner un moteur de Ferrari sur une petite voiture citadine sans la faire surchauffer.
• Gérer les géants (Mémoire) : Parfois, la bibliothèque de données est si grande qu'elle ne rentre pas dans la mémoire de la carte graphique. Grâce à une puce spéciale (Grace-Hopper) qui relie le processeur et la carte graphique très étroitement, ils peuvent utiliser la mémoire du processeur comme un "sac de stockage" supplémentaire sans ralentir le tout.

4. Pourquoi c'est important ?

Imaginez que vous vouliez prédire la forme de 350 millions de protéines (ce que les scientifiques veulent faire pour découvrir de nouveaux médicaments).

• Avec les anciennes méthodes : Il faudrait 500 ans avec un seul ordinateur.
• Avec cette nouvelle méthode : Il faudrait 4 mois et demi.

C'est la différence entre attendre toute une vie pour un résultat et pouvoir le faire pendant votre pause café. Cela ouvre la porte à la création de nouveaux médicaments, de matériaux écologiques et à la compréhension de la vie, beaucoup plus rapidement.

En résumé : Les auteurs ont pris deux étapes lentes et lourdes, et les ont transformées en une course de relais ultra-rapide grâce à de meilleurs outils de recherche et un moteur d'intelligence artificielle optimisé. Ils ont rendu la prédiction de protéines accessible, rapide et efficace, du petit ordinateur portable aux plus grands supercalculateurs.

Résumé technique

1. Problématique

La prédiction de la structure des protéines, bien que révolutionnée par AlphaFold2, fait face à un goulot d'étranglement majeur : l'explosion du volume des bases de données de séquences protéiques et la fin potentielle de la loi de Moore. Les pipelines actuels reposent sur deux étapes coûteuses en calcul :

1. Génération d'alignements de séquences multiples (MSA) : Recherche d'homologie dans des bases de données massives (métagénomiques).
2. Inférence par apprentissage profond (DL) : Prédiction de la structure 3D à partir des MSA générés.

Les solutions existantes (comme AlphaFold2 standard ou ColabFold) peinent à maintenir un débit élevé (throughput) lorsque les bases de données dépassent la mémoire GPU disponible ou lorsque le matériel n'est pas optimisé pour ces charges de travail spécifiques. Il existe un besoin urgent d'une co-conception matériel-logiciel pour accélérer l'inférence sans sacrifier la précision, que ce soit sur des serveurs de datacenter ou des systèmes compacts.

2. Méthodologie

Les auteurs ont développé une approche d'accélération combinant des optimisations logicielles et matérielles sur deux piliers principaux :

• Accélération de l'inférence DL (OpenFold-TRT) :

  • Utilisation de TensorRT pour compiler le modèle OpenFold (une réimplémentation open source d'AlphaFold2 basée sur PyTorch).
  • Mise en œuvre d'une précision mixte : utilisation de la précision TF32 pour le module ExtraMSA et BF16 pour le cœur Evoformer.
  • Support des formes dynamiques (Dynamic Shapes) : Utilisation de TorchDynamo pour exporter le modèle en ONNX avec des profils dynamiques, permettant de gérer des séquences de longueurs variables sans recompilation.
  • Fusion de noyaux (Kernel Fusion) : Consolidation des opérations d'attention en noyaux GPU optimisés pour réduire le trafic mémoire.

• Accélération de la recherche d'homologie (MMseqs2-GPU) :

  • Optimisation des noyaux GPU pour l'architecture NVIDIA Blackwell (série RTX PRO 6000) en exploitant les nouvelles instructions de programmation dynamique (DPX) pour le traitement des matrices de programmation dynamique.
  • Optimisations spécifiques pour l'architecture ARM (NVIDIA Grace-Hopper, DGX Spark) : Réécriture des opérations vectorielles SIMD (remplacement des instructions SSE x86 par des instructions NEON natives) et support des opérations vectorielles 256 bits pour améliorer l'efficacité du CPU.
  • Gestion de la mémoire unifiée sur les systèmes ARM pour permettre la recherche au-delà de la mémoire GPU physique.

3. Contributions Clés

1. OpenFold-TRT : Un pipeline d'inférence DL optimisé qui est le plus rapide parmi les solutions de type AlphaFold2, offrant une accélération significative par rapport aux implémentations JAX (AlphaFold2) et PyTorch (OpenFold standard).
2. Optimisations MMseqs2-GPU :
   • Support natif des instructions DPX sur les GPU Blackwell, permettant une génération de MSA jusqu'à 1,4x plus rapide que sur les GPU L40S précédents.
   • Optimisations ARM permettant une exécution efficace sur des systèmes à faible consommation (DGX Spark) et une gestion transparente de bases de données dépassant la mémoire GPU via la mémoire partagée CPU-GPU (Grace-Hopper).
3. Pipeline End-to-End : Intégration complète de la génération de MSA (MMseqs2-GPU) et de l'inférence DL (OpenFold-TRT) sur différentes plateformes matérielles (x86 + Blackwell, ARM + Grace-Hopper).

4. Résultats

Les tests ont été réalisés sur 20 cibles difficiles du concours CASP14 et sur des bases de données de référence de 30 millions de séquences.

• Vitesse d'inférence DL :

  • OpenFold-TRT est 2,54x plus rapide que OpenFold-PyTorch standard.
  • Il est 20,69x plus rapide qu'AlphaFold2 (JAX) et 6,13x plus rapide que ColabFold-batch sur un serveur x86 équipé d'une RTX PRO 6000.
  • Le temps moyen d'inférence DL est de 5,4 secondes sur DGX GH200 et 5,6 secondes sur RTX PRO 6000.

• Vitesse de génération de MSA :

  • Sur un serveur x86 avec RTX PRO 6000, MMseqs2-GPU est 191,4x plus rapide que le pipeline de base AlphaFold2 (JackHMMER+HHblits) et 5,8x plus rapide que ColabFold-search (CPU).
  • Le système Grace-Hopper (GH200) maintient un débit constant même lorsque la taille de la base de données dépasse la mémoire GPU (96 Go), grâce à l'interconnexion CPU-GPU à haut débit (450 Go/s), là où les systèmes x86+GPU voient leur performance chuter.

• Performance Globale (End-to-End) :

  • La combinaison RTX PRO 6000 + MMseqs2-GPU + OpenFold-TRT atteint le débit le plus élevé, avec un temps total moyen de 15,93 secondes par protéine.
  • Cela représente une accélération de 131,4x par rapport au pipeline AlphaFold2 de base et 5,94x par rapport à ColabFold (CPU).
  • Précision : Les scores TM (Template Modeling) restent cohérents avec les méthodes de référence (environ 0,67-0,71), confirmant qu'il n'y a pas de perte de précision malgré les accélérations.

• Évolutivité :

  • Sur un serveur unique, prédire les 350 millions de séquences de la base de données AlphaFold prendrait environ 500 ans avec l'ancienne méthode la plus rapide. Avec OpenFold-TRT, ce temps est réduit à 4,5 mois avec la même infrastructure.

5. Signification

Ce travail démontre que l'optimisation conjointe du logiciel (TensorRT, MMseqs2) et du matériel (Blackwell, Grace-Hopper) peut surmonter les limites actuelles de la prédiction de structure protéique.

• Accessibilité : Ces accélérations permettent d'effectuer des inférences à haut débit sur des systèmes compacts (DGX Spark) et des serveurs standard, rendant la prédiction de structures à grande échelle plus accessible.
• Impact sur la recherche : La réduction drastique du temps de calcul (de plusieurs années à quelques mois pour des bases de données massives) est cruciale pour le développement de nouvelles méthodes de conception de protéines génératives et pour l'exploration de l'espace des séquences protéiques.
• Durabilité : L'efficacité énergétique améliorée sur les architectures ARM (Grace-Hopper) offre une voie pour des calculs scientifiques plus durables.

L'ensemble des accélérations a été intégré dans les dépôts open source de MMseqs2, OpenFold et TensorRT, permettant une reproduction immédiate des résultats par la communauté scientifique.