🧬 biology

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

MegaFold is een nieuw cross-platform systeem dat de geheugen- en computationele knelpunten van het trainen van volgende generatie 3D-attention proteïnemodellen overwint door geheugenefficiënte kernels, geoptimaliseerde sharding, gefuseerde operatoren en een determinisme-bewuste pipeline te combineren om aanzienlijk langere sequentielengtes en snellere trainingstijden te bereiken op zowel NVIDIA- als AMD-GPU's.

Oorspronkelijke auteurs: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Gepubliceerd 2026-06-16

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ Dit is een AI-gegenereerde uitleg van een preprint die niet peer-reviewed is. Dit is geen medisch advies. Neem geen gezondheidsbeslissingen op basis van deze inhoud. Lees de volledige disclaimer

Stel je voor dat je probeert een enorme, ingewikkelde 3D-puzzel van een eiwit te bouwen. Een lange tijd gebruikten wetenschappers een specifieke set regels (zoals de oude "AlphaFold2") die goed werkten. Maar de nieuwste versie, AlphaFold3, is als de upgrade van een platte 2D-puzzel naar een complexe, meerlagige 3D-structuur waarbij elk stukje met elk ander stukje interageert in een 3D-ruimte.

Hoewel dit nieuwe model wetenschappelijk gezien geweldig is, heeft het een groot probleem: het vreet computermemori als een zwart gat en draait ongelooflijk traag.

Dit paper introduceert MegaFold, een nieuw "systeem" (een set gespecialiseerde tools en regels) dat het mogelijk maakt om deze nieuwe 3D-eiwitmodellen te trainen op moderne computers zonder dat ze vastlopen of eeuwig duren.

Hier is hoe MegaFold de vier hoofdproblemen oplost, uitgelegd met eenvoudige analogieën:

Het Probleem: Waarom de oude manier faalt

Het nieuwe AlphaFold3-model bekijkt eiwitten op een manier die een "kubische" explosie van data creëert.

De Analogie: Stel je voor dat je een feestje organiseert. In het oude model hoefde je alleen maar te weten wie vrienden met wie was (een simpele lijst). In het nieuwe model moet je een 3D-kaart bijhouden van hoe elke gast met elke andere gast tegelijkertijd interageert. Als je 100 gasten hebt, heeft de oude manier 10.000 aantekeningen nodig; de nieuwe manier heeft 1.000.000 aantekeningen nodig. Als je al die aantekeningen op één enkel notitieblok probeert te schrijven (het geheugen van de computer), ontploft het blokje.

De Oplossing: MegaFold's Vier Tools

MegaFold lost dit op door vier specifieke upgrades te introduceren:

1. Het "Slimme Notitieblok" (EvoFlash-3D)

Het Probleem: De computer probeert de volledige 3D-interactiekaart in één keer op te schrijven, waardoor het geheugen direct vol raakt.
De MegaFold Fix: In plaats van de hele kaart op het notitieblok te schrijven, gebruikt MegaFold een "kladblok" (snel, tijdelijk geheugen) om de interacties in kleine stukjes te berekenen, één voor één, en gooit de aantekeningen daarna direct weer weg.
Het Resultaat: Het hoeft nooit de hele enorme kaart tegelijk in het geheugen te houden. Dit stelt de computer in staat om veel langere eiwitketens aan te kunnen zonder dat de ruimte opraakt.

2. De "Teamestafette" (EvoSP-3D)

