🧬 MegaFold: Efficient AlphaFold-3 Training Across NVIDIA and AMD GPUs MegaFold will be presented at ISC High Performance 2026! 🗺️Research Paper Session: ML Systems and Frameworks 🗓️ Wednesday, June 24 · 1:00–2:00 PM (CEST) 📍Hall E, 2nd Floor, Hamburg, Germany Come find out how we train AlphaFold-3 on 3.36× longer sequences, up to 1.73× faster, on both NVIDIA and AMD GPUs. Using MegaFold, we can successfully train on H200s, MI250 and MI325X. 🚀 Think training AlphaFold-3 (AF3) is just like training an LLM? Think again. Despite having only ~500M parameters, AF3 is up to two orders of magnitude slower to train than a size-equivalent LLM and OOMs on sequences as short as 590 residues. Why? One operator: 3D attention over 2D pairwise representations. It makes compute and memory scale cubically with sequence length, blowing up activation memory, fragmenting execution into tens of thousans of tiny kernels per iteration, and starving the GPU while the CPU grinds through evolutionary-context preprocessing. In our latest work, MegaFold, we diagnose exactly why AF3 training is slow and build an end-to-end system that fixes it across the entire stack. The results: ✅ 3.36× longer trainable sequences on 32 GPUs ✅ Up to 1.73× faster per-iteration training (1.62× on AMD) ✅ 23% lower peak memory ✅ Runs on both NVIDIA H200, AMD MI250 & AMD MI325x, fully cross-platform 🔧 Four key innovations: - EvoFlash-3D: a cross-platform Triton kernel that computes 3D attention without ever materializing the cubic attention map (cutting EvoAttention memory by 4.24×). - EvoSP-3D: a communication-efficient sharding strategy built for alternating-axis 2D pairwise representations - EvoFusion: a fused operator stack that tames AF3's kernel fragmentation - EvoPipe: a determinism-aware host-device pipeline that eliminates preprocessing stalls And it's verified: MegaFold matches baseline training loss step-for-step, so you get all the speed and memory savings with zero impact on convergence. Huge thanks to our students Hoa La and Ahan Gupta for working hard on this project, and to our collaborators Jianlin Cheng and Alex Morehead for their valuable contributions, feedback, and support throughout the work! MegaFold is fully open source. 🎉 💻 Code: github.com/Supercomputing-Sy… 🔗 Project page: supercomputing-system-ai-lab… 📄 Paper: arxiv.org/pdf/2506.20686 📝 Blog: supercomputing-system-ai-lab… See you at ISC, Hall E! 👋

🔬Interested in training AlphaFold3 faster, at scale, and beyond NVIDIA GPU? Now you can. AlphaFold3 is a major leap in biomolecular modeling, but behind the scenes, it introduces severe system bottlenecks: 🧠 2D EvoAttention spikes memory usage 📉 Retrieval-augmented training pipeline causes long GPU idle time ⛔ Frequent but memory-intensive ops slow everything down Today, I'm excited to announce MegaFold, a fully open-source system to make AlphaFold3 training fast, scalable, and cross-platform on both NVIDIA and AMD GPUs. MegaFold delivers: ⚡ Up to 1.73x / 1.62x faster training on NVIDIA H100 / AMD MI250 🧬 Up to 1.35× longer sequences compared to PyTorch baseline Key features: 🚀 Memory-Efficient EvoAttention via portable Triton kernels 💡 Ahead-of-Time Caching to eliminate GPU idle time in retrieval pipelines 🔗 DeepFusion for reducing overhead of small but frequent memory-intensive AF3 ops 📘 Project page: supercomputing-system-ai-lab… 📄 Paper: arxiv.org/pdf/2506.20686 💻 Code: github.com/Supercomputing-Sy… 🤝 MegaFold is developed in collaboration between UIUC SSAIL Lab and researchers from University of Missouri and Lawrence Berkeley National Laboratory. Kudos to the brilliant team: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng #AlphaFold3 #AI #ProteinFolding #Bioinformatics #AMD #Triton #CrossPlatform #OpenSource