使用 NVIDIA BioNeMo 中的上下文并行扩展生物分子建模

几十年来，计算生物学一直在一种还原论式的折中方案下运行。为了将复杂的生物系统装入单个 GPU 有限的内存中，研究人员不得不将它们拆解成孤立的片段——单个蛋白质或小结构域。这造成了上下文缺口，即由于 GPU 硬件内存限制，较大的蛋白质或复合物无法以零样本方式进行折叠。

现在，NVIDIA BioNeMo 团队推出的一种新的上下文并行（CP）框架正在打破结构生物学的内存壁垒，使系统的整体建模成为可能。

本文介绍如何在不同于标准 Transformer 的生物分子架构中实现 CP。如果你是结构生物学家、计算化学家或机器学习工程师，希望在不牺牲全局上下文的情况下对大型生物分子复合物进行建模，请继续阅读。

要使用本文概述的解决方案，你需要：

• 熟悉 AlphaFold3 或 Boltz-2 等几何深度学习基础模型。
• 理解 PyTorch Distributed（DTensor）操作和自定义 autograd 函数。
• 能够访问 NVIDIA H100 或 B200 GPU 集群，因为该框架在执行百亿亿次级任务时高度依赖其互连带宽和 Transformer Engine 加速。

更多详情请参见 Fold-CP: A Context Parallelism Framework for Biomolecular Modeling。

将单个大型分子系统分片到多块 GPU 上

在没有 CP 的情况下，折叠大型复合物（通常超过 1,000–3,000 个残基）需要采用还原论方法，即将系统在物理上或计算上拆解为可管理的块。这些方法使研究人员能够保持在单块 GPU 严格的 VRAM 限制内，但往往会牺牲全局结构准确性。

针对超大蛋白质最常见的变通方法是将序列切分成更小的重叠片段。片段之间必须有显著重叠，以确保局部二级结构在切分点两侧保持一致。这种方法会破坏长程信息。例如，研究人员无法对整个复合物范围内的变构作用或信号转导进行建模。

另一种常见的变通方法是分块；与物理上的序列片段化不同，分块发生在模型架构内部，用于在推理期间节省 VRAM。Boltz 等模型使用激进的分块方式，将大型矩阵处理为更小的瓦片。FastFold 等其他技术则采用自动分块，以动态调整分块策略并改善峰值内存使用。若要了解更多信息，请参见 FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours。

所有这些技术都天然存在缺乏全局上下文的问题，尤其是在训练过程中。NVIDIA BioNeMo CP 框架通过将一个大型分子系统分片到多个 GPU 上来克服这些限制。与传统的数据并行不同，传统数据并行会为每个 GPU 分配一个不同的蛋白质进行折叠，而 CP 会将一个超大样本拆分到多个 GPU 上。

BioNeMo 上下文并行实现

NVIDIA BioNeMo CP 实现基于 Torch 分布式 API，用于 GPU 到 GPU 的通信。该架构自底向上构建，从低级通信协议开始，逐步上升到高级的模型特定工作流。本文使用 Boltz 作为示例代码库。

为实现线性容量扩展——即系统能力随 GPU 数量线性增长——该框架采用了多维分片策略。这确保没有任何单个设备持有生物分子的完整全局状态，否则将违背 CP 的内存目标。针对不同模块构建了类似乐高的自定义组件，以提高实现效率并更容易移植到其他架构。

成对表示的二维平铺

该框架将全局 (N x N) 矩阵划分为由多个块组成的网格。对于一个包含 10,000 个残基的复合物，即代表 1 亿次相互作用，每个 GPU 只管理一个特定的子块。这将每个设备的内存占用从 O(N2) 局部化为 O(N2/P)。

重叠计算与通信

该框架实现了分布式原语，用于协调本地计算与异步点对点传输。当 GPU 正在计算本地更新时，它会同时向其在行环和列环中的邻居发送数据并从中接收数据。随着生物学问题规模的增长，计算与通信的比例得到改善，使系统在更大规模下更加高效。

高效的原子序列局部注意力

Alphafold3 序列局部注意力将原子注意力限制在大小为 32 乘 128 的局部注意力窗口中，这些窗口的堆叠会以批处理方式进行注意力计算。上述原子特征的平铺需要针对这种基于窗口批处理的原子注意力的分布式版本重新分区。NVIDIA 团队实现了基于 halo exchange 的分布式原语来划分原子特征，使得后续的窗口批处理注意力无需 GPU 间通信。

三角乘法的上下文并行实现

