Efficient protein structure prediction fromcompact computers to datacenters withOpenFold-TRT

Deze paper introduceert OpenFold-TRT, een versnelling voor deep learning-inferentie die met TensorRT en MMseqs2-GPU op zowel compacte x86-systemen als krachtige datacentersystemen tot 131 keer snellere eiwitstructuurvoorspelling mogelijk maakt dan AlphaFold2 zonder in te leveren op nauwkeurigheid.

De kern

Stel je voor dat het voorspellen van de vorm van een eiwit (een bouwsteen van het leven) als het oplossen van een gigantische, driedimensionale puzzel is. Vroeger duurde het jaren om deze puzzels op te lossen, en zelfs de slimste computers van vandaag (zoals die van AlphaFold2) kunnen er soms dagen over doen als je duizenden puzzels tegelijk wilt maken.

Deze paper vertelt het verhaal van een nieuwe, razendsnelle manier om deze puzzels op te lossen, ontwikkeld door een team van NVIDIA en universiteiten. Ze noemen hun oplossing OpenFold-TRT.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het probleem: De twee stappen van de reis

Om de vorm van een eiwit te voorspellen, moet een computer twee grote taken uitvoeren:

• Stap 1: De zoektocht (MSA): De computer moet in een enorme bibliotheek van miljarden oude eiwitten zoeken naar familieleden van het eiwit dat je wilt bestuderen. Dit is als het zoeken naar een specifiek boek in een bibliotheek die groter is dan de aarde.
• Stap 2: De constructie (DL Inference): Zodra de familieleden gevonden zijn, bouwt een slim algoritme (een "neuraal netwerk") de 3D-vorm op. Dit is als het daadwerkelijk bouwen van het model op basis van de gevonden blauwdrukken.

Vroeger waren beide stappen traag. De bibliotheek werd steeds groter, maar de computers werden niet snel genoeg om mee te houden.

2. De oplossing: Een nieuwe raceauto en een super-efficiënte routeplanner

Het team heeft twee dingen gedaan om dit te versnellen:

A. De nieuwe motor: TensorRT en Blackwell
Ze hebben de software (OpenFold) herschreven om te werken met een nieuwe technologie van NVIDIA genaamd TensorRT.

• De analogie: Stel je voor dat je eerder een oude, zware vrachtwagen gebruikte om een pakketje te bezorgen. Nu hebben ze die vrachtwagen vervangen door een Formule 1-raceauto die speciaal is afgesteld voor deze route.
• Ze hebben ook gebruikgemaakt van de nieuwste grafische kaarten (de RTX PRO 6000 met de "Blackwell" chip). Deze kaarten zijn als een superkrachtige motor die veel meer werk per seconde kan verzetten dan de oude modellen.

B. De slimme routeplanner: ARM en Grace Hopper
Voor de zoektocht in de bibliotheek (Stap 1) hebben ze de software ook aangepast voor een ander type computerchips (ARM-architectuur, zoals in de DGX Spark en Grace Hopper).

• De analogie: Normaal gesproken moet je in een bibliotheek eerst naar de afdeling A, dan B, dan C... en als de bibliotheek te groot is, moet je wachten tot de volgende bus (de CPU) je boeken haalt.
• Met de nieuwe ARM-chips is het alsof de bibliothecaris (de CPU) en de boekenplanken (de GPU) direct met elkaar verbonden zijn door een snelle tunnel. Ze kunnen samenwerken alsof ze in dezelfde kamer zitten, waardoor ze veel grotere bibliotheken kunnen doorzoeken zonder vast te lopen.

3. De resultaten: Van dagen naar seconden

De resultaten zijn verbluffend:

• Snelheid: Op een enkele krachtige server kunnen ze nu eiwitten voorspellen die 131 keer sneller zijn dan de oude AlphaFold2 methode.
  • Vergelijking: Als de oude methode 2424 seconden (bijna 40 minuten) nodig had voor één eiwit, doet de nieuwe methode dit in slechts 15 seconden.
