🧬 biology

MegaFold: Efficient Training of Next-Generation 3D Attention Protein Models on Cross-Platform GPUs

MegaFold ist ein neuartiges plattformübergreifendes System, das die Speicher- und Rechenengpässe beim Training von 3D-Attention-Proteinfunktionsmodellen der nächsten Generation durch die Kombination von speichereffizienten Kernels, optimiertem Sharding, fusionierten Operatoren und einer Determinismus-bewussten Pipeline überwindet, um signifikant längere Sequenzlängen und schnellere Trainingszeiten sowohl auf NVIDIA- als auch auf AMD-GPUs zu erreichen.

Ursprüngliche Autoren: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Veröffentlicht 2026-06-16

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ Dies ist eine KI-generierte Erklärung eines Preprints, das nicht peer-reviewed wurde. Dies ist kein medizinischer Rat. Treffen Sie keine Gesundheitsentscheidungen auf Grundlage dieses Inhalts. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, ein massives, kompliziertes 3D-Puzzle eines Proteins zusammenzubauen. Lange Zeit nutzten Wissenschaftler einen bestimmten Satz von Regeln (wie das alte „AlphaFold2“), der gut funktionierte. Aber die neueste Version, AlphaFold3, ist wie der Übergang von einem flachen 2D-Puzzle zu einer komplexen, mehrschichtigen 3D-Struktur, bei der jedes einzelne Teil mit jedem anderen Teil im 3D-Raum interagiert.

Obwohl dieses neue Modell wissenschaftlich brillant ist, hat es ein riesiges Problem: Es verschlingt Computerarbeitsspeicher wie ein Schwarzes Loch und läuft unglaublich langsam.

Dieses Paper stellt MegaFold vor, ein neues „System“ (eine Sammlung spezialisierter Werkzeuge und Regeln), das es möglich macht, diese neuen 3D-Proteinstrukturen auf modernen Computern zu trainieren, ohne dass diese abstürzen oder ewig dauern.

Hier ist, wie MegaFold die vier Hauptprobleme löst, erklärt mit einfachen Analogien:

Das Problem: Warum der alte Weg scheitert

Das neue AlphaFold3-Modell betrachtet Proteine auf eine Weise, die eine „kubische“ Datenexplosion erzeugt.

Die Analogie: Stellen Sie sich vor, Sie organisieren eine Party. Im alten Modell mussten Sie nur wissen, wer mit wem befreundet ist (eine einfache Liste). Im neuen Modell müssen Sie eine 3D-Karte davon erstellen, wie jeder Gast gleichzeitig mit jedem anderen Gast interagiert. Wenn Sie 100 Gäste haben, benötigt der alte Weg 10.000 Notizen; der neue Weg benötigt 1.000.000 Notizen. Wenn Sie versuchen, all diese Notizen auf ein einziges Notizbuch zu schreiben (den Arbeitsspeicher des Computers), explodiert das Notizbuch.

Die Lösung: Die vier Werkzeuge von MegaFold

MegaFold behebt dies durch vier spezifische Upgrades:

1. Das „Schlaue Notizbuch“ (EvoFlash-3D)

Das Problem: Der Computer versucht, die gesamte 3D-Interaktionskarte auf einmal aufzuschreiben, was seinen Speicher sofort füllt.
Die MegaFold-Lösung: Anstatt die ganze Karte auf das Notizbuch zu schreiben, nutzt MegaFold einen „Schmierblock“ (schnellen, temporären Speicher), um die Interaktionen in kleinen Häppchen, eins nach dem anderen, zu berechnen und die Notizen sofort wieder wegzuwerfen.
Das Ergebnis: Er muss nie die gesamte massive Karte gleichzeitig im Speicher halten. Dies ermöglicht es dem Computer, viel längere Proteinketten zu verarbeiten, ohne den Platz zu verlieren.

2. Das „Staffellauf-Team“ (EvoSP-3D)

Das Problem: Wenn Sie mehrere Computer (GPUs) nutzen, um zusammenzuarbeiten, teilen diese die Arbeit normalerweise dadurch auf, dass jeder Computer eine andere Gästeliste erhält. Aber da die Daten von AlphaFold3 ein 2D-Gitter aus Interaktionen sind, funktioniert das einfache Aufteilen nicht; die Computer werden verwirrt darüber, wer mit wem spricht.
Die MegaFold-Lösung: MegaFold erfindet einen neuen Weg, wie die Computer Daten übergeben können. Anstatt nur die Liste aufzuteilen, teilen sie das Gitter der Interaktionen auf. Sie geben Teile des Puzzles in einem spezifischen, effizienten Muster (wie bei einem Staffellauf) hin und her, sodass jeder Computer genau weiß, für welchen Teil der 3D-Karte er verantwortlich ist, ohne Zeit mit dem Warten auf andere zu verschwenden.

