NVIDIA BioNeMo nutzt Context Parallelism für Biomodelle
Seit Jahrzehnten war die computergestützte Biologie durch hardwarebedingte Speicherbeschränkungen gezwungen, komplexe Biomoleküle in isolierte Fragmente zu zerlegen. Diese Reduktionismus-Methode erzeugte eine Lücke im globalen Kontext, was es unmöglich machte, große Proteinkomplexe vollständig und ohne Vorwissen zu modellieren. NVIDIA BioNeMo hat nun mit einem neuen Framework für Context Parallelism (CP) diese Barrieren durchbrochen. Dieses System ermöglicht die ganzheitliche Modellierung molekularer Systeme, indem es ein einziges großes Molekül über mehrere GPUs verteilt, anstatt es in kleinere Teile zu schneiden. Traditionelle Workarounds wie das Überlappen von Sequenzsegmenten oder das Chunking innerhalb der Modelarchitektur bewahren zwar den lokalen Speicher, zerstören aber Informationen über weit entfernte Wechselwirkungen wie Allostery oder Signalübertragung. Das BioNeMo CP-Framework umgeht dies, indem es die Repräsentation eines einzelnen Proteins über einen 2D-Gitterbereich aufteilt. Ein Beispiel: Bei einem Komplex mit 10.000 Resten, was 100 Millionen Interaktionen entspricht, verwaltet jede GPU nur einen spezifischen Teilblock. Dadurch sinkt der Speicherbedarf pro Gerät von O(N²) auf O(N²/P). Die Implementierung nutzt Torch Distributed APIs, um eine multidimensionale Sharding-Strategie zu ermöglichen. Ein wesentlicher Mechanismus ist die Überlappung von Berechnung und Kommunikation. Während eine GPU eine lokale Aktualisierung berechnet, führt sie gleichzeitig asynchrone Datenübertragungen mit benachbarten Geräten aus. Dies verbessert die Effizienz bei wachsender Problemgröße. Auch spezielle Aufmerksamkeitsmechanismen wie die Atomsequenz-Beobachtung wurden angepasst, um mit Hilfe von Halo-Austausch-Primitiven keine zusätzliche Kommunikation zwischen GPUs für Fenster-Batching zu benötigen. Die praktische Anwendung des Frameworks hat bereits bemerkenswerte Ergebnisse gezeigt. Das Team konnte das Boltz-Modell nutzen, um ein System mit 3.605 Resten, bestehend aus vier Ketten, auf vier NVIDIA H100-GPUs zu falten. Dies übertrifft die ursprüngliche Trainingsgröße von 768 Resten bei weitem. Die Generierung von fünf Strukturproben dauerte weniger als fünf Minuten, wobei alle langreichweitigen Kontakte zwischen den Untereinheiten korrekt modelliert wurden. Ohne zusätzliches Training oder Feinabstimmung gelang dies dank der verteilten Speicherkapazität. Die Technologie wird bereits von externen Partnern wie Rezo Therapeutics, Proxima und Earendil Labs integriert. Rezo Therapeutics nutzt das Framework zur Vorhersage riesiger Protein-Protein-Wechselwirkungen mit bis zu 6.500 Resten, was eine dreifach höhere Entdeckung neuartiger Komplexe im Vergleich zu herkömmlichen Methoden ermöglicht. Proxima hat es in ihr atomares generatives Modell integriert, um therapeutisch relevante Wechselwirkungen aufzudecken, während Earendil Labs die Sequenzlängen in ihrem eigenen Modell erweitert hat, um bisher rechnerisch zu komplexe Systeme zu analysieren. Trotz der technischen Durchbrüche warnt die NVIDIA-Gruppe davor, dass physische Kapazität allein keine biologische Genauigkeit garantiert. Da viele Modelle auf kleinen Fragmenten trainiert wurden, ist es entscheidend, diese mit längeren Sequenzabschnitten nachzufinieren, um die Logik langreichweitiger Interaktionen korrekt zu erfassen. Um dieses Datenproblem zu lösen, arbeitet das Team an der Nutzung von NVIDIA-Beschleunigungssoftware, um die AlphaFold Protein Structure Database mit hochdurchsatzfähigen Vorhersagen großer Komplexe zu füllen. Dies schafft die notwendige Grundlage für neue Foundation-Modelle, die größere biologische Systeme repräsentieren können. Das Open-Source-Code-Framework steht Entwicklern und Forschern nun zur Verfügung, um die Skalierbarkeit der biomolekularen Modellierung voranzutreiben.
