🧬 biology

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

MegaFold est un nouveau système multiplateforme qui surmonte les goulots d'étranglement de la mémoire et du calcul lors de l'entraînement des modèles de protéines d'attention 3D de nouvelle génération en combinant des noyaux économes en mémoire, un partitionnement optimisé, des opérateurs fusionnés et un pipeline sensible au déterminisme afin d'atteindre des longueurs de séquence nettement plus longues et des temps d'entraînement plus rapides sur les GPU NVIDIA et AMD.

Auteurs originaux : Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Publié 2026-06-16

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ Ceci est une explication générée par l'IA d'un preprint qui n'a pas été évalué par des pairs. Ce n'est pas un avis médical. Ne prenez pas de décisions de santé basées sur ce contenu. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de construire un puzzle 3D massif et complexe d'une protéine. Pendant longtemps, les scientifiques ont utilisé un ensemble de règles spécifiques (comme l'ancien « AlphaFold2 ») qui fonctionnaient bien. Mais la version la plus récente, AlphaFold3, est comme passer d'un puzzle 2D plat à une structure 3D complexe et multicouche où chaque pièce interagit avec toutes les autres dans un espace 3D.

Bien que ce nouveau modèle soit scientifiquement incroyable, il présente un énorme problème : il dévore la mémoire informatique comme un trou noir et fonctionne incroyablement lentement.

