NVIDIA stellt BioIR für beschleunigte Strukturvorhersage vor

NVIDIA stellt BioIR für beschleunigte Strukturvorhersage vor

NVIDIA hat mit BioNeMo Inference Runtime (BioIR) eine Python-Bibliothek vorgestellt, die die Inferenz von Biomolekül-Strukturmodellen auf NVIDIA-GPUs beschleunigt. Sie ist für Hochdurchsatz auf Proteom-Skala optimiert. Hintergrund ist, dass Strukturvorhersage zunehmend auf Proteom-Ebene ausgeführt wird, wo es darum geht, eine gesamte Arbeitsliste effizient durch die Pipeline zu bewegen. Allgemeine Inferenz-Stacks optimieren die spezialisierten Operationen solcher Modelle nicht vollständig. BioIR setzt hier an und bietet zwei Nutzungswege: einen End-to-End-Prozessor, der eine Eingabeanforderung durch Parsing, Tokenisierung, Merkmalsgenerierung, GPU-Inferenz und PDB- oder mmCIF-Ausgabe führt, sowie eine direkte PyTorch-Integration, bei der unterstützte Modelle als gewöhnliche torch.nn.Module konstruiert oder ausgewählte Module in eigenem Code wiederverwendet werden können. Die Bibliothek zielt auf spezialisierte Operationen wie Evoformer, Pairformer, Triangle Operations, Pairwise Attention, Diffusion Transformers und Atom-Level-Module, die in allgemeinen Stacks nicht vollständig optimiert werden. BioIR ist Teil des BioNeMo Frameworks und wurde bereits in realen Proteom-Projekten eingesetzt.

Optimierte Kernel, CUDA-Graphs und Ray-Parallelisierung

BioIR setzt auf drei Ebenen an: Kernel, Module/Layer und End-to-End-Pipeline. Die unterste Ebene bilden spezialisierte CuTeDSL-Kernels. Diese sind gezielt für architekturspezifische Operationen wie Evoformer-Stacks, Triangle-Operationen, Pairformer-Stacks, Aufmerksamkeit mit Paar-Bias (Pairwise Attention), Diffusion-Transformer und Atom-Level-Module implementiert. Allgemeine Inferenz-Stacks führen diese Operationen mit generischen PyTorch-Operationen aus, die nicht auf die speziellen Datenstrukturen und Rechenmuster dieser Biomolekül-Architekturen zugeschnitten sind. Die CuTeDSL-Kernels ersetzen diese generischen Pfade durch hochoptimierte GPU-Implementierungen, die den Durchsatz pro Operation erhöhen und gleichzeitig den GPU-Speicherbedarf senken, weil sie Zwischenergebnisse effizienter verwalten und weniger temporären Speicher benötigen.

Auf der zweiten Ebene bündelt BioIR wiederkehrende Abläufe mit CUDA-Graphs. Bei tiefen neuronalen Netzen entsteht durch das einzelne Starten vieler kleiner Kernel-Aufrufe erheblicher CPU-Launch-Overhead. CUDA-Graphs fassen diese Aufrufe zu einem Graphen zusammen, sodass die GPU sie als eine Einheit ausführen kann. Das reduziert den Overhead und steigert den Durchsatz, insbesondere bei Modellen mit vielen aufeinanderfolgenden Schichten.

Für große Mengen unabhängiger Eingaben kann Ray auf der Pipeline-Ebene eine vollständige Modellreplik auf jeder GPU eines einzelnen Knotens ausführen. Dadurch werden mehrere Eingaben parallel auf verschiedenen GPUs verarbeitet, statt sie sequenziell auf einer GPU abzuarbeiten. Das erhöht den Gesamtdurchsatz und senkt den Speicherdruck pro GPU, weil jede GPU nur eine Replik des Modells halten muss.

Trotz dieser Optimierungen bleibt ein BioIR-Modell ein gewöhnliches torch.nn.Module. Es gibt keinen Engine-Build, keinen Export-Schritt und kein separates Build-Artefakt zwischen Checkpoint und Vorwärtsdurchlauf. Entwickler können Tensoren inspizieren, Standard-Python-Stacktraces nutzen und die optimierten Layer über einfache Konfigurationsfelder auswählen, mit dokumentierten PyTorch-Fallbacks für unterstützte Pfade. Das erhält den vertrauten Python-Entwicklungsworkflow und erleichtert Debugging und Anpassung. Unterstützt werden Modelle wie OpenFold3, Protenix, Boltz-2 und OpenFold2.

