Efficient protein structure prediction fromcompact computers to datacenters withOpenFold-TRT

Die Arbeit stellt Beschleunigungen für OpenFold und TensorRT vor, die eine bis zu 131-fach schnellere Proteinstrukturvorhersage im Vergleich zu AlphaFold2 auf verschiedenen Hardware-Plattformen – von kompakten Systemen bis zu Rechenzentren – ohne Genauigkeitsverlust ermöglichen.

Das Wesentliche

Stellen Sie sich vor, Sie sind ein Architekt, der versuchen muss, ein riesiges, komplexes Schloss aus nur einer Liste von Buchstaben (der DNA-Sequenz) zu bauen. Das ist im Grunde das, was Wissenschaftler tun, wenn sie die Struktur von Proteinen vorhersagen. Proteine sind die kleinen Baumeister unseres Körpers, und um zu verstehen, wie sie funktionieren, müssen wir wissen, wie sie gefaltet sind.

Bis vor kurzem war dieser Prozess wie das Bauen eines Schlosses mit einem sehr langsamen, alten Werkzeugkasten. Es dauerte Stunden oder sogar Tage, und man brauchte dafür riesige, teure Supercomputer.

Dieser neue Bericht von NVIDIA und Partnern stellt eine revolutionäre Werkzeug-Upgrade vor, das diesen Prozess um das 131-fache beschleunigt. Hier ist die Erklärung, wie das funktioniert, ganz einfach und mit ein paar Bildern:

1. Das alte Problem: Die Bibliothek und der Baumeister

Der Prozess hat zwei Hauptschritte:

• Schritt A (Die Bibliothek): Zuerst muss man in einer riesigen Bibliothek nach ähnlichen Buchstabenfolgen suchen, um Hinweise zu finden. Das ist wie das Suchen nach alten Bauplänen in einer Bibliothek, die so groß ist wie ein ganzer Kontinent. Bisher war das Suchen in dieser Bibliothek sehr langsam.
• Schritt B (Der Baumeister): Dann nimmt ein künstlicher Intelligenz-Algorithmus (der "Baumeister") diese Hinweise und baut das 3D-Modell. Auch das war bisher rechenintensiv und langsam.

2. Die Lösung: Ein Turbo-Kit für beide Schritte

Die Forscher haben zwei Hauptwerkzeuge verbessert: MMseqs2 (für die Bibliothek) und OpenFold (für den Baumeister).

Das neue Suchwerkzeug (MMseqs2-GPU)

Stellen Sie sich vor, früher mussten Sie jede einzelne Seite in der riesigen Bibliothek einzeln durchblättern.

• Die alte Methode: Ein einzelner Bibliothekar (CPU) läuft durch die Gänge.
• Die neue Methode: Sie haben jetzt einen ganzen Schwarm von Robotern (die GPU), die alle gleichzeitig arbeiten.
• Der Blackwell-Turbo: Die Forscher haben diese Roboter auf den neuesten NVIDIA-Chips (Blackwell) laufen lassen. Es ist, als hätten sie den Bibliothekar nicht nur durch Roboter ersetzt, sondern sie mit einem Lichtgeschwindigkeits-Schreibsystem ausgestattet. Auf einem Server mit einer einzigen neuen Grafikkarte (RTX PRO 6000) finden sie die Hinweise so schnell, dass es 191-mal schneller ist als die alte Standardmethode.

Der neue Baumeister (OpenFold-TRT)

Jetzt kommt der eigentliche Architekt (die KI).

• Das alte Werkzeug: Der Architekt arbeitete sehr sorgfältig, aber langsam, weil er jede Berechnung mit extrem hoher Genauigkeit (wie mit einem Mikroskop) durchführte, was viel Zeit kostete.
• Das neue Werkzeug (TensorRT): Die Forscher haben dem Architekten ein Super-Setzzeug gegeben. Sie haben die Berechnungen so optimiert, dass der Architekt immer noch genau baut, aber jetzt mit einem Turbo-Modus arbeitet. Er nutzt eine spezielle "Misch-Genauigkeit", die schnell ist, aber das Ergebnis nicht verschlechtert.
• Das Ergebnis: Der Architekt baut das Schloss 20-mal schneller als das alte Modell (AlphaFold2) und 2,5-mal schneller als die bisherige beste Open-Source-Version.

3. Die Hardware: Von der Garage zum Rechenzentrum

Das Tolle an dieser Arbeit ist, dass es nicht nur für riesige Supercomputer funktioniert, sondern auch für kleinere, effizientere Maschinen:

• Der "Super-Chip" (Grace-Hopper): Stellen Sie sich vor, CPU (das Gehirn) und GPU (die Hände) sind nicht mehr in zwei verschiedenen Räumen, sondern direkt miteinander verbunden, wie ein Gehirn und seine Hände, die sich die Hand reichen. Das ermöglicht es, auch riesige Datenbanken zu durchsuchen, ohne dass der Speicherplatz ausgeht.
• Der "Mini-Computer" (DGX Spark): Selbst auf einem kleinen, energieeffizienten Gerät, das so groß wie ein Schuhkarton ist, funktioniert das System erstaunlich gut. Es ist wie ein Rennwagen, der auch in der Stadt fährt.