• Schaalbaarheid: Ze kunnen nu zelfs zoeken in bibliotheken die groter zijn dan het geheugen van de computer zelf. De nieuwe systemen kunnen "uitwaaieren" naar het geheugen van de processor zonder dat het traag wordt.
• Kwaliteit: Het allerbelangrijkste: hoewel het zo veel sneller is, is het antwoord even nauwkeurig. Ze hebben de puzzel niet sneller opgelost door er op los te gissen; ze hebben gewoon de manier waarop ze zoeken en bouwen geoptimaliseerd.

4. Waarom is dit belangrijk?

Stel je voor dat je een fabriek hebt die nieuwe medicijnen of materialen ontwerpt.

• Met de oude methode zou het 500 jaar duren om alle bekende eiwitten in de wereld te analyseren met één computer.
• Met deze nieuwe methode duurt het slechts 4,5 maanden.

Dit opent de deur voor een nieuwe revolutie in de biologie. Onderzoekers kunnen nu in een paar maanden doen wat ze voorheen jaren nodig hadden. Ze kunnen duizenden nieuwe eiwitten ontwerpen om ziektes te bestrijden of schone energie te maken, allemaal dankzij een stukje software dat is "getuned" om de hardware van de computer perfect te benutten.

Kortom: Ze hebben de sleutel gevonden om de snelste computers ter wereld te laten werken als een goed geoliede machine, waardoor het voorspellen van eiwitstructuren niet langer een bottleneck is, maar een fluitje van een cent wordt.

Technische samenvatting

Probleemstelling

De voorspelling van eiwitstructuren, hoewel revolutionair door AlphaFold2, blijft een rekenkundig intensieve taak die twee hoofdwerklasten omvat:

1. Homologie-retrieval: Het genereren van Multiple Sequence Alignments (MSA) via zoekopdrachten in enorme databanken.
2. Deep Learning (DL) inferentie: Het voorspellen van de 3D-structuur op basis van de gegenereerde MSA's.

