🧬 biology

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

MegaFold 是一种新颖的跨平台系统，它通过结合显存高效内核、优化的分片技术、融合算子以及确定性感知流水线，克服了训练下一代 3D 注意力蛋白质模型时的显存与计算瓶颈，从而在 NVIDIA 和 AMD GPU 上均实现了显著更长的序列长度和更快的训练速度。

原作者： Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

发布于 2026-06-16

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ⚕️ 这是一篇未经同行评审的预印本的AI生成解释。这不是医疗建议。请勿根据此内容做出健康决定。阅读完整免责声明

想象一下，你正试图构建一个巨大且复杂的蛋白质 3D 拼图。长期以来，科学家们一直使用一套特定的规则（比如旧版的 “AlphaFold2”），这些规则效果很好。但最新版本 AlphaFold3 就像是从一个扁平的 2D 拼图升级到了一个复杂的、多层级的 3D 结构，其中每一个碎片都与 3D 空间中的每一个其他碎片相互作用。

虽然这个新模型在科学上非常惊人，但它有一个巨大的问题：它像黑洞一样吞噬着计算机内存，并且运行速度极其缓慢。

这篇论文介绍了一种名为 MegaFold 的新“系统”（一套专门的工具和规则），旨在让这些新型 3D 蛋白质模型能够在现代计算机上进行训练，而不会导致崩溃或耗时过长。

以下是 MegaFold 如何通过简单的类比来解决四个主要问题的：

问题所在：为什么旧的方法会失败

新的 AlphaFold3 模型观察蛋白质的方式会产生“立方级”的数据爆炸。

类比： 想象你正在组织一场派对。在旧模型中，你只需要知道谁和谁是朋友（一份简单的名单）。而在新模型中，你必须追踪一张 3D 地图，记录每一位宾客如何同时与每一位其他宾客进行互动。如果你有 100 位宾客，旧方法需要 10,000 条笔记；而新方法则需要 1,000,000 条笔记。如果你试图把所有这些笔记都写在一本笔记本上（计算机的内存），这本笔记本就会爆炸。

解决方案：MegaFold 的四件工具

MegaFold 通过引入四项具体的升级来修复这个问题：

1. “智能笔记本” (EvoFlash-3D)

问题： 计算机试图一次性写下整个 3D 交互图，瞬间填满了它的内存。
MegaFold 的修复方案： MegaFold 不再尝试把整张地图写在笔记本上，而是使用一个“草稿本”（快速、临时的内存）来逐个、逐块地计算交互，然后立即将笔记丢弃。
结果： 它永远不需要在内存中同时持有整个庞大的地图。这使得计算机能够处理更长的蛋白质链，而不会耗尽空间。

2. “团队接力” (EvoSP-3D)

问题： 当你使用多台计算机（GPU）协同工作时，它们通常通过给每台计算机分配不同的“宾客名单”来分工。但由于 AlphaFold3 的数据是一个 2D 交互网格，简单的拆分并不奏效；计算机会对谁在和谁说话感到困惑。
MegaFold 的修复方案： MegaFold 发明了一种新的数据传递方式。它们不仅是拆分名单，而是拆分了交互的“网格”。它们以一种特定的、高效的模式（类似于接力赛）将拼图碎片来回传递，这样每台计算机都能准确知道自己负责 3D 地图的哪一部分，而不会因为等待他人而浪费时间。

3. “流水线” (EvoFusion)

问题： 计算机之前的操作步骤既细碎又低效。它先计算一个数字，停下来，保存，加载一个新工具，再次计算，停下来，再保存，如此循环。这就像一位厨师切完洋葱，然后走去冰箱拿刀，然后再走回来继续切。
MegaFold 的修复方案： MegaFold 将这些微小的步骤合并为一个巨大的、流畅的动作。它将“切菜”、“搅拌”和“烹饪”融合为一个连续的动作。
结果： 计算机不再浪费时间在任务的启动和停止上。由于它能保持“流水线”不间断地运转，因此运行速度更快。

4. “预备好的食材” (EvoPipe)

