🧬 biology

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

MegaFold es un sistema multiplataforma novedoso que supera los cuellos de botella de memoria y computación del entrenamiento de modelos de proteínas de atención 3D de próxima generación al combinar núcleos eficientes en memoria, fragmentación optimizada, operadores fusionados y un flujo de trabajo consciente de la determinación para lograr longitudes de secuencia significativamente más largas y tiempos de entrenamiento más rápidos tanto en GPUs NVIDIA como AMD.

Autores originales: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Publicado 2026-06-16

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ Esta es una explicación generada por IA de un preprint que no ha sido revisado por pares. No es consejo médico. No tome decisiones de salud basándose en este contenido. Leer descargo de responsabilidad completo

Imagina que estás intentando construir un rompecabezas 3D masivo e intrincado de una proteína. Durante mucho tiempo, los científicos utilizaron un conjunto específico de reglas (como el antiguo "AlphaFold2") que funcionaban bien. Pero la versión más reciente, AlphaFold3, es como actualizar de un rompecabezas 2D plano a una estructura 3D compleja y multicapa donde cada una de las piezas interactúa con todas las demás en un espacio 3D.

Si bien este nuevo modelo es científicamente asombroso, tiene un gran problema: consume la memoria de la computadora como un agujero negro y funciona increíblemente lento.

Este artículo presenta MegaFold, un nuevo "sistema" (un conjunto de herramientas y reglas especializadas) diseñado para hacer posible el entrenamiento de estos nuevos modelos de proteínas 3D en computadoras modernas sin que se bloqueen o tarden una eternidad.

Aquí te explicamos cómo MegaFold resuelve los cuatro problemas principales, utilizando analogías sencillas:

El Problema: Por qué el método antiguo falla

El nuevo modelo AlphaFold3 observa las proteínas de una manera que crea una explosión de datos "cúbica".

La Analogía: Imagina que estás organizando una fiesta. En el modelo antiguo, solo necesitabas saber quién era amigo de quién (una lista simple). En el nuevo modelo, tienes que rastrear un mapa 3D de cómo cada invitado está interactuando con cada otro invitado simultáneamente. Si tienes 100 invitados, el método antiguo necesita 10,000 notas; el nuevo método necesita 1,000,000 de notas. Si intentas escribir todas esas notas en un solo bloc de notas (la memoria de la computadora), el bloc de notas explota.

La Solución: Las cuatro herramientas de MegaFold

MegaFold soluciona esto introduciendo cuatro mejoras específicas:

1. El "Bloc de Notas Inteligente" (EvoFlash-3D)

El Problema: La computadora intenta escribir todo el mapa de interacción 3D a la vez, llenando su memoria instantáneamente.
La Solución de MegaFold: En lugar de escribir todo el mapa en el bloc de notas, MegaFold utiliza un "borrador" (memoria rápida y temporal) para calcular las interacciones en pequeños fragmentos, uno por uno, y luego desecha las notas inmediatamente.
El Resultado: Nunca necesita mantener todo el mapa masivo en la memoria al mismo tiempo. Esto permite a la computadora manejar cadenas de proteínas mucho más largas sin quedarse sin espacio.

2. El "Relevo en Equipo" (EvoSP-3D)

El Problema: Cuando utilizas múltiples computadoras (GPUs) para trabajar en conjunto, normalmente dividen el trabajo dándole a cada computadora una lista de invitados diferente. Pero debido a que los datos de AlphaFold3 son una cuadrícula 2D de interacciones, dividirlo de esa manera simplemente no funciona; las computadoras se confunden sobre quién está hablando con quién.
La Solución de MegaFold: MegaFold inventa una nueva forma para que las computadoras pasen datos. En lugar de solo dividir la lista, dividen la cuadrícula de interacciones. Pasan piezas del rompecabezas de un lado a otro en un patrón específico y eficiente (como una carrera de relevos) para que cada computadora sepa exactamente qué parte del mapa 3D es su responsabilidad, sin perder tiempo esperando a las demás.

