Fold-CP: A Context Parallelism Framework for Biomolecular Modeling

• URL: http://arxiv.org/abs/2603.14806v1
• Date: Mon, 16 Mar 2026 04:20:01 GMT
• Title: Fold-CP: A Context Parallelism Framework for Biomolecular Modeling
• Abstract summary: We present NVIDIA BioNeMo Fold-CP, a context parallelism framework.<n>We use the Boltz models as open source reference architectures.<n>We implement custom multidimensional primitives that efficiently parallelize both the dense triangular updates and the irregular, data-dependent pattern of window-batched local attention.
• Score: 11.44600151695527
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Multimodal Molecular Representation Learning with Graph Neural Networks, Deep & Cross Networks, and SMILES Embeddings [2.9593641578445258]
  We introduce a parameter-efficient Tri-Branch Modular Fusion Network that synthesizes three modalities: 3D spatial geometry (SchNet), discrete topological grammar (SMILES via ChemBERTa), and explicit macroscopic physicochemical descriptors (Deep & Cross Network)<n> operating with fewer than one million parameters, this architecture decisively surpasses the sub-chemical accuracy threshold and yields a substantial 20.6% reduction over a strictly controlled geometric baseline.
  arXiv  Detail & Related papers  (2026-07-07T01:47:44Z)
• Scaling Parallel Sequence Models to Foundation-Scale Vision Encoders [123.58723804218151]
  Vision foundation models are bottlenecked by the quadratic cost of self-attention.<n>We present C-GSPN based on 2D spatial propagation propagation.
  arXiv  Detail & Related papers  (2026-05-30T14:29:43Z)
• Scalable Training of Mixture-of-Experts Models with Megatron Core [26.9162079065285]
  Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models.<n>Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation.<n>We address these challenges for MoE training through integrated optimizations spanning memory (fine-dimensional recomputation, offloading), communication, and computation.
  arXiv  Detail & Related papers  (2026-03-08T15:42:43Z)
• Space Filling Curves is All You Need: Communication-Avoiding Matrix Multiplication Made Simple [42.09057806159106]
  General Matrix multiplication is the cornerstone of Deep Learning and HPC workloads.<n>Modern platforms with matrix multiplication accelerators exhibit high FLOP/Byte machine balance.<n>In this work we revisit space filling curves (SFC) to alleviate the problem of this cumbersome tuning.<n>We obtain platform-oblivious and shape-oblivious matrix-multiplication schemes that exhibit inherently high degree of data locality.
  arXiv  Detail & Related papers  (2026-01-22T19:56:16Z)
• GSPN-2: Efficient Parallel Sequence Modeling [101.33780567131716]
  Generalized Spatial Propagation Network (GSPN) addresses this by replacing quadratic self-attention with a line-scan propagation scheme.<n>GSPN-2 establishes a new efficiency frontier for modeling global spatial context in vision applications.
  arXiv  Detail & Related papers  (2025-11-28T07:26:45Z)
• Preparation of Fractal-Inspired Computational Architectures for Advanced Large Language Model Analysis [50.11146543029802]
  It introduces FractalNet, a fractal-inspired computational architectures for advanced large language model analysis.<n>The new set-up involves a template-driven generator, runner, and evaluation framework that, through systematic permutations of convolutional, normalization, activation, and dropout layers, can create more than 1,200 variants of neural networks.<n>The paper positions fractal design as a feasible and resource-efficient method of automated architecture exploration.
  arXiv  Detail & Related papers  (2025-11-10T17:31:39Z)
• Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
  We introduce the ''Three Taxes'' (Bulk Synchronous, Inter- Kernel Data Locality, and Kernel Launch Overhead) as an analytical framework.<n>We propose moving beyond the rigid BSP model to address key inefficiencies in distributed GPU execution.<n>We observe a 10-20% speedup in end-to-end latency over BSP-based approaches.
  arXiv  Detail & Related papers  (2025-11-04T01:15:44Z)
• Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction [76.62155593340763]
  Equivariant Graph Neural Networks (eGNNs) trained on density-functional theory (DFT) data can potentially perform electronic structure prediction at unprecedented scales.<n>However, the graph representations required for this task tend to be densely connected.<n>We present a distributed eGNN implementation which leverages direct GPU communication and introduce a partitioning strategy of the input graph.
  arXiv  Detail & Related papers  (2025-07-04T23:53:47Z)
• Unified Cross-Scale 3D Generation and Understanding via Autoregressive Modeling [32.45851798752336]
  We propose Uni-3DAR, a unified autoregressive framework for cross-scale 3D generation and understanding.<n>At its core is a coarse-to-fine tokenizer based on octree data structures, which compresses diverse 3D structures into compact 1D token sequences.<n>To address the challenge of dynamically varying token positions introduced by compression, we introduce a masked next-token prediction strategy.
  arXiv  Detail & Related papers  (2025-03-20T16:07:04Z)
• fVDB: A Deep-Learning Framework for Sparse, Large-Scale, and High-Performance Spatial Intelligence [55.582429009401956]
  fVDB is a novel framework for deep learning on large-scale 3D data.<n>Our framework is fully integrated with PyTorch enabling interoperability with existing pipelines.
  arXiv  Detail & Related papers  (2024-07-01T20:20:33Z)
• DeformUX-Net: Exploring a 3D Foundation Backbone for Medical Image Segmentation with Depthwise Deformable Convolution [26.746489317083352]
  We introduce 3D DeformUX-Net, a pioneering volumetric CNN model. We revisit volumetric deformable convolution in depth-wise setting to adapt long-range dependency with computational efficiency. Our empirical evaluations reveal that the 3D DeformUX-Net consistently outperforms existing state-of-the-art ViTs and large kernel convolution models.
  arXiv  Detail & Related papers  (2023-09-30T00:33:41Z)
• Distributed Out-of-Memory NMF on CPU/GPU Architectures [1.0051474951635875]
  We propose an efficient out-of-memory implementation of the Non-negative Matrix Factorization (NMF) algorithm for HPC systems. Benchmark results show significant improvement of 32X to 76x speedup with the new implementation using GPU over the CPU-based NMFk.
  arXiv  Detail & Related papers  (2022-02-19T03:49:21Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.