🧬 biology

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

O MegaFold é um sistema cross-platform inovador que supera os gargalos de memória e computação do treinamento de modelos de atenção de proteínas 3D de próxima geração ao combinar kernels eficientes em memória, fragmentação otimizada, operadores fundidos e um pipeline consciente de determinismo para alcançar comprimentos de sequência significativamente maiores e tempos de treinamento mais rápidos tanto em GPUs NVIDIA quanto AMD.

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

Publicado 2026-06-16

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

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

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ Esta é uma explicação gerada por IA de um preprint que não foi revisado por pares. Não é aconselhamento médico. Não tome decisões de saúde com base neste conteúdo. Ler aviso legal completo

Imagine que você está tentando construir um quebra-cabeça 3D massivo e intrincado de uma proteína. Por muito tempo, os cientistas usaram um conjunto específico de regras (como o antigo "AlphaFold2") que funcionava bem. Mas a versão mais recente, o AlphaFold3, é como atualizar de um quebra-cabeça 2D plano para uma estrutura 3D complexa e multicamadas, onde cada peça interage com todas as outras peças em um espaço 3D.

Embora este novo modelo seja cientificamente incrível, ele tem um grande problema: ele consome memória de computador como um buraco negro e roda incrivelmente devagar.

Este artigo apresenta o MegaFold, um novo "sistema" (um conjunto de ferramentas e regras especializadas) projetado para tornar possível o treinamento desses novos modelos de proteínas 3D em computadores modernos sem que eles travem ou demorem uma eternidade.

Aqui está como o MegaFold resolve os quatro problemas principais, explicados com analogias simples:

O Problema: Por que o modo antigo falha

O novo modelo AlphaFold3 observa as proteínas de uma forma que cria uma explosão "cúbica" de dados.

A Analogia: Imagine que você está organizando uma festa. No modelo antigo, você só precisava saber quem era amigo de quem (uma lista simples). No novo modelo, você tem que rastrear um mapa 3D de como cada convidado está interagindo com todos os outros convidados simultaneamente. Se você tiver 100 convidados, o modo antigo precisa de 10.000 notas; o novo modo precisa de 1.000.000 de notas. Se você tentar escrever todas essas notas em um único bloco de notas (a memória do computador), o bloco de notas explode.

A Solução: As Quatro Ferramentas do MegaFold

O MegaFold corrige isso introduzindo quatro atualizações específicas:

1. O "Bloco de Notas Inteligente" (EvoFlash-3D)

O Problema: O computador tenta escrever todo o mapa de interação 3D de uma vez, preenchendo sua memória instantaneamente.
A Correção do MegaFold: Em vez de escrever todo o mapa no bloco de notas, o MegaFold usa um "rascunho" (memória rápida e temporária) para calcular as interações em pequenos blocos, um por um, e depois descarta as notas imediatamente.
O Resultado: Ele nunca precisa manter todo o mapa massivo na memória ao mesmo tempo. Isso permite que o computador lide com cadeias de proteínas muito mais longas sem ficar sem espaço.

2. O "Revezamento de Equipe" (EvoSP-3D)

O Probleão: Quando você usa vários computadores (GPUs) para trabalharem juntos, eles geralmente dividem o trabalho dando a cada computador uma lista de convidados diferente. Mas como os dados do AlphaFold3 são uma grade 2D de interações, dividir isso simplesmente não funciona; os computadores ficam confusos sobre quem está falando com quem.
A Correção do MegaFold: O MegaFold inventa uma nova maneira para os computadores passarem dados. Em vez de apenas dividir a lista, eles dividem a grade de interações. Eles passam pedaços do quebra-cabeça de um para o outro em um padrão específico e eficiente (como uma corrida de revezamento), para que cada computador saiba exatamente qual parte do mapa 3D é de sua responsabilidade, sem perder tempo esperando pelos outros.

3. A "Linha de Montagem" (EvoFusion)

O Problema: O computador estava fazendo o trabalho em etapas minúsculas e ineficientes. Ele calculava um número, parava, salvava, carregava uma nova ferramenta, calculava novamente, parava e salvava de novo. Isso é como um chef picando uma cebola, depois caminhando até a geladeira para pegar uma faca, e depois voltando para picar novamente.
A Correção do MegaFold: O MegaFold combina essas etapas minúsculas em um único movimento grande e suave. Ele funde o "picar", o "misturar" e o "cozinhar" em uma única ação contín로.
O Resultado: O computador para de perder tempo começando e parando tarefas. Ele roda muito mais rápido porque mantém a "linha de montagem" em movimento, sem interrupções.

4. Os "Ingredientes Pré-Preparados" (EvoPipe)

