MegaFold是一个跨平台系统，用于加速蛋白质结构预测模型（例如 AlphaFold3、AlphaFold2）。
为什么选择 MegaFold？
- 跨平台支持：通过优化的基于 Triton 的内核，支持在异构设备上执行，包括 NVIDIA GPU 和 AMD GPU。
- 易于使用：只需更改几行代码即可获得巨大的性能提升
- 速度提升：每次迭代训练时间加快高达 1.73 倍
- 减少内存：将训练期间的峰值内存减少高达 1.23 倍
- 序列长度扩展：支持对 1.35 倍序列长度进行训练
1. MegaFold 针对 AlphaFold3 的关键训练瓶颈，显著节省内存并提升速度，使得训练 1.35 倍长的蛋白质序列时不会出现内存溢出错误。该算法已在 NVIDIA 和 AMD GPU 上进行评估，以证明其广泛的适用性。
2. MegaFold 的核心是对 AlphaFold3 的 EvoAttention 进行内存高效的重新实现——EvoAttention 是一种昂贵的二维注意力机制，具有立方内存复杂度。MegaFold 通过在快速暂存器中动态重新计算中间逻辑来避免存储它们。
3. MegaFold 为常用但内存受限的运算符（例如层范数和线性层）引入了一种深度核融合策略。这种“DeepFusion”策略减少了核启动次数，避免了不必要的内存流量，从而提升了性能。
4. 该系统通过基于缓存的提前数据加载器消除了繁重的 CPU 端预处理（例如 MSA 和模板检索）导致的 GPU 空闲时间。仅此一项就将预处理时间减少了 3.34 倍，并加快了端到端迭代速度。
5. 与 PyTorch 基线相比，MegaFold 在 NVIDIA H200 和 AMD MI250 上的每次迭代加速分别提高了 1.73 倍和 1.62 倍，同时将峰值内存使用量降低了 1.23 倍。
6. 与 PyTorch-Inductor 或 FastFold 等现有系统相比，MegaFold 是第一个支持 AlphaFold3 跨平台训练（NVIDIA + AMD）的系统，并且将 H200 上的序列长度可行性从 512 显著扩展到 768。
7. 与仅针对标准注意力机制的 FlashAttention 不同，MegaFold 融合的 EvoAttention 支持加性偏差和动态 softmax 归一化。这使得它在保持生物学准确性的同时实现了高计算效率。
8. AlphaFold3 的扩散模块改进最为显著，其中数据增强和二维输入加剧了内存问题。MegaFold 将内存占用减少了 1.38 倍，尽管复杂度较高，但执行速度仍然加快。
9. 在 8 GPU 设置下，MegaFold 相比 PyTorch 保持了 1.46 倍的每次迭代加速，表明其优势在规模化情况下依然有效，并且不会引入分布式通信开销。
10. MegaFold 的代码完全开源，并建立在流行的基于 PyTorch 的 AlphaFold3 实现之上，从而支持计算生物学中可重复和可扩展的系统研究。
代码：https://github.com/Supercomputing-System-AI-Lab/MegaFold
论文：https://arxiv.org/abs/2506.20686v1