3. Das „Fließband“ (EvoFusion)

Das Problem: Der Computer erledigte die Aufgabe in winzigen, ineffizienten Schritten. Er berechnete eine Zahl, stoppte, speicherte sie, lud ein neues Werkzeug, berechnete erneut, stoppte und speicherte wieder. Das ist so, als würde ein Koch eine Zwiebel schneiden, dann zum Kühlschrank gehen, um ein Messer zu holen, und dann zurückgehen, um weiterzuschneiden.
Die MegaFold-Lösung: MegaFold kombiniert diese winzigen Schritte zu einer einzigen, flüssigen Bewegung. Es verschmilzt das „Schneiden“, „Mischen“ und „Kochen“ zu einer einzigen, kontinuierlichen Aktion.
Das Ergebnis: Der Computer hört auf, Zeit mit dem Stoppen und Starten von Aufgaben zu verschwenden. Er läuft viel schneller, weil er die „Fließbandarbeit“ ohne Unterbrechung fortsetzt.

4. Die „Vorbereiteten Zutaten“ (EvoPipe)

Das Problem: Bevor der Computer überhaupt mit dem Bau des Proteins beginnen kann, muss ein Mensch (oder eine langsame CPU) viel mühsame Recherche betreiben, um die richtigen Daten zu sammeln (wie die Suche nach der Evolutionsgeschichte). Das dauert lange und lässt den leistungsstarken Computer untätig herumsitzen, während er auf die Daten wartet.
Die MegaFold-Lösung: MegaFold erkennt, dass ein Teil dieser Recherche für dasselbe Protein immer gleich bleibt. Es erledigt die harte Arbeit einmal im Voraus und speichert sie in einer „Speisekammer“ (einem Cache). Wenn der Computer die Daten braucht, greift er einfach sofort auf die vorbereiteten Zutaten zu.
Das Ergebnis: Der leistungsstarke Computer muss nie warten. Er ist zu 100 % ausgelastet.

Die Ergebnisse

Das Paper hat MegaFold auf zwei Arten von leistungsstarken Computerchips (NVIDIA und AMD) getestet. Hier ist, was sie herausgefunden haben:

Längere Proteine: MegaFold ermöglichte es dem Team, auf Proteinsequenzen zu trainieren, die 3,36-mal länger waren als zuvor.
Schnellere Geschwindigkeit: Es machte den Trainingsprozess auf NVIDIA-Chips 1,73-mal schneller und auf AMD-Chips 1,62-mal schneller.
Weniger Speicher: Es verbrauchte signifikant weniger Speicher und verhinderte die „Out of Memory“-Abstürze, die bei diesen Modellen normalerweise auftreten.

Zusammenfassung

Betrachten Sie AlphaFold3 als einen superstarken Motor, der jedoch zu schwer für sein Auto-Chassis (das Computersystem) war. MegaFold hat nicht nur ein größeres Auto gebaut; es hat den Motor, das Getriebe und das Kraftstoffsystem neu gestaltet, damit der superstarke Motor tatsächlich effizient laufen kann. Dies ermöglicht es Wissenschaftlern, diese nächsten Generationen von Modellen auf einer breiteren Palette von Computern zu trainieren, was die Wissenschaft der Proteinfaltung schneller und zugänglicher macht.

Technisches Resümee: MegaFold

Problemstellung

Jüngste Fortschritte in der biomolekularen Modellierung, insbesondere AlphaFold3 (AF3), haben einen Paradigmenwechsel von Standard-Transformer-Architekturen hin zu Modellen eingeführt, die 3D-Attention über 2D-Paarweisen-Repräsentationen nutzen. Während AF3 präzise Vorhersagen von Protein-Protein-, Protein-DNA/RNA- und Protein-Ligand-Interaktionen ermöglicht, führt diese architektonische Änderung schwerwiegende systemseitige Engpässe ein, die durch bestehende Optimierungen für Large Language Models (LLMs) nicht adressiert werden können.

Das Paper identifiziert vier fundamentale Herausforderungen, die auf der kubischen Skalierung der 3D-Attention beruhen:

