NVIDIA BioNeMo beschreibt ein Context-Parallelism-Framework (CP), das die Speichergrenzen einzelner GPUs bei der Strukturvorhersage großer biomolekularer Systeme umgehen soll. Statt große Proteine oder Komplexe in Fragmente zu zerlegen, verteilt CP ein einzelnes großes Sample über mehrere GPUs, sodass der globale Kontext erhalten bleibt.
Als bisherige Workarounds nennt der Beitrag die Aufteilung langer Sequenzen in überlappende Segmente sowie modellinternes Chunking zur Reduktion des VRAM-Bedarfs. Beide Ansätze helfen bei großen Systemen, verlieren laut NVIDIA aber insbesondere bei langen Distanzen wichtigen Kontext. Genannt werden etwa Einschränkungen bei der Modellierung von Allosterie oder Signalübertragung über einen gesamten Komplex hinweg.
Die BioNeMo-Implementierung baut auf den verteilten Torch-APIs für GPU-zu-GPU-Kommunikation auf und wird im Beitrag am Beispiel von Boltz erläutert. Für lineare Kapazitätsskalierung verwendet das Framework eine mehrdimensionale Sharding-Strategie, bei der kein einzelnes Gerät den vollständigen globalen Zustand des Biomoleküls hält. Eine globale Matrix der Größe N × N wird in ein Gitter aus Blöcken aufgeteilt; bei einem Komplex mit 10.000 Resten und damit 100 Millionen Interaktionen verwaltet jede GPU nur einen Teilblock. Der Speicherbedarf pro Gerät sinkt dadurch von O(N²) auf O(N²/P).
Die verteilten Primitive kombinieren lokale Berechnung mit asynchronen Peer-to-Peer-Transfers. Während eine GPU lokale Updates berechnet, sendet und empfängt sie gleichzeitig Daten an Nachbarn in Zeilen- und Spalten-Ringen. Laut NVIDIA verbessert sich mit wachsender Problemgröße das Verhältnis von Rechenaufwand zu Kommunikation, was das System auf größeren Skalen effizienter macht.
Für die in AlphaFold3 verwendete lokale Atom-Attention mit Fenstern der Größe 32 × 128 musste die Kachelstruktur der Atom-Features in der verteilten Variante neu partitioniert werden. NVIDIA setzt dafür auf verteilte Primitive auf Basis von Halo Exchange, sodass die nachfolgende Window-Batch-Attention ohne weitere Inter-GPU-Kommunikation auskommt.
Der Beitrag zeigt außerdem ein Beispiel für eine CP-fähige Schicht zur verteilten Berechnung von Triangle Multiplication. Dazu wird mit dem DistributedManager eine quadratische 2D-Device-Mesh aufgebaut, um symmetrische Kommunikationsmuster entlang von Zeilen und Spalten zu erhalten. Ein Ring2DComm-Handle übernimmt die Peer-to-Peer-Kommunikation, während eine serielle Schicht wie TriangleMultiplicationOutgoing zunächst auf der CPU geladen und anschließend mit DistributedTriangleMultiplication für DTensors in die verteilte Ausführung überführt wird. So bleiben große Aktivierungstensoren über das Device-Gitter verteilt.
Nach Angaben von NVIDIA lassen sich mit CP Skalierungsgesetze für die Token-Kapazität auch für biomolekulare Architekturen nutzen. Boltz-Vorhersagen seien damit auf bis zu etwa 20.000 Tokens mit 256 GPUs möglich; auf NVIDIA H100-GPUs lasse sich die maximale Token-Länge skalieren, mit beschleunigter Skalierung auf NVIDIA B300-GPUs.
Ohne zusätzliches Training oder Fine-Tuning auf längeren Crop-Längen habe das Team ein TTC7A/PI4KA/FAM126A/EFR3A(700–823)-System mit 3.605 Resten über vier Ketten gefaltet. Das liegt laut Beitrag deutlich über der Boltz-2-Trainings-Crop-Größe von 768 Resten und über der Speicherkapazität einer einzelnen GPU. Mit CP seien auf vier NVIDIA H100-GPUs fünf Struktur-Samples in unter fünf Minuten erzeugt worden, also etwa 54 Sekunden pro Sample, wobei alle langreichweitigen Kontakte zwischen Untereinheiten innerhalb des Modell-Kontextfensters erhalten blieben.
NVIDIA nennt außerdem Arbeiten mit Rezo Therapeutics, Proxima und Earendil Labs. Rezo Therapeutics habe das Framework für Vorhersagen großer Protein-Protein-Interaktionen bis 6.500 Reste integriert und berichte von einer mehr als dreifachen Anreicherung CP-aufgelöster, qualitativ hochwertiger neuer Proteinkomplexe im Vergleich zu Vorhersagen, die nur hochvertrauenswürdige öffentliche PPI-Daten nutzen. Proxima habe CP in das all-atom generative foundation model Neo eingebettet und damit Inferenz auf Assemblies bis 4.000 Tokens ermöglicht. Earendil Labs habe die Eingabesequenzlängen seines proprietären biomolekularen foundation models erweitert, um komplexe Multi-Protein-Systeme zu modellieren.
Der Beitrag betont zugleich Grenzen des Ansatzes. Größere physische Kapazität garantiere keine biologische Genauigkeit, weil aktuelle Modelle bei hochskalierter Faltung oft Schwierigkeiten hätten, wenn sie nur auf kleinen Fragmenten trainiert wurden. Für eine präzisere Erfassung langreichweitiger Interaktionen sei Fine-Tuning mit größeren Crop-Größen erforderlich.
Als weiteres Problem nennt NVIDIA die knappe Datenlage. Daran werde unter anderem durch Beiträge zur AlphaFold Protein Structure Database gearbeitet. Mit Software wie NVIDIA cuEquivariance und NVIDIA TensorRT sollen hochdurchsatzfähige Vorhersagen großer homo- und heteromerer Komplexe zur AFCDB beitragen, um synthetische Komplexdaten für das Training von Foundation Models für größere biologische Systeme aufzubauen.