Het Probleem: Wanneer je meerdere computers (GPU's) gebruikt om samen te werken, verdelen ze de taken meestal door elke computer een andere gastentlijst te geven. Maar omdat de data van AlphaFold3 een 2D-raster van interacties is, werkt het simpelweg splitsen niet; de computers raken in de war over wie met wie praat.
De MegaFold Fix: MegaFold vindt een nieuwe manier voor de computers om data door te geven. In plaats van alleen de lijst te splitsen, splitsen ze het raster van interacties. Ze geven stukjes van de puzzel in een specifief, efficiënt patroon (zoals een estafette) heen en weer, zodat elke computer precies weet welk deel van de 3D-kaart hij verantwoordelijk is, zonder tijd te verspillen aan het wachten op anderen.

3. De "Assemblagelijn" (EvoFusion)

Het Probleem: De computer voerde de taak uit in kleine, inefficiënte stappen. Het berekende een getal, stopte, sloeg het op, laadde een nieuwe tool, berekende opnieuw, stopte en sloeg het weer op. Dit is als een chef die een ui snijdt, dan naar de koelkast loopt om een mes te pakken, en dan weer terugloopt om verder te snijden.
De MegaFold Fix: MegaFold combineert deze kleine stappen tot één grote, vloeiende beweging. Het voegt het "snijden", "mengen" en "koken" samen tot één continue actie.
Het Resultaat: De computer verspilt geen tijd meer aan het stoppen en starten van taken. Het draait veel sneller omdat de "assemblagelijn" zonder onderbreking blijft bewegen.

4. De "Voorbereide Ingrediënten" (EvoPipe)

Het Probleem: Voordat de computer zelfs maar kan beginnen met het bouwen van het eiwit, moet een mens (of een trage CPU) veel rommelig onderzoek doen om de juiste data te verzamelen (zoals het zoeken naar de evolutionaire geschiedenis). Dit duurt lang en laat de krachtige computer ongebruikt zitten, wachtend op de data.
De MegaFold Fix: MegaFold realiseert zich dat sommige van dit onderzoek altijd hetzelfde is voor hetzelfde eiwit. Het doet het zware werk één keer vooraf en slaat het op in een "voorraadkast" (cache). Wanneer de computer de data nodig heeft, pakt hij gewoon direct de voorbereide ingrediënten.
Het Resultaat: De krachtige computer hoeft nooit te wachten. Hij is 100% van de tijd bezig.

De Resultaten

Het paper heeft MegaFold getest op twee soorten krachtige computerchips (NVIDIA en AMD). Dit is wat ze vonden:

Langere Eiwitten: MegaFold stelde het team in staat om te trainen op eiwitsequenties die 3,36 keer langer waren dan voorheen.
Snellere Snelheid: Het maakte het trainingsproces 1,73 keer sneller op NVIDIA-chips en 1,62 keer sneller op AMD-chips.
Minder Geheugen: Het gebruikte aanzienlijk minder geheugen, waardoor de "out of memory"-crashes die normaal gesproken bij deze modellen optreden, werden voorkomen.

Samenvatting

Beschouw AlphaFold3 als een superkrachtige motor die te zwaar was voor zijn auto-chassis (het computersysteem). MegaFold heeft niet alleen een grotere auto gebouwd; het heeft de motor, de transmissie en het brandstofsysteem herontworpen, zodat de superkrachtige motor daadwerkelijk efficiënt kan draaien. Dit stelt wetenschappers in staat om deze volgende generatie modellen te trainen op een breder scala aan computers, wat de wetenschap van eiwitvouwing sneller en toegankelijker maakt.

Technische Samenvatting: MegaFold

Probleemstelling

Recente vooruitgang in biomoleculaire modellering, specifiek AlphaFold3 (AF3), heeft een paradigmaverschuiving geïntroduceerd van standaard transformer-architecturen naar modellen die gebruikmaken van 3D-aandacht over 2D-paarvoorstellingen. Hoewel AF3 nauwkeurige voorspellingen mogelijk maakt voor eiwit-eiwit-, eiwit-DNA/RNA- en eiwit-ligandinteracties, introduceert deze architecturale verandering ernstige systeemniveau-bottlenecks die bestaande optimalisaties voor Large Language Models (LLM's) niet kunnen aanpakken.

Het paper identificeert vier fundamentele uitdagingen die geworteld zijn in de kubische schaling van 3D-aandacht:

Kubische Activatiegeheugen: In tegen tegenoverkomende transformers schaalt aandacht kwadratisch ( $O(N^2)$ ), maar de 3D-aandacht van AF3 over 2D-paarvoorstellingen produceert tensoren die kubisch schalen ( $O(N^3)$ ). Dit leidt ertoe dat activatiegeheugen het GPU-gebruik domineert (tot 96% van het totale geheugen), wat veroorzaakt dat Out-of-Memory (OOM) fouten optreden bij gematigde sequentielengtes (bijv. >590 residuen), zelfs met kleine batchgroottes.
Incompatibiliteit met Sequentie-parallellisme (SP): Standaard SP-strategieën voor transformers vertrouwen op 1D-sequentie-sharding. De 2D-paarvoorstellingen van AF3 en de alternerende reductie-assen over verschillende lagen creëren complexe communicatiepatronen over meerdere assen die conventionele SP niet kan ondersteunen zonder excessieve resharding of synchronisatie.
Kernel-fragmentatie: Het afwisselen van 3D-aandacht, transitielagen en projectie-operaties over dimensies met gemengde dimensies resulteert in duizenden fijnmazige kernels per iteratie. Deze fragmentatie verhindert effectieve kernel-fusie, wat leidt tot een slechte warp-bezetting, hoge launch-overhead en inefficiënt gebruik van de geheugenbandbreedte.
Host-zijde Preprocessing Stalls: Het construeren van de evolutionaire context (bijv. Multiple Sequence Alignments, MSA) die vereist is voor 3D-aandacht, omvat zware, afhankelijkheid-gevoelige CPU-preprocessing. Dit creëk een aanzienlijke GPU-idle tijd die niet verborgen kan worden door conventionele prefetching-pipelines.

Methodologie: MegaFold

MegaFold is een cross-platform systeem ontworpen om deze uitdagingen aan te pakken via een co-design van geheugenefficiënte kernels, communicatiestrategieën, operator-fusie en host-device pipelining. Het is geïmplementeerd bovenop een open-source AF3-trainingsframework en ondersteunt zowel NVIDIA- als AMD-GPU's.

1. EvoFlash-3D: Geheugenefficiënte 3D-aandacht

Om de kubische geheugengroei aan te pakken, introduceert MegaFold EvoFlash-3D, een Triton-gebaseerde kernel voor 3D-aandacht.

Mechanisme: In plaats van de volledige $O(N^3)$ aandachtskaart in het globale geheugen (HBM) te materialiseren, gebruikt de kernel een online, numeriek stabiele softmax-strategie. Het berekent incrementeel de aandacht-logits binnen het snelle scratchpad-geheugen (shared memory) tijdens zowel de forward- als de backward-pass.
Optimalisatie: Dit reduceert de piek activatiegeheugencomplexiteit van $O(N^3)$ naar $O(N^2)$ terwijl de numerieke correctheid behouden blijft. De kernel handelt de unieke additieve gebroadcastede paar-bias term van AF3 af, die multidimensionale data-afhankelijkheden introduceert die niet aanwezig zijn in standaard FlashAttention.

2. EvoSP-3D: Communicatie-efficiënte Sharding

Om training op langere sequenties mogelijk te maken, stelt MegaFold de EvoSP-3D strategie voor, die is afgestemd op 2D-paarvoorstellingen.

Mechanisme: Het maakt gebruik van een hybride Data Parallel (DP) en Sequence Parallel (SP) benadering. Terwijl DP-groepen verschillende batch-inputs afhandelen, partitioneren SP-groepen de 2D-paarvoorstelling ( $[N_{token}, N_{token}, h_{pair}]$ ).
Communicatie: Naarmate de AF3-lagen de aandacht-assen afwisselen (bijv. rij-gewijs naar kolom-gewijs), gebruikt EvoSP-3D all_to_all collectieven om de sharding-layout van de paarvoorstelling tussen de stadia te roteren. Dit stelt opeenvolgende aandacht-lagen in staat om op de juiste 2D-slices te opereren zonder de volledige tensor te materialiseren.
Diffusie-fase: In de diffusie-decoder worden uitgebreide atoomcoördinaten verdeeld over SP-groepen zonder communicatie, terwijl de paarvoorstelling wordt verzameld (all_gather) om de diffusie toe te staan op de volledige context te opereren.

3. EvoFusion: Gefuseerde Operator Stack

Om kernel-fragmentatie te beperken, introduceert MegaFold EvoFusion, dat de interne executie van kritieke operaties herstructureert in plaats van enkel de grenzen te fuseren.

Scope: Richt zich op Layer Normalization, Lineaire Projecties en SwiGLU-activaties, die gezamenlijk verantwoordelijk zijn voor ~20% van het activatiegeheugen en de kernel-aanroepen.
Mechanisme: Door de innerlijke loops van deze operaties te herstructureren, elimineert EvoFusion redundante tussenliggende materialisatie (bijv. het vermijden van het schrijven van een volledige $N \times N$ tensor naar HBM tussen normalisatie en lineaire projectie). Dit vermindert de druk op het shared memory en de kernel launch-overhead, terwijl efficiënte reductie- en GeMM-schema's behouden blijven.

4. EvoPipe: Determinisme-bewuste Host-Device Pipeline

Om GPU-idle tijd te elimineren die wordt veroorzaakt door CPU-preprocessing, implementeert MegaFold EvoPipe.

Mechanisme: Het ontkoppelt deterministische preprocessing-stappen (bijv. RDKit-gebaseerde atoomconstructie, MSA-feature extractie) van de stochastische trainingslus. Deze deterministische features worden offline vooraf berekend en gecached.
Executie: Tijdens de training voert het systeem lichte lookups uit om gecachte features op te halen, waarbij alleen niet-deterministische augmentaties (bijv. random cropping) on-the-fly worden toegepast. Dit zorgt voor een consistente GPU-benutting door variabele host-zijde stalls te verwijderen.

Belangrijkste Resultaten

Het systeem werd geëvalueerd op NVIDIA H200 (141GB) en AMD MI250 (64GB) GPU's met behulp van het AlphaFold3-PyTorch framework.

Schaalbaarheid van Sequentielengte: MegaFold maakt training mogelijk met tot wel 3,36× langere sequentielengtes op 32 GPU's vergeleken met de baseline PyTorch-implementaties. Op 32 AMD GPU's ondersteunt het sequenties tot ~1450 residuen, terwijl baselines falen bij ~590 residuen.
Executietijd: MegaFold vermindert de end-to-end per-iteratie executietijd met tot wel 1,73× op NVIDIA en 1,62× op AMD GPU's.
Geheugenefficiëntie: Het piek GPU-geheugengebruik is met tot wel 1,23× verminderd (of een reductie van 23%), waardoor de geheugencapaciteit kan worden gealloceerd aan langere sequenties.
Model FLOPs Utilisatie (MFU): MegaFold bereikt een aanzienlijk hogere MFU (tot 44,9% bij 1,81B parameters) vergeleken met baselines, die vaak de OOM-limieten bereiken voordat ze grotere schalen bereiken.
Cross-Platform Prestaties: Het systeem vertoont consistente winsten over zowel NVIDIA- als AMD-hardware, wat de draagbaarheid van de Triton-gebaseerde kernel-ontwerpen valideert.

Betekenis en Claims

Het paper claimt dat MegaFold het eerste end-to-end systeem is dat geïntegreerde ondersteuning biedt voor het trainen van de volgende generatie 3D-aandacht eiwitmodellen. De betekenis ligt in het aantonen dat:

LLM-optimalisaties ontoereikend zijn: Standaard transformer-optimalisaties (FlashAttention, standaard SP, compiler-fusie) falen bij het aanpakken van de unieke kubische schaling en dimensie-gemengde afhankelijkheden van AF3-stijl modellen.
Systeemniveau Co-design essentieel is: Efficiënte training van deze wetenschappelijke modellen vereist specifieke systeemondersteuning die gezamenlijk kernel-executie, parallellisatie, operator-fusie en host-device pipelining aanpakt.
De Schaalbaarheidsbarrière Verschuift: Door de specifieke bottlenecks van 3D-aandacht aan te pakken, verschuift MegaFold de OOM-barrière, waardoor het trainen van grotere biomoleculaire complexen mogelijk wordt die voorheen onhaalbaar waren, wat de schaling van opkomende wetenschappelijke workloads op moderne GPU-platforms faciliteert.