Ce document présente MegaFold, un nouveau « système » (un ensemble d'outils et de règles spécialisés) conçu pour rendre l'entraînement de ces nouveaux modèles de protéines 3D possible sur des ordinateurs modernes sans qu'ils ne plantent ou ne mettent une éternité.

Voici comment MegaFold résout les quatre problèmes principaux, expliqués avec des analogies simples :

Le Problème : Pourquoi l'ancienne méthode échoue

Le nouveau modèle AlphaFold3 observe les protéines d'une manière qui crée une explosion « cubique » de données.

L'analogie : Imaginez que vous organisiez une fête. Dans l'ancien modèle, vous aviez juste besoin de savoir qui était ami avec qui (une simple liste). Dans le nouveau modèle, vous devez suivre une carte 3D de la façon dont chaque invité interagit avec chaque autre invité simultanément. Si vous avez 100 invités, l'ancienne méthode nécessite 10 000 notes ; la nouvelle nécessite 1 000 000 de notes. Si vous essayez d'écrire toutes ces notes sur un seul bloc-notes (la mémoire de l'ordinateur), le bloc-notes explose.

La Solution : Les quatre outils de MegaFold

MegaFold corrige cela en introduisant quatre mises à niveau spécifiques :

1. Le « Bloc-notes Intelligent » (EvoFlash-3D)

Le Problème : L'ordinateur essaie d'écrire toute la carte d'interaction 3D d'un coup, remplissant instantanément sa mémoire.
La correction de MegaFold : Au lieu d'écrire toute la carte sur le bloc-notes, MegaFold utilise un « brouillon » (une mémoire temporaire rapide) pour calculer les interactions par petits morceaux, un par un, puis jette les notes immédiatement après.
Le Résultat : Il n'a jamais besoin de conserver toute la carte massive en mémoire à la fois. Cela permet à l'ordinateur de gérer des chaînes de protéines beaucoup plus longues sans manquer d'espace.

2. Le « Relais d'Équipe » (EvoSP-3D)

Le Problème : Lorsque vous utilisez plusieurs ordinateurs (GPU) pour travailler ensemble, ils répartissent généralement le travail en donnant à chaque ordinateur une liste d'invités différente. Mais comme les données d'AlphaFold3 sont une grille 2D d'interactions, la simple division ne fonctionne pas ; les ordinateurs se confondent sur qui parle à qui.
La Correction de MegaFold : MegaFold invente une nouvelle façon pour les ordinateurs de transmettre les données. Au lieu de simplement diviser la liste, ils divisent la grille d'interactions. Ils se passent les morceaux du puzzle d'avant en arrière selon un schéma spécifique et efficace (comme une course de relais) afin que chaque ordinateur sache exactement quelle partie de la carte 3D est sous sa responsabilité, sans perdre de temps à attendre les autres.

3. La « Chaîne de Montage » (EvoFusion)

Le Problème : L'ordinateur effectuait le travail en étapes minuscules et inefficaces. Il calculait un nombre, s'arrêtait, sauvegardait, chargeait un nouvel outil, calculait à nouveau, s'arrêtait, puis sauvegardait à nouveau. C'est comme un chef qui coupe un oignon, puis marche jusqu'au frigo pour chercher un couteau, puis revient pour couper à nouveau.
La Correction de MegaFold : MegaFold combine ces étapes minuscules en un seul mouvement géant et fluide. Il fusionne le « hachage », le « mélange » et la « cuisson » en une seule action continue.
Le Résultat : L'ordinateur arrête de perdre du temps à démarrer et arrêter des tâches. Il fonctionne beaucoup plus vite car il maintient la « chaîne de montage » en mouvement sans interruption.

4. Les « Ingrédients Préparés à l'Avance » (EvoPipe)

Le Problème : Avant même que l'ordinateur puisse commencer à construire la protéine, un humain (ou un processeur CPU lent) doit effectuer beaucoup de recherches fastidieuses pour rassembler les bonnes données (comme trouver l'histoire évolutive). Cela prend beaucoup de temps et laisse l'ordinateur puissant inactif, attendant les données.
La Correction de MegaFold : MegaFold réalise que certaines de ces recherches sont toujours les mêmes pour une même protéine. Il effectue le travail difficile une seule fois à l'avance et le stocke dans un « garde-manger » (un cache). Quand l'ordinateur a besoin des données, il saisit simplement les ingrédients pré-préparés instantanément.
Le Résultat : L'ordinateur puissant n'a jamais besoin d'attendre. Il reste occupé 100 % du temps.

Les Résultats

L'article a testé MegaFold sur deux types de puces informatiques puissantes (NVIDIA et AMD). Voici ce qu'ils ont trouvé :

Protéines plus longues : MegaFold a permis à l'équipe de s'entraîner sur des séquences de protéines 3,36 fois plus longues qu'auparavant.
Vitesse accrue : Cela a rendu le processus d'entraînement 1,73 fois plus rapide sur les puces NVIDIA et 1,62 fois plus rapide sur les puces AMD.
Moins de mémoire : Il a utilisé nettement moins de mémoire, évitant les plantages de type « out of memory » (mémoire insuffisante) qui surviennent habituellement avec ces modèles.

Résumé

Considérez AlphaFold3 comme un moteur super puissant qui était trop lourd pour son châssis de voiture (le système informatique). MegaFold n'a pas seulement construit une voiture plus grande ; il a redessiné le moteur, la transmission et le système d'alimentation pour que ce moteur super puissant puisse réellement fonctionner efficacement. Cela permet aux scientifiques d'entraîner ces modèles de nouvelle génération sur une plus grande variété d'ordinateurs, rendant la science du repliement des protéines plus rapide et plus accessible.

Résumé Technique : MegaFold

Énoncé du Problème

Les avancées récentes dans la modélisation biomoléculaire, spécifiquement AlphaFold3 (AF3), ont introduit un changement de paradigme par rapport aux architectures transformer standards, en utilisant des modèles basés sur l'attention 3D sur des représentations paires 2D. Bien qu'AF3 permette une prédiction précise des interactions protéine-protéine, protéine-ADN/ARN et protéine-ligand, ce changement architectural introduit de graves goulots d'étranglement au niveau du système que les optimisations actuelles des Grands Modèles de Langage (LLM) ne peuvent résoudre.

Le document identifie quatre défis fondamentaux ancrés dans la mise à l'échelle cubique de l'attention 3D :

Mémoire d'Activation Cubique : Contrairement aux transformers où l'attention suit une échelle quadratique ( $O(N^2)$ ), l'attention 3D d'AF3 sur des représentations paires 2D produit des tenseurs dont l'échelle est cubique ( $O(N^3)$ ). Cela entraîne une dominance de la mémoire d'activation dans l'utilisation des GPU (jusqu'à 96 % de la mémoire totale), provoquant des erreurs de type Out-of-Memory (OOM) pour des longueurs de séquences modérées (par exemple, >590 résidus), même avec de petites tailles de lots (batch sizes).
Incompatibilité avec le Parallélisme de Séquence (SP) : Les stratégies de SP standard pour les transformers reposent sur un partitionnement de séquence 1D. Les représentations paires 2D d'AF3 et les axes de réduction alternés à travers les couches créent des modèles de communication multi-axes complexes que le SP conventionnel ne peut supporter sans un resharding ou une synchronisation excessive.
Fragmentation de Noyaux (Kernel Fragmentation) : L'entrelacement de l'attention 3D, des couches de transition et des opérations de projection sur des tenseurs de dimensions mixtes résulte en des milliers de noyaux (kernels) fins par itération. Cette fragmentation empêche une fusion de noyaux efficace, entraînant une faible occupation des warps, un overhead de lancement élevé et une utilisation inefficace de la bande passante mémoire.
Blocages de Prétraitement Côté Hôte : La construction du contexte évolutif (par exemple, les alignements de séquences multiples, MSA) requis pour l'attention 3D implique un prétraitement CPU lourd et dépendant de nombreuses variables. Cela crée un temps d'inactivité significatif du GPU qui ne peut être masqué par les pipelines de préchargement (prefetching) conventionnels.