问题： 在计算机开始构建蛋白质之前，人类（或缓慢的 CPU）必须进行大量的繁琐研究来收集正确的数据（例如寻找进化史）。这需要很长时间，并导致强大的计算机处于闲置状态，等待数据的到来。
MegaFold 的修复方案： MegaFold 意识到对于相同的蛋白质，其中一些研究工作总是重复的。它会提前完成这些繁重的工作并将数据存储在“储藏室”（缓存）中。当计算机需要数据时，它可以立即抓取这些预备好的食材。
结果： 强大的计算机永远不必等待。它能保持 100% 的忙碌状态。

实验结果

论文在两种类型的强大计算机芯片（NVIDIA 和 AMD）上测试了 MegaFold。以下是发现：

更长的蛋白质： MegaFold 让团队能够训练比以前长 3.36 倍的蛋白质序列。
更快的速度： 它使 NVIDIA 芯片上的训练过程快了 1.73 倍，AMD 芯片上的训练过程快了 1.62 倍。
更少的内存： 它使用了显著更少的内存，防止了通常会在这些模型中发生的“内存溢出”崩溃。

总结

可以把 AlphaFold3 看作是一个动力极强但引擎过重的发动机（相对于其汽车底盘/计算机系统而言）。MegaFold 不仅仅是造了一辆更大的车；它重新设计了引擎、变速器和燃油系统，使得这个超强力的引擎能够高效运行。这使得科学家能够在更广泛的计算机上训练这些下一代模型，让蛋白质折叠的研究变得更快、更普及。

技术摘要：MegaFold

问题陈述

生物分子建模领域的最新进展，特别是 AlphaFold3 (AF3)，引入了从标准 Transformer 架构向利用 2D 成对表示（2D pairwise representations）进行 3D 注意力计算的模型范式的转变。虽然 AF3 能够实现精确的蛋白质-蛋白质、蛋白质-DNA/RNA 以及蛋白质-配体相互作用预测，但这种架构变化引入了现有大语言模型（LLM）优化手段无法解决的严重系统级瓶颈。

本文识别了四个源于 3D 注意力立方级缩放（cubic scaling）的基础性挑战：

立方级激活内存（Cubic Activation Memory）： 与注意力机制随序列长度平方级（ $O(N^2)$ ）增长的 Transformer 不同，AF3 在 2D 成对表示上的 3D 注意力产生的张量规模呈立方级（ $O(N^3)$ ）增长。这导致激活内存占据了极高的 GPU 使用率（高达总内存的 96%），即使在小批量（batch size）的情况下，在处理中等长度序列（例如 >590 个残基）时也会导致显存溢出（OOM）错误。
与序列并行（Sequence Parallelism, SP）的不兼容性： 标准 Transformer 的 SP 策略依赖于 1D 序列分片。而 AF3 的 2D 成对表示以及跨层交替的归约轴（reduction axes）创建了复杂的多轴通信模式，传统的 SP 无法在不进行过度重分片（resharding）或同步的情况下支持这些模式。
内核碎片化（Kernel Fragmentation）： 3D 注意力、转换层（transition layers）以及针对混合维度张量的投影操作之间的交错，导致每次迭代产生数以千计的细粒度内核（kernels）。这种碎片化阻碍了有效的内核融合（kernel fusion），导致了较低的 Warp 占用率、高昂的启动开销以及低效的内存带宽利用。
主机端预处理停顿（Host-Side Preprocessing Stalls）： 构建 3D 注意力所需的进化上下文（例如多序列比对，MSA）涉及繁重的、具有高度依赖性的 CPU 预处理。这造成了显著的 GPU 空闲时间，且无法通过传统的预取流水线进行隐藏。

方法论：MegaFold

MegaFold 是一个跨平台的系统，旨在通过内存高效内核、通信策略、算子融合以及主机-设备流水线的协同设计来解决上述挑战。它构建在开源的 AF3 训练框架之上，并同时支持 NVIDIA 和 AMD GPU。

1. EvoFlash-3D：内存高效的 3D 注意力

为了解决立方级内存增长问题，MegaFold 引入了基于 Triton 的 3D 注意力内核 EvoFlash-3D。

机制： 该内核不再在全局内存（HBM）中实例化完整的 $O(N^3)$ 注意力图，而是采用一种在线、数值稳定的 Softmax 策略。它在正向和反向传播过程中，在快速暂存内存（shared memory）中增量地计算注意力逻辑值（logits）。
优化： 这将峰值激活内存复杂度从 $O(N^3)$ 降低到 $O(N^2)$ ，同时保持了数值正确性。该内核处理了 AF3 特有的加性广播对偏置项（additive broadcasted pair-bias term），该项引入了标准 FlashAttention 中不存在的多维数据依赖。