Kubische Aktivierungs-Speicherkapazität: Im Gegensatz zu Transformern, bei denen die Attention quadratisch skaliert ( $O(N^2)$ ), erzeugt die 3D-Attention von AF3 über 2D-Paarweisen-Repräsentationen Tensoren, die kubisch skalieren ( $O(N^3)$ ). Dies führt dazu, dass der Aktivierungs-Speicher die GPU-Nutzung dominiert (bis zu 96 % des Gesamtspeichers) und selbst bei kleinen Batch-Größen bei moderaten Sequenzlängen (z. B. >590 Residuen) zu Out-of-Memory (OOM)-Fehlern führt.
Inkompatibilität mit Sequenz-Parallelität (SP): Standardmäßige SP-Strategien für Transformer basieren auf 1D-Sequenz-Sharding. Die 2D-Paarweisen-Repräsentationen von AF3 und die alternierenden Reduktionsachsen über die Layer hinweg erzeugen komplexe Multi-Achsen-Kommunikationsmuster, die konventionelle SP ohne übermäßiges Resharding oder Synchronisation nicht unterstützen kann.
Kernel-Fragmentierung: Das Verweben von 3D-Attention, Transitions-Layern und Projektionsoperationen über gemischte Dimensionen hinweg resultiert in Tausenden von feinkörnigen Kernels pro Iteration. Diese Fragmentierung verhindert eine effektive Kernel-Fusion, was zu einer schlechten Warp-Occupancy, hohen Launch-Overheads und einer ineffizienten Nutzung der Speicherbandbreite führt.
Host-seitige Preprocessing-Stalls: Die Konstruktion des evolutionären Kontextes (z. B. Multiple Sequence Alignments, MSA), der für die 3D-Attention erforderlich ist, erfordert rechenintensive, abhängigkeitsreiche CPU-Vorverarbeitung. Dies erzeugt signifikante GPU-Leerlaufzeiten, die nicht durch konventionelle Prefetching-Pipelines verborgen werden können.

Methodik: MegaFold

MegaFold ist ein plattformübergreifendes System, das darauf ausgelegt ist, diese Herausforderungen durch ein Co-Design aus speichereffizienten Kernels, Kommunikationsstrategien, Operator-Fusion und Host-Device-Pipelining zu adressieren. Es ist auf einem Open-Source AF3-Trainingsframework implementiert und unterstützt sowohl NVIDIA- als auch AMD-GPUs.

1. EvoFlash-3D: Speichereffiziente 3D-Attention

Um das kubische Speicherwachstum zu adressieren, führt MegaFold EvoFlash-3D ein, einen Triton-basierten Kernel für 3D-Attention.

Mechanismus: Anstatt die vollständige $O(N^3)$ Attention-Map im globalen Speicher (HBM) zu materialisieren, nutzt der Kernel eine Online-Strategie für numerisch stabile Softmax-Berechnungen. Er berechnet die Attention-Logits inkrementell innerhalb eines schnellen Scratchpad-Speichers (Shared Memory) während des Forward- und Backward-Passes.
Optimierung: Dies reduziert die Spitzenkomplexität des Aktivierungs-Speichers von $O(N^3)$ auf $O(N^2)$ , während die numerische Korrektheit gewahrt bleibt. Der Kernel handhabt den einzigartigen additiv broadcasteten Pair-Bias-Term von AF3, der multidimensionale Datenabhängigkeiten einführt, die in Standard-FlashAttention nicht vorhanden sind.

2. EvoSP-3D: Kommunikations-effizientes Sharding

Um das Training auf längeren Sequenzen zu ermöglichen, schlägt MegaFold EvoSP-3D vor, eine auf 2D-Paarweisen-Repräsentationen zugeschnittene Sharding-Strategie.

Mechanismus: Es nutzt einen hybriden Data Parallel (DP) und Sequence Parallel (SP) Ansatz. Während DP-Gruppen unterschiedliche Batch-Inputs verarbeiten, partitionieren SP-Gruppen den 2D-Paarweisen-Tensor ( $[N_{token}, N_{token}, h_{pair}]$ ).
Kommunikation: Da die AF3-Layer die Attention-Achsen alternieren (z. B. zeilenweise zu spaltenweise), nutzt EvoSP-3D all_to_all-Kollektive, um das Sharding-Layout des Paarweisen-Tensors zwischen den Stufen zu rotieren. Dies ermöglicht es den aufeinanderfolgenden Attention-Layern, auf den korrekten 2D-Slices zu operieren, ohne den vollständigen Tensor zu materialisieren.
Diffusions-Stadium: Im Diffusions-Decoder werden augmentierte Atom-Koordinaten über SP-Gruppen hinweg verteilt, ohne Kommunikation zu benötigen, während die Paarweisen-Repräsentation mittels all_gather zusammengeführt wird, um die Diffusion auf dem vollen Kontext operieren zu lassen.

3. EvoFusion: Gefügter Operator-Stack

