突破蛋白复合物建模的内存瓶颈,支持超3万残基结构预测。
Fold-CP: A Context Parallelism Framework for Biomolecular Modeling
- 通过多卡上下文并行,拆分模型计算与存储负担。
- 64块B300显卡可处理超3万残基系统,内存开销随卡数线性下降。
- 适用于大规模蛋白复合物建模,助力虚拟细胞研究。
解析细胞机器需对大型生物分子复合物进行原子级重构。然而,如AlphaFold 3等模型受硬件内存限制,单卡仅能处理数千残基,形成实际瓶颈。本文提出NVIDIA BioNeMo Fold-CP,一种基于上下文并行的框架,将共折叠模型的推理与训练流程分布到多块GPU上。采用Boltz作为开源参考架构,实现定制化的多维原语,高效并行化密集三角更新与窗口批处理的局部注意力不规则数据依赖模式。该方法实现高效内存扩展:对于输入长度为N、分布在P块GPU上的任务,单卡内存消耗为$O(N^2/P)$,使64块NVIDIA B300 GPU可处理超过30,000残基的系统。通过开发者案例验证其科学价值:成功评分90%以上的哺乳动物蛋白复合物数据库(CORUM),并完整折叠了与无序区域结合的疾病相关PI4KA脂质激酶复合物,无需截断。Fold-CP为全全局上下文下建模巨型系统提供了可扩展路径,是迈向虚拟细胞的重要一步。
原文摘要 · Abstract (English)
Understanding cellular machinery requires atomic-scale reconstruction of large biomolecular assemblies. However, predicting the structures of these systems has been constrained by hardware memory requirements of models like AlphaFold 3, imposing a practical ceiling of a few thousand residues that can be processed on a single GPU. Here we present NVIDIA BioNeMo Fold-CP, a context parallelism framework that overcomes this barrier by distributing the inference and training pipelines of co-folding models across multiple GPUs. We use the Boltz models as open source reference architectures and implement custom multidimensional primitives that efficiently parallelize both the dense triangular updates and the irregular, data-dependent pattern of window-batched local attention. Our approach achieves efficient memory scaling; for an N-token input distributed across P GPUs, per-device memory scales as $O(N^2/P)$, enabling the structure prediction of assemblies exceeding 30,000 residues on 64 NVIDIA B300 GPUs. We demonstrate the scientific utility of this approach through successful developer use cases: Fold-CP enabled the scoring of over 90% of Comprehensive Resource of Mammalian protein complexes (CORUM) database, as well as folding of disease-relevant PI4KA lipid kinase complex bound to an intrinsically disordered region without cropping. By providing a scalable pathway for modeling massive systems with full global context, Fold-CP represents a significant step toward the realization of a virtual cell.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。