2. EvoSP-3D：通信高效的分片策略

为了支持更长序列的训练，MegaFold 提出了针对 2D 成对表示定制的分片策略 EvoSP-3D。

机制： 它采用了混合数据并行（DP）和序列并行（SP）的方法。当 DP 组处理不同的批次输入时，SP 组则对 2D 成对张量（ $[N_{token}, N_{token}, h_{pair}]$ ）进行分区。
通信： 由于 AF3 层会交替改变注意力轴（例如，行向到列向），EvoSP-3D 通过 all_to_all 集合通信操作在不同阶段旋转 2D 成对张量的分片布局。这使得连续的注意力层能够在不实例化完整张量的情况下，在正确的 2D 切片上进行操作。
扩散阶段： 在扩散解码器中，增强后的原子坐标分布在各个 SP 组中而不进行通信，而成对表示则通过 all_gather 进行聚合，以允许扩散过程在完整的上下文中运行。

3. EvoFusion：融合算子栈

为了缓解内核碎片化，MegaFold 引入了 EvoFusion，它通过重构关键算子的内部执行逻辑，而非仅仅是融合边界来发挥作用。

范围： 目标涵盖层归一化（Layer Normalization）、线性投影（Linear Projections）和 SwiGLU 激活函数，这些算子合计占据了约 20% 的激活内存和内核调用。
机制： 通过重构这些算子的内层循环，EvoFusion 消除了冗余的中间状态物化（例如，避免在归一化和线性投影之间将完整的 $N \times N$ 张量写入 HBM）。这降低了共享内存压力和内核启动开销，同时保持了高效的归约（reduction）和 GeMM 调度。

4. EvoPipe：确定性感知的主机-设备流水线

为了消除由 CPU 预处理引起的 GPU 空闲时间，MegaFold 实现了 EvoPipe。

机制： 它将确定性的预处理步骤（例如基于 RDKit 的原子构建、MSA 特征提取）与随机的训练循环解耦。这些确定性特征被预先计算并进行离线缓存。
执行： 在训练期间，系统执行轻量级的查找以检索缓存的特征，仅在运行时应用非确定性的增强操作（例如随机裁剪）。这通过消除可变的宿主侧停顿，确保了持续的 GPU 利用率。

关键结果

该系统在 NVIDIA H200 (141GB) 和 AMD MI250 (64GB) GPU 上使用 AlphaFold3-PyTorch 框架进行了评估。

序列长度可扩展性： 与基准 PyTorch 实现相比，MegaFold 在 32 块 GPU 上支持的训练序列长度最高可达 3.36 倍。在 32 块 AMD GPU 上，它支持高达 ~1450 个残基的序列，而基准实现会在 ~590 个残基处失败。
执行时间： MegaFold 将端到端的单次迭代执行时间分别降低了高达 1.73× (NVIDIA) 和 1.62× (AMD)。
内存效率： 峰值 GPU 内存使用量降低了高达 1.23×（即减少了 23%），从而可以将释放出的内存容量重新分配给更长的序列。
模型 FLOPs 利用率 (MFU)： 与基准实现相比，MegaFold 实现了显著更高的 MFU（在 1.81B 参数下高达 44.9%），而基准实现往往在达到更大规模之前就触发了 OOM 错误。
跨平台性能： 该系统在 NVIDIA 和 AMD 硬件上均表现出一致的增益，验证了其基于 Triton 的内核设计的可移植性。

意义与主张

论文声称 MegaFold 是第一个为下一代 3D 注意力蛋白质模型提供集成训练支持的端到端系统。其重要性在于证明了：

LLM 优化是不充分的： 标准的 Transformer 优化（如 FlashAttention、标准 SP、编译器融合）无法解决 AF3 类模型特有的立方级缩放和混合维度依赖问题。
系统级协同设计至关重要： 高效训练此类科学模型需要专门的系统支持，以共同解决内核执行、并行化、算子融合以及主机-设备流水线的问题。
可扩展性障碍的转移： 通过解决 3D 注意力的特定瓶颈，MegaFold 转移了 OOM 屏障，使得训练此前无法实现的更大规模生物分子复合物成为可能，从而促进了新兴科学工作负载在现代 GPU 平台上的规模化发展。