Fold-CP: A Context Parallelism Framework for Biomolecular Modeling

.gitignore

@@ -1,3 +1,6 @@
+# Project specific
+tests/test_data
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -163,4 +166,4 @@ cython_debug/
 
 # Boltz prediction outputs
 #   All result files generated from a boltz prediction call
-boltz_results_*/
+boltz_results_*/

.pre-commit-config.yaml

@@ -0,0 +1,31 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-yaml
+        exclude: "mkdocs.yml"
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.1
+    hooks:
+      - id: ruff
+        # Don't fail on docstring checks in pre-commit, and don't remove these unused noqa flags.
+        args: ["--fix", "--ignore", "D", "--ignore", "RUF100"]
+      - id: ruff-format
+  - repo: https://github.com/Yelp/detect-secrets
+    rev: v1.5.0
+    hooks:
+      - id: detect-secrets
+        name: detect-secrets (everything but notebooks)
+        args: ['--baseline', './pre-commit/.secrets.baseline', '--exclude-files', '(.*\.ipynb|.*\.baseline|.*\.a3m)$', ]
+        exclude: package.lock.json
+  - repo: local
+    hooks:
+      - id: license-header-check
+        name: Run license-check script
+        entry: python pre-commit/license_check.py --license-header ./pre-commit/license_header --modify
+        language: python
+        additional_dependencies: ["click==8.1.7"]
+        pass_filenames: false
+        always_run: true

CHANGELOG.md

@@ -0,0 +1,13 @@
+# Boltz Context Parallelism Changelog
+
+## 0.1.0 (initial release)
+
+### New Features
+
+- Distributed inference with DTensor context parallelism (`src/boltz/distributed/predict.py`)
+- Distributed training with DTensor context parallelism (`src/boltz/distributed/train.py`)
+- 2D CP mesh support (Shard x Shard context parallelism)
+- Data parallelism combined with context parallelism (DP x CP)
+- Multiple attention kernel backends: cuEquivariance, trifast, FlexAttention
+- CUDA memory profiling support
+- Preprocessed input format for distributed inference

CONTRIBUTING.md

@@ -0,0 +1 @@
+This project is currently not accepting contributions.

Dockerfile

@@ -0,0 +1,60 @@
+# Boltz-1 Dockerfile
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:25.10-py3
+
+FROM ${BASE_IMAGE} AS boltz-base
+
+# Install core apt packages.
+RUN --mount=type=cache,id=apt-cache,target=/var/cache/apt,sharing=locked \
+  --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
+  bash -c '\
+  apt-get update -qy && \
+  apt-get install -qyy \
+  libsndfile1 \
+  ffmpeg \
+  git \
+  curl \
+  pre-commit \
+  lsof \
+  git-lfs \
+  sudo && \
+  apt-get upgrade -qyy \
+  rsync && \
+  rm -rf /tmp/* /var/tmp/*'
+
+RUN apt-get install -y gnupg
+
+RUN mkdir -p /workspace/boltz/
+
+# Fix for duplicate triton installation due to pytorch_triton renaming
+RUN bash -c ' \
+  cd /usr/local/lib/python3*/dist-packages/ && \
+  PTRITON_DIR=$(ls -d pytorch_triton-*.dist-info 2>/dev/null | head -n 1) && \
+  if [ -n "$PTRITON_DIR" ]; then \
+    VERSION=${PTRITON_DIR#pytorch_triton-} && \
+    VERSION=${VERSION%.dist-info} && \
+    NEW_DIR="triton-${VERSION}.dist-info" && \
+    cp -r "$PTRITON_DIR" "$NEW_DIR" && \
+    sed -i "s/Name: pytorch-triton/Name: triton/" "$NEW_DIR/METADATA" && \
+    echo "Successfully aliased pytorch-triton to triton version $VERSION"; \
+  fi'
+
+ENV NVIDIA_TF32_OVERRIDE=0
+
+RUN bash -c 'echo "ubuntu ALL=(root) NOPASSWD:ALL" > /etc/sudoers.d/ubuntu && \
+    chmod 0440 /etc/sudoers.d/ubuntu'
+
+FROM boltz-base AS dev
+
+# Install boltz-1
+COPY ./README.md /workspace/boltz/README.md
+COPY ./pyproject.toml /workspace/boltz/pyproject.toml
+COPY ./src /workspace/boltz/src
+COPY ./tests /workspace/boltz/tests
+COPY ./scripts /workspace/boltz/scripts
+COPY ./examples /workspace/boltz/examples
+
+WORKDIR /workspace/boltz
+RUN bash -c 'find . -name __pycache__ -type d -print | xargs rm -rf'
+RUN pip install --no-build-isolation --editable .[lint,test,cuda,dev]
+
+ENV NVIDIA_TF32_OVERRIDE=0