O Problema: Antes que o computador possa sequer começar a construir a proteína, um humano (ou uma CPU lenta) tem que fazer muita pesquisa bagunçada para reunir os dados corretos (como encontrar o histórico evolutivo). Isso leva muito tempo e deixa o computador potente parado, esperando pelos dados.
A Correção do MegaFold: O MegaFold percebe que parte dessa pesquisa é sempre a mesma para a mesma proteína. Ele faz o trabalho duro uma vez antecipadamente e o armazena em uma "despensa" (cache). Quando o computador precisa dos dados, ele apenas pega os ingredientes pré-preparados instantaneamente.
O Resultado: O computador poderoso nunca precisa esperar. Ele permanece ocupado 100% do tempo.

Os Resultados

O artigo testou o MegaFold em dois tipos de chips de computador poderosos (NVIDIA e AMD). Aqui está o que eles descobriram:

Proteínas Mais Longas: O MegaFold permitiu que a equipe treinasse em sequências de proteínas 3,36 vezes mais longas do que antes.
Velocidade Maior: Tornou o processo de treinamento 1,73 vezes mais rápido em chips NVIDIA e 1,62 vezes mais rápido em chips AMD.
Menos Memória: Usou significativamente menos memória, evitando os travamentos de "falta de memória" (out of memory) que costumam acontecer com esses modelos.

Resumo

Pense no AlphaFold3 como um motor superpotente que era pesado demais para o chassi do carro (o sistema do computador). O MegaFold não apenas construiu um carro maior; ele redesenhou o motor, a transmissão e o sistema de combustível para que o motor superpotente pudesse realmente rodar de forma eficiente. Isso permite que os cientistas treinem esses modelos de próxima geração em uma variedade maior de computadores, tornando a ciência do dobramento de proteínas mais rápida e acessível.

Resumo Técnico: MegaFold

Declaração do Problema

Avanços recentes na modelagem biomolecular, especificamente o AlphaFold3 (AF3), introduziram uma mudança de paradigma dos modelos transformer padrão para modelos que utilizam atenção 3D sobre representações pareadas 2D. Embora o AF3 permita a previsão precisa de interações proteína-proteína, proteína-DNA/RNA e proteína-ligante, essa mudança arquitetônica introduz gargalos severos em nível de sistema que as otimizações existentes de Large Language Models (LLMs) não conseguem abordar.

O artigo identifica quatro desafios fundamentais enraizados no escalonamento cúbico da atenção 3D:

Memória de Ativação Cúbica: Diferente dos transformers onde a atenção escala quadraticamente ( $O(N^2)$ ), a atenção 3D do AF3 sobre representações pareadas 2D produz tensores que escalam cubicamente ( $O(N^3)$ ). Isso faz com que a memória de ativação domine o uso da GPU (até 96% da memória total), causando erros de Out-of-Memory (OOM) em comprimentos de sequência moderados (ex: >590 resíduos), mesmo com tamanhos de lote pequenos.
Incompatibilidade com Paralelismo de Sequência (SP): Estratégias padrão de SP para transformers dependem de fragmentação (sharding) de sequência 1D. As representações pareadas 2D do AF3 e os eixos de redução alternados entre camadas criam padrões de comunicação multi-eixo complexos que o SP convencional não consegue suportar sem excessiva refragmentação ou sincronização.
Fragmentação de Kernel: O entrelaçamento de atenção 3D, camadas de transição e operações de projeção sobre tensores de dimensões mistas resulta em milhares de kernels granulares por iteração. Essa fragmentação impede a fusão de kernels eficaz, levando a uma baixa ocupação de warps, alto overhead de lançamento e utilização ineficiente da largura de banda de memória.
Paradas de Pré-processamento no Lado do Host: A construção do contexto evolutivo (ex: Alinhamentos de Múltiplas Sequências, MSA) necessário para a atenção 3D envolve um pré-processamento pesado no CPU, dependente de diversas dependências. Isso cria um tempo de ociosidade da GPU que não pode ser ocultado por pipelines de prefetching convencionais.

Metodologia: MegaFold

O MegaFold é um sistema cross-platform projetado para abordar esses desafios através de um co-design de kernels eficientes em memória, estratégias de comunicação, fusão de operadores e pipelining host-device. Ele é implementado sobre um framework de treinamento de AF3 de código aberto e suporta GPUs NVIDIA e AMD.

1. EvoFlash-3D: Atenção 3D Eficiente em Memória

Para abordar o crescimento cúbico de memória, o MegaFold introduz o EvoFlash-3D, um kernel baseado em Triton para atenção 3D.

Mecanismo: Em vez de materializar o mapa de atenção completo de $O(N^3)$ na memória global (HBM), o kernel utiliza uma estratégia de softmax numericamente estável e online. Ele computa incrementalmente os logits de atenção dentro da memória rápida de scratchpad (shared memory) durante as passagens forward e backward.
Otimização: Isso reduz a complexidade da memória de ativação de pico de $O(N^3)$ para $O(N^2)$ , preservando a correção numérica. O kernel lida com o termo de viés de par broadcastado aditivo único do AF3, que introduz dependências de dados multidimensionais não presentes na FlashAttention padrão.