Benchmarks: Deutliche Vorteile auf Datacenter-GPUs

Die veröffentlichten Benchmarks vergleichen BioIR mit zwei Open-Source-PyTorch-Baselines: PyTorch eager und torch.compile, auf jeder getesteten GPU. Die Zahlen sind reproduzierbar, nicht Spitzenwerte, und es wurde keine Konfiguration pro GPU getunt. Gegenüber PyTorch eager liegt BioIR auf allen getesteten Datacenter-GPUs vorn, mindestens mit Faktor 1,19 – außer auf GB10, wo Protenix 0,90x und OpenFold3 0,99x erreicht. Gegenüber torch.compile hängt das Ergebnis von der Architektur ab: 1,5- bis 3,0-fach auf Ampere, Hopper und Ada, 1,0- bis 2,1-fach auf Blackwell und 0,8- bis 1,7-fach auf GB10. torch.compile ist bei 21 Prozent der Einzelmessungen vorn, und das nicht zufällig: nie unter 256 Residuen, steigend mit der Länge auf rund 29 Prozent über 769 Residuen – und nur bei Boltz-2 und OpenFold3, nie bei OpenFold2 oder Protenix. Lange Eingaben auf Blackwell sind der Fall, den man vor einer einzelnen Kennzahl prüfen sollte. Die Ausnahme GB10 erklärt sich als Desktop-Teil mit 128 GB vereinheitlichtem Speicher, das anders skaliert. Da das Zitieren nur einer Baseline das Ergebnis in eine Richtung über- oder untertreiben würde, werden auch Werte unter 1,00x ausgewiesen statt gestrichen.

Einsatz bei AlphaFold Database: 1,81 Millionen neue Vorhersagen

BioIR wurde bereits in realer Proteom-Arbeit eingesetzt, konkret bei der jüngsten Erweiterung der AlphaFold Database. Dabei beschleunigte die Laufzeitumgebung die Erzeugung von Protein-Komplex-Strukturen über 4.777 Proteome hinweg. Insgesamt wurden etwa 31 Millionen Kandidatenkomplexe verarbeitet, von denen 1,81 Millionen als hochkonfidente Vorhersagen veröffentlicht wurden. Der Einsatz zeigt, dass BioIR für Durchsatz auf Proteom-Skala ausgelegt ist: Statt einzelne Strukturen isoliert zu berechnen, geht es darum, eine gesamte Arbeitsliste effizient durch die Pipeline zu bewegen. Die Verarbeitung von 31 Millionen Kandidatenkomplexen verdeutlicht den Umfang. Die Zahlen stammen aus der NVIDIA-Dokumentation und belegen den praktischen Nutzen der Bibliothek in einem der größten öffentlichen Strukturdatenbestände.

Hardware-Anforderungen und Open-Source-Lizenz

Für den Betrieb von BioIR-basierten NIM-Microservices gelten unterschiedliche Mindestanforderungen. Der OpenFold3 NIM benötigt mindestens 64 GB RAM, eine CPU mit mindestens 8 Kernen und 80 GB Speicherplatz. Das Container-Image ist etwa 11 GB groß beim Download und rund 40 GB nach dem Entpacken; Modellparameter werden beim ersten Start in den NIM-Cache geladen und benötigen weitere 15 GB. Unterstützte GPUs sind getestet, jede Vorhersage läuft auf einer einzelnen GPU. Der ProteinMPNN NIM verlangt eine GPU mit mindestens 3 GB Grafikspeicher und Compute Capability größer 7.0; empfohlen werden getestete Konfigurationen, zusätzlich sind 10 GB NVMe-SSD-Speicher nötig. Für lokale BioNeMo-Installationen wird ein System mit NVIDIA L40S oder H100 und schnellem lokalem NVMe-Speicher empfohlen. Die Tutorial-Schätzung liegt bei rund 700 GB für eine UniRef30-only MSA-Datenbank, davon etwa 490 GB für die Datenbank selbst, plus 30 bis 40 GB für Boltz-2- und OpenFold3-Container. Die vollständige MSA-Datenbank ist mit etwa 1,4 TB deutlich größer. Das BioNeMo Framework, zu dem BioIR gehört, ist unter der Apache-2.0-Lizenz veröffentlicht und kostenlos nutzbar, auch für kommerzielle Zwecke.