Méthodologie : MegaFold

MegaFold est un système cross-plateforme conçu pour répondre à ces défis via un co-design de noyaux économes en mémoire, de stratégies de communication, de fusion d'opérateurs et de pipelining hôte-périphérique. Il est implémenté sur la base d'un framework d'entraînement AF3 open-source et supporte les GPU NVIDIA et AMD.

1. EvoFlash-3D : Attention 3D Économe en Mémoire

Pour traiter la croissance cubique de la mémoire, MegaFold introduit EvoFlash-3D, un noyau basé sur Triton pour l'attention 3D.

Mécanisme : Au lieu de matérialiser la carte d'attention complète $O(N^3)$ dans la mémoire globale (HBM), le noyau utilise une stratégie de softmax en ligne et numériquement stable. Il calcule de manière incrémentale les logits d'attention dans une mémoire scratchpad rapide (mémoire partagée) lors des passes avant et arrière.
Optimisation : Cela réduit la complexité de la mémoire d'activation de pointe de $O(N^3)$ à $O(N^2)$ tout en préservant la justesse numérique. Le noyau gère le terme de biais de paire diffusé additivement unique d'AF3, qui introduit des dépendances de données multidimensionnelles absentes de la FlashAttention standard.

2. EvoSP-3D : Partitionnement (Sharding) Économe en Communication

Pour permettre l'entraînement sur des séquences plus longues, MegaFold propose EvoSP-3D, une stratégie de partitionnement adaptée aux représentations paires 2D.

Mécanisme : Il utilise une approche hybride de Parallélisme de Données (DP) et de Séquence (SP). Tandis que les groupes DP gèrent différentes entrées de lots, les groupes SP partitionnent le tenseur pair 2D ( $[N_{token}, N_{token}, h_{pair}]$ ).
Communication : Comme les couches d'AF3 alternent les axes d'attention (par exemple, de la ligne vers la colonne), EvoSP-3D utilise des collectifs all_to_all pour faire pivoter la disposition du partitionnement du tenseur pair entre les étapes. Cela permet aux couches d'attention successives d'opérer sur les tranches 2D correctes sans matérialiser le tenseur complet.
Étape de Diffusion : Dans le décodeur de diffusion, les coordonnées atomiques augmentées sont distribuées à travers les groupes SP sans communication, tandis que la représentation paire est rassemblée (all_gather) pour permettre à la diffusion d'opérer sur le contexte complet.

3. EvoFusion : Pile d'Opérateurs Fusionnés

