MegaFold: An Open-Sourced AlphaFold-3 Training System (supercomputing-system-ai-lab.github.io)

MegaFold is an open‑source end‑to‑end training system for AlphaFold‑3 (AF‑3) that reworks AF‑3’s inefficient training pipeline with low‑level Triton kernels and system optimizations to substantially reduce runtime and activation memory. The authors profiled AF‑3 and found it far slower and much more memory‑hungry than size‑equivalent transformers (AF‑3 ≈ 500M params vs. BLOOM‑560M): AF‑3 launches tens of thousands of kernels per iteration (~37k linear layers, ~13k LNs, ~13k activations) versus <100 calls per kernel in transformers, and activation storage dominates (≈97% of memory). These bottlenecks limit single‑GPU context lengths (≈380 tokens) and inflate training cost for scientific protein‑folding workloads that require 2‑D pairwise representations and diffusion decoders. MegaFold’s core technical fixes are kernel fusions and memory‑aware operator reimplementation targeted at AF‑3’s “science‑informed” ops: a fused EvoAttention avoids materializing huge intermediate O(N^3) logits by streaming attention computation across thread blocks, and a fused Transition (LN+Linear) computes row statistics then applies layer‑norm and the first MLP multiply in one pass, cutting activation storage from O(Nd) to O(N). Implemented in Triton for performance portability, these fused kernels reduce kernel launch overhead and activation footprint during training, enabling faster iterations, larger effective contexts, and lower GPU memory requirements—making AF‑3 training more practical and accessible for the AI/ML and computational biology communities.