NVIDIA BioNeMo 使用上下文并行扩展生物分子建模
realtime news Apr 28, 2026 19:53
NVIDIA 的 BioNeMo 引入了上下文并行，通过克服 GPU 内存限制，支持大规模系统的生物分子建模。
几十年来，计算生物学领域的研究人员一直在与一个关键限制作斗争：GPU 的内存容量。对包含数千个残基的蛋白质复合物等大型生物分子系统进行建模，通常需要将它们分割成更小的片段，从而牺牲了理解生物学相互作用所需的全局上下文。NVIDIA 的 BioNeMo 团队现在引入了一项突破性技术：上下文并行（CP），这是一种通过跨多个 GPU 分片数据实现对大规模生物分子系统的整体建模的新框架。
通过上下文并行突破内存障碍
传统的大型蛋白质折叠方法依赖于分割序列或使用诸如分块这样的激进内存节省技术。虽然这些方法可以有效地将数据适配到单个 GPU，但往往会损害远程结构信息。NVIDIA BioNeMo 的 CP 框架通过将单个大型生物分子系统分割到多个 GPU 上，而不是为每个 GPU 分配单独的任务，消除了这一权衡。此方法在扩展计算能力的同时，保留了全局结构上下文，且计算能力与 GPU 数量线性增长。
CP 的实现利用了 NVIDIA 的先进 GPU 技术，特别是 H100 和 B300 集群，以及 PyTorch Distributed API。通过将蛋白质的结构数据分片到一组 GPU 上，本地化内存使用，且没有单个 GPU 承担全部计算负载。这使得研究人员能够建模包含数万个残基的系统——远远超出了传统方法的限制。
CP 框架中的技术创新
CP 框架引入了多项创新以优化性能：
- 二维划分：将蛋白质相互作用矩阵划分为子块，将内存需求从 O(N2) 减小到 O(N2/P)，其中 P 是 GPU 的数量。
- 重叠计算与通信：GPU 在执行本地计算的同时，与相邻 GPU 异步交换数据，随着问题规模增加提高效率。
- 高效的本地注意力机制：分布式原语在本地注意力计算中最小化 GPU 之间的通信，这对于处理大规模的标记长度至关重要。
在一个概念验证中，NVIDIA 展示了该框架的能力，在四个 GPU 上在五分钟内折叠了一个超过 3600 个残基的复杂生物分子系统，同时保持结构准确性。这标志着建模能力的显著飞跃。
实际应用及行业影响
一些行业参与者已经利用 CP 框架解决了此前无法克服的挑战：
- Rezo Therapeutics：使用 CP 对多达 6500 个残基的蛋白质-蛋白质相互作用进行建模，从而发现了新的复合物。
- Proxima：将 CP 集成到其 Neo 生成模型中，使治疗相关相互作用的结构分辨率更加详细。
- Earendil Labs：扩展 CP 以建模高度复杂的多蛋白系统，加速了生物治疗的发现时间表。
生物分子建模的下一步
尽管 CP 已经打破了内存障碍，NVIDIA 承认仅靠物理容量并不能保证生物学的准确性。目前的模型是在较小的蛋白质片段上训练的，需要使用更大的数据集进行微调，以完全捕捉远程相互作用。NVIDIA 正通过对 AlphaFold 蛋白质结构数据库的贡献来解决这一问题，并使用如 cuEquivariance 和 TensorRT 等加速软件工具来增强训练未来模型的数据可用性。
有兴趣探索 CP 框架的研究人员可以通过 Boltz CP GitHub 仓库 获取开源文档，或通过 Fold-CP 研究论文 深入了解技术细节。
Image source: Shutterstock.jpg)