Pour atténuer la fragmentation des noyaux, MegaFold introduit EvoFusion, qui restructure l'exécution interne des opérateurs critiques plutôt que de simplement fusionner les frontières.

Portée : Cible la Normalisation de Couche (Layer Normalization), les Projections Linéaires et les activations SwiGLU, qui représentent collectivement environ 20 % de la mémoire d'activation et des appels de noyaux.
Mécanisme : En restructurant les boucles internes de ces opérateurs, EvoFusion élimine la matérialisation intermédiaire redondante (par exemple, éviter d'écrire un tenseur complet $N \times N$ dans l'HBM entre la normalisation et la projection linéaire). Cela réduit la pression sur la mémoire partagée et l'overhead de lancement de noyau tout en maintenant des schémas de réduction et de GeMM efficaces.

4. EvoPipe : Pipeline Hôte-Périphérique Déterministe

Pour éliminer le temps d'inactivité du GPU causé par le prétraitement CPU, MegaFold implémente EvoPipe.

Mécanisme : Il découple les étapes de prétraitement déterministes (par exemple, la construction d'atomes via RDKit, l'extraction de caractéristiques MSA) des boucles d'entraînement stochastiques. Ces caractéristiques déterministes sont précalculées et mises en cache hors ligne.
Exécution : Pendant l'entraînement, le système effectue des recherches légères pour récupérer les caractéristiques mises en cache, n'appliquant que les augmentations non déterministes (par exemple, le recadrage aléatoire) à la volée. Cela garantit une utilisation constante du GPU en supprimant les blocages variables côté hôte.

Résultats Clés

Le système a été évalué sur des GPU NVIDIA H200 (141 Go) et AMD MI250 (64 Go) en utilisant le framework AlphaFold3-PyTorch.

Scalabilité de la Longueur de Séquence : MegaFold permet l'entraînement avec des longueurs de séquences jusqu'à 3,36× plus longues sur 32 GPU par rapport aux implémentations PyTorch de base. Sur 32 GPU AMD, il supporte des séquences allant jusqu'à environ 1450 résidus, là où les versions de base échouent à environ 590 résidus.
Temps d'Exécution : MegaFold réduit le temps d'exécution par itération de pointe jusqu'à 1,73× sur NVIDIA et 1,62× sur AMD.
Efficacité Mémoire : L'utilisation de la mémoire GPU de pointe est réduite jusqu'à 1,23× (soit une réduction de 23 %), permettant de réallouer la capacité de mémoire à des séquences plus longues.
Utilisation des FLOPs du Modèle (MFU) : MegaFold atteint un MFU nettement plus élevé (jusqu'à 44,9 % à 1,81 milliard de paramètres) par rapport aux bases, qui atteignent souvent des limites OOM avant d'atteindre de plus grandes échelles.
Performance Cross-Platform : Le système démontre des gains constants sur les matériels NVIDIA et AMD, validant la portabilité des conceptions de noyaux basées sur Triton.

Signification et Revendications

L'article affirme que MegaFold est le premier système de bout en bout à fournir un support intégré pour l'entraînement des modèles de protéines à attention 3D de nouvelle génération. Sa signification réside dans la démonstration que :

Les Optimisations LLM sont Insuffisantes : Les optimisations de transformer standards (FlashAttention, SP standard, fusion de compilateur) échouent à traiter l'échelle cubique unique et les dépendances de dimensions mixtes des modèles de type AF3.
Le Co-design au Niveau Système est Essentiel : L'entraînement efficace de ces modèles scientifiques nécessite un support système dédié qui traite conjointement l'exécution des noyaux, la parallélisation, la fusion d'opérateurs et le pipelining hôte-périphérique.
Déplacement de la Barrière de Scalabilité : En adressant les goulots d'étranglement spécifiques de l'attention 3D, MegaFold déplace la barrière OOM, permettant l'entraînement de complexes biomoléculaires plus grands qui étaient auparavant irréalisables, facilitant ainsi le passage à l'échelle des charges de travail scientifiques émergentes sur les plateformes GPU modernes.