LICENSE

@@ -0,0 +1 @@
+licenses/LICENSE

README.md

@@ -18,6 +18,18 @@ Boltz is a family of models for biomolecular interaction prediction. Boltz-1 was
 
 All the code and weights are provided under MIT license, making them freely available for both academic and commercial uses. For more information about the model, see the [Boltz-1](https://doi.org/10.1101/2024.11.19.624167) and [Boltz-2](https://doi.org/10.1101/2025.06.14.659707) technical reports. To discuss updates, tools and applications join our [Slack channel](https://boltz.bio/join-slack).
 
+## Fold-CP: A Context Parallelism Framework for Biomolecular Modeling
+
+This repo also contains context parallelism (CP) for distributed inference and training for biomolecular folding models across multiple GPUs using a 2D CP mesh combined with data parallelism, demonstrated with the Boltz model. See [this README](src/boltz/distributed/README.md) for detail.
+
+### Copyright and License Compliance
+
+- The context parallel code is licensed under the terms and conditions as written in [the license file](licenses/LICENSE)
+
+- The original Boltz code is licensed under their respective MIT license (See the [third-party-attr.txt](licenses/third-party-attr.txt))
+
+- This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use
+
 ## Installation
 
 > Note: we recommend installing boltz in a fresh python environment

SECURITY.md

@@ -0,0 +1,24 @@
+## Security
+
+NVIDIA is dedicated to the security and trust of our software products and services, including all source code repositories managed through our organization.
+
+If you need to report a security issue, please use the appropriate contact points outlined below. **Please do not report security vulnerabilities through GitHub.** If a potential security issue is inadvertently reported via a public issue or pull request, NVIDIA maintainers may limit public discussion and redirect the reporter to the appropriate private disclosure channels.
+
+## Reporting Potential Security Vulnerability in an NVIDIA Product
+
+To report a potential security vulnerability in any NVIDIA product:
+- Web: [Security Vulnerability Submission Form](https://www.nvidia.com/object/submit-security-vulnerability.html)
+- E-Mail: psirt@nvidia.com
+    - We encourage you to use the following PGP key for secure email communication: [NVIDIA public PGP Key for communication](https://www.nvidia.com/en-us/security/pgp-key)
+    - Please include the following information:
+   	 - Product/Driver name and version/branch that contains the vulnerability
+     - Type of vulnerability (code execution, denial of service, buffer overflow, etc.)
+   	 - Instructions to reproduce the vulnerability
+   	 - Proof-of-concept or exploit code
+   	 - Potential impact of the vulnerability, including how an attacker could exploit the vulnerability
+
+While NVIDIA currently does not have a bug bounty program, we do offer acknowledgement when an externally reported security issue is addressed under our coordinated vulnerability disclosure policy. Please visit our [Product Security Incident Response Team (PSIRT)](https://www.nvidia.com/en-us/security/psirt-policies/) policies page for more information.
+
+## NVIDIA Product Security
+
+For all security-related concerns, please visit NVIDIA's Product Security portal at https://www.nvidia.com/en-us/security