4. Warum ist das wichtig?

Stellen Sie sich vor, Sie wollen 350 Millionen verschiedene Proteine analysieren (wie es die aktuelle Datenbank versucht).

• Mit den alten Methoden würde das auf einem einzelnen Server 500 Jahre dauern.
• Mit den neuen Methoden dauert es nur ein halbes Jahr.

Das ist ein riesiger Unterschied. Es bedeutet, dass Wissenschaftler viel schneller neue Medikamente entwickeln, Enzyme für umweltfreundliche Prozesse designen oder verstehen können, wie Krankheiten entstehen.

Zusammenfassung in einem Satz

Die Forscher haben die "Suche" und das "Bauen" von Proteinstrukturen so stark optimiert, dass ein einzelner moderner Server heute so viel Leistung hat wie ein ganzes Rechenzentrum von vor ein paar Jahren – und das alles, ohne dass die Qualität der Ergebnisse leidet. Es ist, als hätten sie aus einem langsamen Fahrrad ein Formel-1-Auto gemacht, das trotzdem sicher fährt.

Technische Zusammenfassung

1. Problemstellung

Die Vorhersage von Proteinstrukturen mittels Deep Learning (DL), insbesondere mit AlphaFold2, hat die Strukturbiologie revolutioniert. Allerdings stoßen bestehende Pipelines an Grenzen, da die Nachfrage nach Vorhersagen exponentiell wächst, während das Mooresche Gesetz (die Leistungssteigerung von Hardware) möglicherweise nicht mehr ausreicht, um diesem Bedarf zu folgen.
Die aktuellen Workflows bestehen aus zwei Hauptphasen:

1. Homologiesuche (MSA-Generierung): Erstellung von Multiple Sequence Alignments (MSAs) durch Abgleich mit riesigen Datenbanken.
2. Deep-Learning-Inferenz: Vorhersage der 3D-Struktur basierend auf den MSAs.

Bestehende Lösungen sind oft rechenintensiv, langsam und nicht optimal auf moderne Hardware-Architekturen (wie NVIDIA Blackwell oder ARM-basierte Systeme) abgestimmt. Es besteht ein dringender Bedarf an einer Hardware-Software-Co-Design-Lösung, die den Durchsatz bei der Inferenz maximiert, ohne die Genauigkeit zu beeinträchtigen.

2. Methodik

Die Autoren entwickelten eine optimierte Pipeline, die OpenFold (eine Open-Source-Reimplementierung von AlphaFold2) und TensorRT (NVIDIAs Hochleistungs-Inferenz-Engine) kombiniert, ergänzt durch optimierte Homologiesuche mit MMseqs2-GPU.

Die wichtigsten technischen Optimierungen umfassen:

• OpenFold-TRT (TensorRT-Optimierung):

  • Präzisionsmanagement: Nutzung von gemischter Präzision (Mixed Precision). Der Evoformer-Kern (der rechenintensivste Teil) wird in BF16 (Brain Floating Point) ausgeführt, während der ExtraMSA-Teil in TF32 läuft. Dies beschleunigt die Berechnung erheblich bei Beibehaltung der Genauigkeit.
  • Dynamische Formen (Dynamic Shapes): Anstatt Modelle für feste Sequenzlängen neu zu kompilieren, wurden dynamische Profile erstellt, die Sequenzen von 16 bis zu maximalen Längen (z. B. 949 Aminosäuren) effizient verarbeiten.
  • Kernel-Fusion: Mehrstufige Attention-Operationen wurden zu einzelnen optimierten GPU-Kernels fusioniert, um den Speicherverkehr zu reduzieren und die arithmetische Intensität zu erhöhen.
  • Export: Nutzung von PyTorch's TorchDynamo für den ONNX-Export, um die dynamische Natur der Attention-Mechanismen zu erhalten.

• MMseqs2-GPU-Optimierungen:

  • Blackwell-Optimierung: Anpassung der gapless-Alignment-Kernel für die neue NVIDIA RTX PRO 6000 (Blackwell) Architektur. Nutzung der neuen DPX-Instruktionen (Dynamic Programming Instructions) und optimierte Verteilung der DP-Zellen pro Thread. Dies ermöglichte die Verarbeitung längerer Sequenzen und höhere Durchsätze.
  • ARM-Optimierungen: Für ARM-basierte Systeme (wie NVIDIA Grace-Hopper und DGX Spark) wurden die SIMD-Operationen (Single Instruction, Multiple Data) neu implementiert. Statt einer direkten Abbildung von x86 SSE auf ARM NEON wurden native NEON-Instruktionen (z. B. UMINV, UMAXV) und 256-bit Vektoroperationen genutzt, um die CPU-Effizienz zu verdoppeln.

