For decades, computational biology has operated under a reductionist compromise. To fit complex biological systems into the limited memory of a single GPU, researchers have had to deconstruct them into isolated fragments—single proteins or small domains. This created a context gap, where larger proteins or complexes could not be folded zero-shot due to GPU hardware memory constraints.
Now, a new context parallelism (CP) framework from the NVIDIA BioNeMo team is shattering the memory barriers of structural biology, enabling the holistic modeling of systems.
This post explains how to achieve CP in biomolecular architectures that diverge from standard Transformers. If you’re a structural biologist, computational chemist, or machine learning engineer seeking to model massive biomolecular complexes without sacrificing global context, read on.
To use the solution outlined in this post, you’ll need:
- Familiarity with geometric deep learning foundation models like AlphaFold3 or Boltz-2.
- Understanding of PyTorch Distributed (DTensor) operations and custom autograd functions.
- Access to an NVIDIA H100 or B200 GPU cluster, as the framework relies heavily on its interconnect bandwidth and Transformer Engine acceleration for exascale tasks.
For more details, see Fold-CP: A Context Parallelism Framework for Biomolecular Modeling.
Sharding a single large molecular system across multiple GPUs
In the absence of CP, folding large complexes (typically exceeding 1,000–3,000 residues) requires a reductionist approach where the system is physically or computationally deconstructed into manageable chunks. These methods enable researchers to stay within the strict VRAM limits of single GPUs, but often sacrifice global structural accuracy.
The most common workaround for massive proteins is to slice the sequence into smaller, overlapping segments. Fragments must overlap significantly to ensure that local secondary structures are consistent across the split points. This method destroys long-range information. For example, researchers cannot model allostery or signal transduction across the entire complex.
The other common workaround is chunking, which, unlike physical sequence fragmentation, occurs within the model architecture to save VRAM during inference. Models like Boltz use aggressive chunking to process large matrices in smaller tiles. Other techniques such as FastFold employ autochunking to dynamically adjust the chunking strategy and improve peak memory usage. To learn more, see FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours.
All these techniques inherently suffer from a lack of global context, especially during training. The NVIDIA BioNeMo CP framework overcomes these limits by sharding a single large molecular system across multiple GPUs. Unlike traditional data parallelism, which assigns each GPU a different protein to fold, CP splits a single massive sample across GPUs.
BioNeMo context parallelism implementation
The NVIDIA BioNeMo CP implementation is built on Torch distributed APIs for GPU-to-GPU communications. The architecture is built from the bottom up, starting with low-level communication protocols and moving up to high-level model-specific workflows. This post uses Boltz as the example codebase.
To achieve linear capacity scaling—where the capability of the system grows linearly with the number of GPUs—the framework implements a multidimensional sharding strategy. This ensures that no single device holds the full global state of the biomolecule, which would defeat the memory objectives of CP. Custom lego-like pieces are built for different modules, to make the implementation efficient and easier to port to other architectures.
2D tiling of the pair representation
The framework partitions the global (N x N) matrix into a grid of blocks. For a 10,000-residue complex, which represents 100 million interactions, each GPU manages only a specific sub-block. This localizes the memory footprint from O(N2) to O(N2/P) per device.
Overlapping computation and communication
The framework implements distributed primitives to orchestrate local computation with asynchronous peer-to-peer transfers. While a GPU is computing a local update, it is simultaneously sending and receiving data to and from its neighbors in the row and column rings. As the biological problem size grows, the ratio of computation to communication improves, making the system more efficient at larger scales.
Efficient atom sequence local attention
Alphafold3 sequence local attention limits the atom attention to local attention windows of size 32 by 128, where the stack of such windows go through the attention computation in a batch. The aforementioned tiling of atom features needs to be repartitioned for the distributed version of this window-batching-based atom attention. The NVIDIA team implemented halo-exchange-based distributed primitives to partition the atom features, so that subsequent window-batch attention requires no inter-GPU communication.
Context parallelism implementation for triangle multiplication
The following example shows how to make a CP-aware layer that computes the triangle multiplication in a distributed way.
# … 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()
This code initializes a CP environment by using a DistributedManager to set up a square 2D device mesh—a specific architectural requirement that ensures row-wise and column-wise communication patterns remain symmetrical and efficient. The Ring2DComm handle is then instantiated to manage specialized peer-to-peer communication, which allows for the circulation of data blocks in a continuous loop. This ring approach is critical because it enables the overlapping of local computation with data transfers, ensuring that the \(O(N^2)\) pair representation tensors never exceed the memory capacity of a single GPU.
The second half of the script handles the transition from a standard serial model to DTensor-based one. A standard layer, such as TriangleMultiplicationOutgoing, is loaded on the CPU before being wrapped by DistributedTriangleMultiplication, which implements a distributed version of the algorithm adapted for the NVIDIA BioNeMo CP framework. By processing inputs as distributed tensors (DTensors), the model ensures that the large activation tensors are sharded across the grid.
Unlocking token scaling for structural biology
Figure 2 shows that token capacity scaling laws are now unlocked for biomolecular architectures with the introduction of CP. Boltz predictions can be run on up to ~20,000 tokens using 256 GPUs and can scale the maximum token length on NVIDIA H100 GPUs, with accelerated scaling on NVIDIA B300 GPUs.
Without any additional training or fine-tuning with longer crop lengths, the team folded a TTC7A/PI4KA/FAM126A/EFR3A(700–823) system that contains 3,605 residues across four chains—far exceeding the Boltz-2 training crop size of 768 residues and the memory capacity of a single GPU. Using CP enabled the generation of five structural samples in under five minutes (∼54 seconds per sample), on four NVIDIA H100 GPUs—while maintaining all long-range inter-subunit contacts within the model context window.
In parallel, the team is also working to push this next-generation frontier alongside collaborators like Rezo Therapeutics, Proxima, and Earendil Labs who have deeply contributed to the CP framework development.
Rezo Therapeutics integrated the CP framework to predict massive protein-protein interactions (PPIs) spanning up to 6,500 residues, enabling structural prediction for the vast majority of known protein complexes and unlocking the rapid discovery of novel complexes. In fact, there is greater than 3x enrichment of CP-resolved, high-quality novel protein complexes when compared to predictions made using only high-confidence PPIs in the public domain.
Proxima embedded CP within their all-atom generative foundation model, Neo, enabling inference on assemblies up to 4,000 tokens to structurally resolve therapeutically relevant interactions across the proteome mapped by their unique mass spectrometry platform, thereby helping their drug discovery efforts in developing molecular glues and other proximity-based therapeutics.
Earendil Labs integrated the framework into their proprietary biomolecular foundation model, successfully extending the input sequence lengths to model complex, multi-protein systems that were previously computationally prohibitive. Earendil Labs demonstrated that CP has the potential to maintain high-fidelity structural predictions even as sequence complexity scales and also shrink discovery timelines for next-generation biotherapeutics.
Get started with context parallelism for biomolecular modeling
While initial proof-of-concepts demonstrate that CP shatters previous memory barriers to enable modeling structures of an unlimited size, physical capacity alone does not guarantee biological accuracy. Current models often struggle to perform high-fidelity folding at scale because they were trained on small fragments. Fine-tuning with larger crop sizes is essential to accurately capture the emergent logic of long-range interactions.
As part of this effort, one of the primary bottlenecks in this field, data scarcity, is being addressed through contributions to the AlphaFold Protein Structure Database. By leveraging NVIDIA accelerated computing software such as NVIDIA cuEquivariance and NVIDIA TensorRT to populate AFCDB with high-throughput predictions of massive homomeric and heteromeric complexes, this is the groundwork to build the synthetic complex data that may be valuable to train foundation models that represent larger systems in biology.
To learn more, see the Boltz CP code open-source documentation and check out Fold-CP: A Context Parallelism Framework for Biomolecular Modeling.