3. La "Línea de Ensamblaje" (EvoFusion)

El Problema: La computadora estaba realizando el trabajo en pasos diminutos e ineficientes. Calculaba un número, se detenía, lo guardaba, cargaba una nueva herramienta, calculaba de nuevo, se detenía y guardaba de nuevo. Esto es como un chef picando una cebolla, luego caminando hacia la nevera para buscar un cuchillo, y luego volviendo para seguir picando.
La Solución de MegaFold: MegaFold combina estos pasos diminutos en un único movimiento gigante y fluido. Fusiona el "picar", el "mezclar" y el "cocinar" en una sola acción continua.
El Resultado: La computadora deja de perder tiempo iniciando y deteniendo tareas. Funciona mucho más rápido porque mantiene la "línea de ensamblaje" en movimiento sin interrupciones.

4. Los "Ingredientes Pre-Preparados" (EvoPipe)

El Problema: Antes de que la computadora pueda siquiera empezar a construir la proteína, un humano (o una CPU lenta) tiene que realizar mucha investigación desordenada para reunir los datos correctos (como encontrar la historia evolutiva). Esto toma mucho tiempo y deja a la poderosa computadora sentada ociosa, esperando los datos.
La Solución de MegaFold: MegaFold se da cuenta de que parte de esta investigación siempre es la misma para la misma proteína. Realiza el trabajo difícil una sola vez con antelación y lo almacena en una "despensa" (caché). Cuando la computadora necesita los datos, simplemente toma los ingredientes pre-preparados al instante.
El Resultado: La poderosa computadora nunca tiene que esperar. Se mantiene ocupada el 100% del tiempo.

Los Resultados

El artículo probó MegaFold en dos tipos de chips de computadora potentes (NVIDIA y AMD). Esto fue lo que encontraron:

Proteínas más largas: MegaFold permitió al equipo entrenar con secuencias de proteínas 3.36 veces más largas que antes.
Velocidad mayor: Hizo que el proceso de entrenamiento fuera 1.73 veces más rápido en chips NVIDIA y 1.62 veces más rápido en chips AMD.
Menos Memoria: Utilizó significativamente menos memoria, evitando los errores de "fuera de memoria" (out of memory) que suelen ocurrir con estos modelos.

Resumen

Piensa en AlphaFold3 como un motor superpotente que era demasiado pesado para su chasis de automóvil (el sistema de la computadora). MegaFold no solo construyó un auto más grande; rediseñó el motor, la transmisión y el sistema de combustible para que el motor superpotente pudiera funcionar de manera eficiente. Esto permite a los científicos entrenar estos modelos de próxima generación en una variedad más amplia de computadoras, haciendo que la ciencia del plegamiento de proteínas sea más rápida y accesible.

Resumen Técnico: MegaFold

Planteamiento del Problema

Los avances recientes en el modelado biomolecular, específicamente AlphaFold3 (AF3), han introducido un cambio de paradigma de las arquitecturas transformer estándar hacia modelos que utilizan atención 3D sobre representaciones pareadas 2D. Si bien AF3 permite la predicción precisa de interacciones proteína-proteína, proteína-ADN/ARN y proteína-ligando, este cambio arquitectónico introduce severos cuellos de botella a nivel de sistema que las optimizaciones existentes de los Modelos de Lenguaje de Gran Escala (LLM) no pueden abordar.

El artículo identifica cuatro desafíos fundamentales derivados del escalamiento cúbico de la atención 3D:

Memoria de Activación Cúbica: A diferencia de los transformers donde la atención escala cuadráticamente ( $O(N^2)$ ), la atención 3D de AF3 sobre representaciones pareadas 2D produce tensores que escalan cúbicamente ( $O(N^3)$ ). Esto provoca que la memoria de activación domine el uso de la GPU (hasta el 96% de la memoria total), causando errores de falta de memoria (OOM) con longitudes de secuencia moderadas (por ejemplo, >590 residuos) incluso con tamaños de lote pequeños.
Incompatibilidad con el Paralelismo de Secuencia (SP): Las estrategias estándar de SP para transformers dependen del particionamiento de secuencias 1D. Las representaciones pareadas 2D de AF3 y los ejes de reducción alternos a través de las capas crean patrones de comunicación multiaxiales complejos que el SP convencional no puede soportar sin un exceso de re-particionamiento o sincronización.
Fragmentación de Kernels: El entrelazado de la atención 3D, las capas de transición y las operaciones de proyección sobre tensores de dimensiones mixtas resulta en miles de kernels de grano fino por iteración. Esta fragmentación impide una fusión de kernels efectiva, lo que conduce a una baja ocupación de warps, alta sobrecarga de lanzamiento y un uso ineficiente del ancho de banda de la memoria.
Retrasos en el Preprocesamiento del Host: La construcción del contexto evolutivo (por ejemplo, Alineamientos de Múltiples Secuencias, MSA) requerido para la atención 3D implica un preprocesamiento pesado en la CPU con muchas dependencias. Esto crea un tiempo de inactividad significativo de la GPU que no puede ser ocultado por los canales de prelectura (prefetching) convencionales.

Metodología: MegaFold

MegaFold es un sistema multiplataforma diseñado para abordar estos desafíos mediante un codiseño de kernels eficientes en memoria, estrategias de comunicación, fusión de operadores y pipelining host-dispositivo. Está implementado sobre un marco de entrenamiento de AF3 de código abierto y es compatible tanto con GPUs NVIDIA como AMD.

1. EvoFlash-3D: Atención 3D Eficiente en Memoria

Para abordar el crecimiento cúbico de la memoria, MegaFold introduce EvoFlash-3D, un kernel basado en Triton para la atención 3D.

Mecanismo: En lugar de materializar el mapa de atención completo de $O(N^3)$ en la memoria global (HBM), el kernel utiliza una estrategia de softmax numéricamente estable y en línea. Computa incrementalmente los logits de atención dentro de la memoria rápida de scratchpad (memoria compartida) durante las pasadas hacia adelante y hacia atrás.
Optimización: Esto reduce la complejidad de la memoria de activación pico de $O(N^3)$ a $O(N^2)$ preservando la corrección numérica. El kernel gestiona el término de sesgo de par transmitido aditivo único de AF3, el cual introduce dependencias de datos multidimensionales que no están presentes en la FlashAttention estándar.

2. EvoSP-3D: Particionamiento Eficiente en Comunicación

Para permitir el entrenamiento en secuencias más largas, MegaFold propone EvoSP-3D, una estrategia de particionamiento adaptada para representaciones pareadas 2D.

Mecanismo: Utiliza un enfoque híbrido de Paralelismo de Datos (DP) y Paralelismo de Secuencia (SP). Mientras que los grupos DP manejan diferentes entradas de lote, los grupos SP particionan el tensor pareado 2D ( $[N_{token}, N_{token}, h_{pair}]$ ).
Comunicación: Dado que las capas de AF3 alternan los ejes de atención (por ejemplo, de fila a columna), EvoSP-3D utiliza colectivos all_to_all para rotar el diseño de particionamiento del tensor pareado entre etapas. Esto permite que las sucesivas capas de atención operen sobre las rebanadas 2D correctas sin materializar el tensor completo.
Etapa de Difusión: En el decodificador de difusión, las coordenadas de átomos aumentadas se distribuyen a través de los grupos SP sin comunicación, mientras que la representación pareada se recolecta (all_gather) para permitir que la difusión opere sobre el contexto completo.

3. EvoFusion: Pila de Operadores Fusionados

Para mitigar la fragmentación de kernels, MegaFold introduce EvoFusion, que reestructura la ejecución interna de los operadores críticos en lugar de simplemente fusionar los límites.