Um die Kernel-Fragmentierung zu mildern, führt MegaFold EvoFusion ein, welches die interne Ausführung kritischer Operatoren restrukturiert, anstatt lediglich deren Grenzen zu fusionieren.

Umfang: Zielt auf Layer Normalization, Lineare Projektionen und SwiGLU-Aktivierungen ab, die zusammen etwa 20 % des Aktivierungs-Speichers und der Kernel-Aufrufe ausmachen.
Mechanismus: Durch die Umstrukturierung der inneren Schleifen dieser Operatoren eliminiert EvoFusion die redundante Materialisierung Zwischenschritte (z. B. das Vermeiden des Schreibens eines vollständigen $N \times N$ Tensors in den HBM zwischen Normalisierung und linearer Projektion). Dies reduziert den Shared-Memory-Druck und den Kernel-Launch-Overhead bei gleichzeitiger Aufrechterhaltung effizienter Reduktions- und GeMM-Schedules.

4. EvoPipe: Deterministik-bewusste Host-Device-Pipeline

Um die durch CPU-Preprocessing verursachten GPU-Leerlaufzeiten zu eliminieren, implementiert MegaFold EvoPipe.

Mechanismus: Es entkoppelt deterministische Vorverarbeitungsschritte (z. B. RDKit-basierte Atomkonstruktion, MSA-Feature-Extraktion) vom stochastischen Trainingsloop. Diese deterministischen Features werden offline vorberechnet und gecacht.
Ausführung: Während des Trainings führt das System leichtgewichtige Lookups durch, um gecachte Features abzurufen, wobei nur nicht-deterministische Augmentationen (z. B. Random Cropping) on-the-fly angewendet werden. Dies gewährleistet eine konsistente GPU-Auslastung durch Beseitigung variabler Host-seitiger Stalls.

Wichtigste Ergebnisse

Das System wurde auf NVIDIA H200 (141 GB) und AMD MI250 (64 GB) GPUs unter Verwendung des AlphaFold3-PyTorch-Frameworks evaluiert.

Skalierbarkeit der Sequenzlänge: MegaFold ermöglicht das Training mit bis zu 3,36-mal längeren Sequenzlängen auf 32 GPUs im Vergleich zu den Baseline-PyTorch-Implementierungen. Auf 32 AMD-GPUs unterstützt es Sequenzen bis zu ~1450 Residuen, während Baselines bereits bei ~590 Residuen scheitern.
Ausführungszeit: MegaFold reduziert die End-to-End-Zeit pro Iteration um bis zu 1,73× auf NVIDIA und 1,62× auf AMD GPUs.
Speichereffizienz: Die Spitzen-GPU-Speichernutzung wird um bis zu 1,23× reduziert (bzw. um 23 % gesenkt), was die Umverteilung der Speicherkapazität auf längere Sequenzen ermöglicht.
Modell-FLOPs-Ausnutzung (MFU): MegaFold erreicht eine signifikant höhere MFU (bis zu 44,9 % bei 1,81 Mrd. Parametern) im Vergleich zu Baselines, die oft OOM-Limits erreichen, bevor sie größere Skalen erreichen.
Plattformübergreifende Performance: Das System zeigt konsistente Gewinne über NVIDIA- und AMD-Hardware hinweg, was die Portabilität der Triton-basierten Kernel-Designs validiert.

Bedeutung und Ansprüche

Das Paper behauptet, dass MegaFold das erste End-to-End-System ist, das integrierte Unterstützung für das Training der nächsten Generation von 3D-Attention-Proteinfunktionen bietet. Seine Bedeutung liegt in dem Nachweis, dass:

LLM-Optimierungen unzureichend sind: Standardmäßige Transformer-Optimierungen (FlashAttention, Standard-SP, Compiler-Fusion) scheitern daran, die einzigartige kubische Skalierung und die gemischt-dimensionalen Abhängigkeiten von AF3-ähnlichen Modellen zu adressieren.
Systemseitiges Co-Design essenziell ist: Effizientes Training dieser wissenschaftlichen Modelle erfordert dedizierte Systemunterstützung, die Kernel-Ausführung, Parallelisierung, Operator-Fusion und Host-Device-Pipelining gemeinsam adressiert.
Die Barriere der Skalierbarkeit verschoben wird: Durch die Adressierung der spezifischen Engpässe der 3D-Attention verschiebt MegaFold die OOM-Barriere und ermöglicht so das Training größerer biomolekularer Komplexe, die zuvor nicht realisierbar waren, wodurch das Scaling aufkommender wissenschaftlicher Workloads auf modernen GPU-Plattformen erleichtert wird.