以下示例展示了如何创建一个支持 CP 的层，以分布式方式计算三角乘法。

 # … torchrun or srun SPMD launcher sets up the environment
   
   # Initialize the grid of devices 
   DistributedManager.initialize(device_type="cuda")
   manager = DistributedManager()

   # Create a square 2D device mesh for symmetrical communication
   size_ring = math.isqrt(manager.world_size)
   DistributedManager.create_grid_group({"dp": 1, "cp": (size_ring, size_ring)})


   # Instantiate the specialized peer-to-peer communication handle
   ring_comm = Ring2DComm(manager.group["cp"], manager.subgroups["cp"][0], manager.layout_subgroups["cp"])


   # Data processing or previous layer output
   x_dtensor, mask_dtensor = … 


   # Instantiate standard layer and load the checkpoint on CPU before distributing across the device grid
   layer_serial = TriangleMultiplicationOutgoing(size_input_embed)
   layer_serial.load_state_dict(layer_state_dict)

   # Map standard layer to GPU device
   layer_serial = layer_serial.to(manager.device)

   # Wrap with BioNeMo CP to handle DTensors
   layer = DistributedTriangleMultiplication(Outgoing, layer_serial, manager.device_mesh_subgroups, ring_comm)

   # The resulting activation tensors are now sharded across the device grid 
   result_dtensor = layer(x_dtensor, mask_dtensor)
   DistributedManager.cleanup()

这段代码通过使用 DistributedManager 设置一个方形二维设备网格来初始化 CP 环境——这是一项特定的架构要求，可确保按行和按列的通信模式保持对称且高效。随后实例化 Ring2DComm 句柄来管理专门的点对点通信，从而允许数据块在连续循环中流转。这种环形方法至关重要，因为它能够将本地计算与数据传输重叠进行，确保 \(O(N^2)\) 的成对表示张量永远不会超出单个 GPU 的内存容量。

脚本的后半部分处理从标准串行模型到基于 DTensor 的模型的转换。标准层（例如 TriangleMultiplicationOutgoing）会先加载到 CPU 上，然后由 DistributedTriangleMultiplication 进行封装，后者实现了适配 NVIDIA BioNeMo CP 框架的分布式版本算法。通过将输入作为分布式张量（DTensors）进行处理，该模型确保大型激活张量在网格中被分片。

释放结构生物学的 token 扩展能力

图 2 显示，随着 CP 的引入，生物分子架构的 token 容量缩放定律如今已被解锁。Boltz 预测可使用 256 个 GPU 在高达约 20,000 个 token 上运行，并可在 NVIDIA H100 GPU 上扩展最大 token 长度，在 NVIDIA B300 GPU 上实现加速扩展。

在没有使用更长裁剪长度进行任何额外训练或微调的情况下，该团队折叠了一个 TTC7A/PI4KA/FAM126A/EFR3A(700–823) 系统，该系统包含四条链上的 3,605 个残基——远超 Boltz-2 训练裁剪大小的 768 个残基以及单个 GPU 的内存容量。使用 CP 后，可在四个 NVIDIA H100 GPU 上于五分钟内生成五个结构样本（每个样本约 54 秒），同时将所有长程亚基间接触保持在模型上下文窗口内。

与此同时，该团队还正与 Rezo Therapeutics、Proxima 和 Earendil Labs 等合作者共同推进这一下一代前沿方向；这些合作者为 CP 框架的发展作出了深入贡献。

Rezo Therapeutics 集成了 CP 框架，用于预测涵盖最多 6,500 个残基的大规模蛋白质-蛋白质相互作用（PPIs），从而能够对绝大多数已知蛋白质复合物进行结构预测，并加速发现新型复合物。事实上，与仅使用公共领域高置信度 PPIs 进行的预测相比，由 CP 解析出的高质量新型蛋白质复合物富集程度提高了 3 倍以上。

Proxima 将 CP 嵌入其全原子生成式基础模型 Neo 中，使其能够对最多 4,000 个 token 的组装体进行推理，从而在结构层面解析其独特质谱平台所绘制的蛋白质组中具有治疗相关性的相互作用，进而助力其在开发分子胶及其他基于邻近性的疗法方面的药物发现工作。

Earendil Labs 将该框架集成到其专有的生物分子基础模型中，成功扩展了输入序列长度，从而能够对以往在计算上难以承受的复杂多蛋白系统进行建模。Earendil Labs 证明，CP 有潜力在序列复杂性扩展的同时保持高保真结构预测，并缩短下一代生物治疗药物的发现周期。

开始使用用于生物分子建模的上下文并行

虽然初步概念验证表明，CP 打破了此前的内存壁垒，使得对无限大小结构进行建模成为可能，但仅有物理容量并不能保证生物学准确性。当前模型往往难以在大规模条件下进行高保真折叠，因为它们是在小片段上训练的。使用更大的裁剪尺寸进行微调，对于准确捕捉长程相互作用的涌现逻辑至关重要。

作为这项工作的一部分，该领域的主要瓶颈之一——数据稀缺——正通过对 AlphaFold Protein Structure Database 的贡献得到解决。通过利用 NVIDIA 加速计算软件（如 NVIDIA cuEquivariance 和 NVIDIA TensorRT），以大规模同源和异源复合物的高通量预测来填充 AFCDB，这为构建合成复合物数据奠定了基础，而这些数据可能有助于训练能够表征生物学中更大系统的基础模型。

要了解更多信息，请参阅 Boltz CP code 的开源文档，并查看 Fold-CP: A Context Parallelism Framework for Biomolecular Modeling。