• Hardware-Plattformen:
  Die Studie evaluierte verschiedene Szenarien:

  • x86 + RTX PRO 6000: Für maximale Geschwindigkeit auf Servern.
  • NVIDIA Grace-Hopper (GH200): Nutzt geteilten CPU-GPU-Speicher und hohe Chip-zu-Chip-Bandbreite (450 GB/s), um Speicherbeschränkungen bei großen Datenbanken zu umgehen.
  • DGX Spark: Ein kompaktes System für energieeffiziente Berechnungen.

3. Wichtige Beiträge

• OpenFold-TRT: Einführung einer TensorRT-basierten Inferenz-Pipeline für OpenFold, die den schnellsten AlphaFold2-artigen Inferenz-Workflow darstellt.
• End-to-End-Beschleunigung: Demonstration einer vollständigen Pipeline (MSA + DL), die auf einer einzigen GPU bis zu 131-fach schneller ist als das Baseline-AlphaFold2-Pipeline.
• Skalierbarkeit über GPU-Speicher: Nachweis, dass ARM-basierte Systeme (Grace-Hopper) durch geteilten Speicher und schnelle C2C-Interconnects Homologiesuche auch bei Datenbanken durchführen können, die größer sind als der verfügbare GPU-Speicher, ohne Geschwindigkeitseinbußen.
• Plattformunabhängige Effizienz: Optimierung für sowohl x86-Server als auch energieeffiziente ARM-Systeme (DGX Spark) und Superchips.

4. Ergebnisse

Die Evaluation erfolgte an 20 „harten" Zielen aus dem CASP14-Wettbewerb.

• Inferenzgeschwindigkeit (DL):

  • OpenFold-TRT auf einer RTX PRO 6000 ist 2,54-mal schneller als das vanilla OpenFold (PyTorch).
  • Im Vergleich zu AlphaFold2 (JAX) ist es 20,69-mal schneller.
  • Im Vergleich zu ColabFold-batch ist es 6,13-mal schneller.
  • Die durchschnittliche Zeit für die DL-Inferenz sank auf ca. 5,4 bis 5,6 Sekunden pro Protein.

• End-to-End-Geschwindigkeit (MSA + DL):

  • Die Kombination aus MMseqs2-GPU und OpenFold-TRT auf einer RTX PRO 6000 erreichte eine Gesamtzeit von 15,93 Sekunden pro Protein.
  • Dies entspricht einer 131,4-fachen Beschleunigung gegenüber dem Baseline-AlphaFold2-Pipeline (ca. 2425 Sekunden).
  • Gegenüber ColabFold-CPU ist es 5,94-mal schneller.

• Genauigkeit:

  • Die TM-Scores (Maß für strukturelle Ähnlichkeit) blieben über alle Optimierungen hinweg konsistent (ca. 0,67–0,71), was beweist, dass die Beschleunigung keine Genauigkeitsverluste mit sich bringt.

• Speicherskalierung (Grace-Hopper):

  • Auf dem GH200 (96 GB VRAM) blieb der Durchsatz konstant, selbst wenn die Datenbankgröße die GPU-Speicherkapazität überschritt (durch Nutzung des Host-Speichers). Auf x86-Systemen mit begrenztem VRAM (z. B. 48 GB L40S) brach der Durchsatz bei Überschreitung der Speichergrenze ein.

• Vergleich mit Boltz-2 (AlphaFold3-Reimplementierung):

  • Boltz-2 auf H100 war in der reinen DL-Inferenz schneller als OpenFold-TRT auf RTX PRO 6000, aber aufgrund langsamerer MSA-Generierung auf H100 war OpenFold-TRT in der End-to-End-Pipeline leicht überlegen (Faktor 1,06).

5. Bedeutung und Ausblick

Diese Arbeit zeigt, dass durch enge Hardware-Software-Integration massive Geschwindigkeitssteigerungen bei der Proteinstrukturvorhersage möglich sind, ohne auf die Genauigkeit von AlphaFold2 zu verzichten.

• Skalierbarkeit: Die beschleunigte Pipeline ermöglicht die Vorhersage von Milliarden von Proteinstrukturen in praktikablen Zeiträumen. Die Autoren berechnen, dass die Analyse der gesamten AlphaFold-Datenbank (350 Mio. Sequenzen) mit der alten Methode ca. 500 Jahre auf einem Server dauern würde, während die neue Pipeline dies in wenigen Monaten erledigen könnte.
• Demokratisierung: Durch die Optimierung für kompakte Systeme (DGX Spark) und energieeffiziente ARM-Architekturen wird High-Performance-Protein-Folding auch für kleinere Rechenzentren und lokale Installationen zugänglich.
• Zukunftsfähigkeit: Die vorgestellten Optimierungen (TensorRT, MMseqs2-GPU, ARM-Unterstützung) sind als Open-Source-Code verfügbar und bilden eine solide Basis für zukünftige generative Modelle im Protein-Design und die Verarbeitung noch größerer biologischer Datenmengen.

Zusammenfassend stellt das Paper einen Meilenstein dar, der die Lücke zwischen theoretischer DL-Leistung und praktischer, hochdurchsatzfähiger Anwendung in der Bioinformatik schließt.