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

arXiv: https://arxiv.org/abs/2506.20686 · HTML: https://mlsanity.com/paper/2506.20686/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 57 references that are in the map.
- First posted on arXiv 24 Jun 2025.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Large-Scale Chemical Language Representations Capture Molecular Structure and Properties](https://mlsanity.com/paper/2106.09553/) (2021)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets](https://mlsanity.com/paper/2203.04810/) (2022)
- [FFCV: Accelerating Training by Removing Data Bottlenecks](https://mlsanity.com/paper/2306.12517/) (2023)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies](https://mlsanity.com/paper/2310.04610/) (2023)
- [FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours](https://mlsanity.com/paper/2203.00854/) (2022)
- [ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours](https://mlsanity.com/paper/2404.11068/) (2024)
- [Blockbuster, Part 1: Block-level AI Operator Fusion](https://mlsanity.com/paper/2505.07829/) (2025)

## Built on by (most cited first, 1 of 1)

- [TurboPairFormer: Fast and Stable Protein Folding Model Training with an Optimized Triangle Attention Kernel](https://mlsanity.com/paper/2610.05854/) (2026)

## Abstract (arXiv)

Recent advances in biomolecular modeling have been catalyzed by models such as AlphaFold3 (AF3), which introduce science-informed changes to the transformer architecture. Unlike transformers, a defining characteristic of AF3-style models is their 3D attention over 2D pairwise representations which produces tensors whose computation and memory costs scale cubically with sequence length. As a result, despite moderate parameter counts, AF3-style models are far more expensive to train than size-equivalent transformers, and are severely constrained by GPU memory capacity. Our characterization shows 3D attention fundamentally changes the training workload, causing massive 3D attention maps, complex inter-operator dependencies, kernel fragmentation, and heavy host-side data pipelines which differ substantially from LLM training, leading to poor utilization on modern GPU systems. Moreover, existing GPU optimizations do not adequately address these challenges due to complex cross-layer inter-operator dependencies introduced by 3D attention. Motivated by these challenges, we introduce MegaFold, a novel cross-platform system for efficient training of next-generation 3D-attention protein models. MegaFold combines a memory-efficient 3D-attention kernel, a communication-efficient sharding strategy for quadratic representations, fused operator implementations for critical execution paths, and a determinism-aware host-device pipeline that eliminates preprocessing stalls. Evaluation on both NVIDIA H200 and AMD MI250 GPUs shows that MegaFold enables training with up to 3.36× longer sequence lengths on 32 GPUs while reducing end-to-end execution time by up to 1.73× (NVIDIA) and 1.62× (AMD).