Hoewel databases exponentieel groeien, stuit de hardware-ontwikkeling (Moore's Law) op zijn grenzen. Bestaande pipelines (zoals de baseline AlphaFold2) zijn traag en afhankelijk van specifieke hardware. Er is een dringende behoefte aan hardware-software co-design om de doorvoer (throughput) te maximaliseren zonder in te leveren op nauwkeurigheid, zodat grootschalige in-silico data-generatie mogelijk wordt voor nieuwe generatieve methoden.

Methodologie

De auteurs introduceren een geoptimaliseerde stack die bestaande open-source tools combineert met specifieke hardware-versnellingen:

1. OpenFold-TRT (DL Inferentie):

   • Ze hebben OpenFold (een open-source herimplementatie van AlphaFold2) geoptimaliseerd met NVIDIA TensorRT.
   • Technische optimalisaties:
     • Mixed Precision: Gebruik van TF32 voor ExtraMSA en BF16 voor de Evoformer-module (de kern van het model).
     • Dynamische Vormen (Dynamic Shapes): Ondersteuning voor variabele sequentielengtes zonder hercompilatie, via TorchDynamo en ONNX-export.
     • Kernel Fusion: Het samenvoegen van multi-stap attention-operaties tot geoptimaliseerde GPU-kernels om geheugentransacties te verminderen.
   • Dit resulteert in een snellere inferentie dan de standaard PyTorch-implementatie of JAX-based AlphaFold2.

2. MMseqs2-GPU Optimalisaties (MSA Generatie):

   • Versnelling van homologie-zoekopdrachten op GPU's, specifiek voor de NVIDIA Blackwell architectuur (RTX PRO 6000) en ARM-systemen.
   • Blackwell: Benutting van de nieuwe DPX-instructieset (Dynamic Programming) en optimalisatie van de verdeling van DP-cellen per CUDA-thread. Dit verdubbelt de integer-operatie doorvoer ten opzichte van eerdere generaties (Ada/L40S).
   • ARM (Grace-Hopper & DGX Spark): Herimplementatie van SIMD-operaties (van x86 SSE naar ARM NEON) en ondersteuning voor 256-bit vectoroperaties. Dit verlaagt de instructielatentie en verbetert de single-thread doorvoer aanzienlijk.

3. Hardware-architecturen:

   • De studie evalueert end-to-end prestaties op diverse platforms: x86-servers met RTX PRO 6000, L40S, H100, en ARM-gebaseerde systemen zoals de NVIDIA Grace-Hopper Superchip (GH200) en de compacte DGX Spark.

Belangrijkste Bijdragen

• OpenFold-TRT: Een nieuwe, snellere inferentie-engine die AlphaFold2-achtige voorspellingen mogelijk maakt met een aanzienlijke snelheidswinst.
• Blackwell & ARM Optimalisaties: Eerste implementaties van MMseqs2-GPU die specifiek zijn afgestemd op de Blackwell-architectuur (via DPX) en ARM-architecturen (via NEON-optimalisaties).
• Schalbaarheid: Bewijs dat ARM-systemen met gedeeld geheugen (Unified Memory) zoals de GH200, MSA-generatie kunnen uitvoeren die groter is dan het beschikbare GPU-geheugen, zonder snelheidsverlies, in tegenstelling tot traditionele x86+GPU-systemen.
• Open Source Beschikbaarheid: Alle versnellingen zijn upstream gegaan in de repositories van MMseqs2, OpenFold en TensorRT, waardoor reproduceerbaarheid gegarandeerd is.

Resultaten

De prestaties zijn getest op 20 "harde" targets van de CASP14-concurrentie.

• Snelheidswinst (DL Inferentie):

  • OpenFold-TRT is 2,54× sneller dan de standaard OpenFold-PyTorch.
  • Het is 20,69× sneller dan de baseline AlphaFold2 (JAX) en 6,13× sneller dan ColabFold-batch.
  • Op de GH200 duurt het inferentieproces gemiddeld slechts 5,4 seconden per eiwit.

• Snelheidswinst (MSA Generatie):

  • Op een x86-server met één RTX PRO 6000 is MMseqs2-GPU 1,4× sneller dan op een L40S.
  • Totaal is deze setup 191,4× sneller dan de baseline AlphaFold2-pipeline (JackHMMER+HHblits) en 5,8× sneller dan ColabFold-search (CPU).

• End-to-End Prestaties:

  • De combinatie van MMseqs2-GPU (MSA) en OpenFold-TRT (DL) op een x86-server met RTX PRO 6000 is de snelste pipeline, met een gemiddelde doorlooptijd van 15,93 seconden per sequentie.
  • Dit is een 131,4× snelheidswinst ten opzichte van de baseline AlphaFold2-pipeline.
  • De GH200 (Grace-Hopper) presteert vergelijkbaar snel (gemiddeld 16,62s totaal), maar heeft een uniek voordeel: het kan databases verwerken die groter zijn dan het GPU-geheugen (96GB) door gebruik te maken van het host-geheugen via de snelle Chip-to-Chip (C2C) verbinding, zonder in te leveren op doorvoer.

• Nauwkeurigheid:

  • De TM-scores (een maatstaf voor structurele overeenkomst) blijven consistent tussen de versnellingen en de baseline, wat aantoont dat de snelheidswinst niet ten koste gaat van de voorspellingskwaliteit.

Significantie

Deze studie toont aan dat het mogelijk is om eiwitstructuurvoorspelling te versnellen tot een schaal die eerder ondenkbaar was, zonder de nauwkeurigheid te verliezen.

• Impact op Schaalbaarheid: Waar het voorspellen van 350 miljoen sequenties (zoals in de AlphaFold-database) met de vorige snelste oplossing ongeveer 500 jaar zou duren op één server, kan dit nu met OpenFold-TRT in ongeveer 4,5 maanden worden voltooid.
• Toekomstige Toepassingen: Deze efficiëntie maakt grootschalige in-silico data-generatie haalbaar, wat essentieel is voor het trainen van nieuwe generatieve modellen voor eiwitontwerp.
• Hardware Diversiteit: Het bewijst dat zowel krachtige datacenter-GPU's (Blackwell) als energiezuinige ARM-systemen (Grace-Hopper, DGX Spark) effectief kunnen worden ingezet voor deze complexe taken, waardoor de drempel voor hoogwaardige bio-informatica verder verlaagt.