Alcance: Se dirige a la Normalización de Capas (Layer Normalization), Proyecciones Lineales y activaciones SwiGLU, que colectivamente representan aproximadamente el 20% de la memoria de activación y las invocaciones de kernel.
Mecanismo: Al reestructurar los bucles internos de estos operadores, EvoFusion elimina la materialización intermedia redundante (evitando escribir un tensor completo de $N \times N$ en la HBM entre la normalización y la proyección lineal). Esto reduce la presión en la memoria compartida y la sobrecarga de lanzamiento de kernels, manteniendo un esquema eficiente de reducción y GeMM.

4. EvoPipe: Pipeline Host-Device con Conciencia de Determinismo

Para eliminar el tiempo de inactividad de la GPU causado por el preprocesamiento en la CPU, MegaFold implementa EvoPipe.

Mecanismo: Desacopla los pasos de preprocesamiento deterministas (por ejemplo, la construcción de átomos basada en RDKit, la extracción de características de MSA) del bucle de entrenamiento estocástico. Estas características deterministas se precomputan y se almacenan en caché fuera de línea.
Ejecución: Durante el entrenamiento, el sistema realiza búsquedas ligeras para recuperar las características almacenadas en caché, aplicando únicamente las aumentaciones no deterministas (por ejemplo, recorte aleatorio) sobre la marcha. Esto asegura una utilización constante de la GPU al eliminar los retrasos variables del host.

Resultados Clave

El sistema fue evaluado en GPUs NVIDIA H200 (141GB) y AMD MI250 (64GB) utilizando el marco de trabajo AlphaFold3-PyTorch.

Escalabilidad de la Longitud de Secuencia: MegaFold permite el entrenamiento con longitudes de secuencia hasta 3.36× más largas en 32 GPUs en comparación con las implementaciones base de PyTorch. En 32 GPUs AMD, soporta secuencias de hasta ~1450 residuos, mientras que las bases fallan a los ~590 residuos.
Tiempo de Ejecución: MegaFold reduce el tiempo de ejecución por iteración de extremo a extremo hasta en 1.73× en NVIDIA y 1.62× en AMD.
Eficiencia de Memoria: El uso de memoria pico de la GPU se reduce hasta en 1.23× (o una reducción del 23%), permitiendo la reasignación de la capacidad de memoria a secuencias más largas.
Utilización de FLOPs del Modelo (MFU): MegaFold logra un MFU significativamente mayor (hasta 44.9% con 1.81B de parámetros) en comparación con las bases, que a menudo alcanzan los límites de OOM antes de llegar a escalas mayores.
Rendimiento Multiplataforma: El sistema demuestra ganancias consistentes en hardware NVIDIA y AMD, validando la portabilidad de los diseños de kernel basados en Triton.

Significado y Reivindicaciones

El artículo afirma que MegaFold es el primer sistema de extremo a extremo que proporciona soporte integrado para el entrenamiento de modelos de proteínas con atención 3D de próxima generación. Su importancia radica en demostrar que:

Las optimizaciones de LLM son insuficientes: Las optimizaciones estándar de transformer (FlashAttention, SP estándar, fusión de compilador) fallan al abordar el escalamiento cúbico único y las dependencias de dimensiones mixtas de los modelos tipo AF3.
El codiseño a nivel de sistema es esencial: El entrenamiento eficiente de estos modelos científicos requiere un soporte de sistema dedicado que aborde conjuntamente la ejecución de kernels, el paralelismo, la fusión de operadores y el pipelining host-dispositivo.
Cambio en la Barrera de Escalabilidad: Al abordar los cuellos de botella específicos de la atención 3D, MegaFold desplaza la barrera de OOM, permitiendo el entrenamiento de complejos biomoleculares más grandes que antes eran inviables, facilitando así el escalamiento de las cargas de trabajo científicas emergentes en plataformas de GPU modernas.