2. EvoSP-3D: Sharding Eficiente em Comunicação

Para permitir o treinamento em sequências mais longas, o MegaFold propõe o EvoSP-3D, uma estratégia de sharding adaptada para representações pareadas 2D.

Mecanismo: Utiliza uma abordagem híbrida de Paralelismo de Dados (DP) e Paralelismo de Sequência (SP). Enquanto os grupos DP lidam com diferentes entradas de lote, os grupos SP particionam o tensor pareado 2D ( $[N_{token}, N_{token}, h_{pair}]$ ).
Comunicação: À medida que as camadas do AF3 alternam os eixos de atenção (ex: de linha para coluna), o EvoSP-3D utiliza coletivos all_to_all para rotacionar o layout de sharding do tensor pareado entre os estágios. Isso permite que as sucessivas camadas de atenção operem nos cortes 2D corretos sem materializar o tensor completo.
Estágio de Difusão: No decodificador de difusão, as coordenadas de átomos aumentadas são distribuídas entre os grupos SP sem comunicação, enquanto a representação pareada é reunida (all_gather) para permitir que a difusão opere sobre o contexto completo.

3. EvoFusion: Stack de Operadores Fundidos

Para mitigar a fragmentação de kernel, o MegaFold introduz o EvoFusion, que reestrutura a execução interna de operadores críticos em vez de apenas fundir fronteiras.

Escopo: Alvos incluem Layer Normalization, Projeções Lineares e ativações SwiGLU, que coletivamente representam ~20% da memória de ativação e das invocações de kernel.
Mecanismo: Ao reestruturar os loops internos desses operadores, o EvoFusion elimina a materialização intermediária redundante (ex: evitando escrever um tensor completo $N \times N$ na HBM entre a normalização e a projeção linear). Isso reduz a pressão na memória compartilhada e o overhead de lançamento de kernel, mantendo esquemas de redução e GeMM eficientes.

4. EvoPipe: Pipeline Host-Device Determinístico

Para eliminar o tempo de ociosidade da GPU causado pelo pré-processamento de CPU, o MegaFold implementa o EvoPipe.

Mecanismo: Desacopla os passos de pré-processamento determinísticos (ex: construção de átomos via RDKit, extração de características de MSA) do loop de treinamento estocástico. Essas características determinísticas são pré-computadas e armazenadas em cache offline.
Execução: Durante o treinamento, o sistema realiza buscas leves para recuperar características em cache, aplicando apenas aumentos não determinísticos (ex: recorte aleatório) em tempo real. Isso garante uma utilização consistente da GPU ao remover paradas variáveis no lado do host.

Principais Resultados

O sistema foi avaliado em GPUs NVIDIA H200 (141GB) e AMD MI250 (64GB) usando o framework AlphaFold3-PyTorch.

Escalabilidade de Comprimento de Sequência: O MegaFold permite o treinamento com sequências até 3.36× mais longas em 32 GPUs comparado às implementações base do PyTorch. Em 32 GPUs AMD, ele suporta sequências de até ~1450 resíduos, enquanto os baselines falham em ~590 resíduos.
Tempo de Execução: O MegaFold reduz o tempo de execução por iteração em até 1.73× na NVIDIA e 1.62× na AMD.
Eficiência de Memória: O uso de pico de memória da GPU é reduzido em até 1.23× (ou uma redução de 23%), permitindo o realocamento da capacidade de memória para sequências mais longas.
Utilização de FLOPs do Modelo (MFU): O MegaFold atinge um MFU significativamente maior (até 44.9% com 1.81B parâmetros) comparado aos baselines, que frequentemente atingem limites de OOM antes de alcançar escalas maiores.
Desempenho Cross-Platform: O sistema demonstra ganhos consistentes em hardware NVIDIA e AMD, validando a portabilidade dos designs de kernel baseados em Triton.

Significância e Alegações

O artigo afirma que o MegaFold é o primeiro sistema end-to-end a fornecer suporte integrado para o treinamento de modelos de proteína com atenção 3D de próxima geração. Sua significância reside em demonstrar que:

Otimizações de LLM são Insuficientes: Otimizações padrão de transformer (FlashAttention, SP padrão, fusão de compilador) falham em abordar o escalonamento cúbico único e as dependências de dimensões mistas de modelos no estilo AF3.
Co-Design de Sistema é Essencial: O treinamento eficiente desses modelos científicos exige suporte de sistema dedicado que trate conjuntamente a execução de kernel, paralelização, fusão de operadores e pipelining host-device.
Mudança da Barreira de Escala: Ao endereçar os gargalos específicos da atenção 3D, o MegaFold desloca a barreira de OOM, permitindo o treinamento de complexos biomoleculares maiores que eram anteriormente inviáveis, facilitando assim o escalonamento de cargas de trabalho científicas emergentes em plataformas modernas de GPU.