Merge pull request #534 from aqlaboratory/pl_upgrades

Update openfold to use pytorch 2 and other updated dependencies

Author: jnwei

Dockerfile

@@ -1,17 +1,20 @@
-FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 
 # metainformation
-LABEL org.opencontainers.image.version = "1.0.0"
-LABEL org.opencontainers.image.authors = "Gustaf Ahdritz"
+LABEL org.opencontainers.image.version = "2.0.0"
+LABEL org.opencontainers.image.authors = "OpenFold Team"
 LABEL org.opencontainers.image.source = "https://github.com/aqlaboratory/openfold"
 LABEL org.opencontainers.image.licenses = "Apache License 2.0"
-LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:10.2-cudnn8-runtime-ubuntu18.04"
+LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:12.4.1-devel-ubuntu22.04"
+
+RUN apt-get update && apt-get install -y wget
 
 RUN apt-key del 7fa2af80
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
+RUN dpkg -i cuda-keyring_1.0-1_all.deb
+
+RUN apt-get install -y libxml2 cuda-minimal-build-12-1 libcusparse-dev-12-1 libcublas-dev-12-1 libcusolver-dev-12-1 git
 
-RUN apt-get update && apt-get install -y wget libxml2 cuda-minimal-build-11-3 libcusparse-dev-11-3 libcublas-dev-11-3 libcusolver-dev-11-3 git
 RUN wget -P /tmp \
     "https://github.com/conda-forge/miniforge/releases/download/23.3.1-1/Miniforge3-Linux-x86_64.sh" \
     && bash /tmp/Miniforge3-Linux-x86_64.sh -b -p /opt/conda \

docs/source/Inference.md

@@ -62,7 +62,7 @@ python3 run_pretrained_openfold.py \
     $TEMPLATE_MMCIF_DIR 
     --output_dir $OUTPUT_DIR \
     --config_preset model_1_ptm \
-    --uniref90_database_path $BASE_DATA_DIR/uniref90 \
+    --uniref90_database_path $BASE_DATA_DIR/uniref90/uniref90.fasta \
     --mgnify_database_path $BASE_DATA_DIR/mgnify/mgy_clusters_2018_12.fa \
     --pdb70_database_path $BASE_DATA_DIR/pdb70 \
     --uniclust30_database_path $BASE_DATA_DIR/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \

docs/source/Installation.md

@@ -4,7 +4,7 @@ In this guide, we will OpenFold and its dependencies.
 
 **Pre-requisites**
 
-This package is currently supported for CUDA 11 and Pytorch 1.12. All dependencies are listed in the [`environment.yml`](https://github.com/aqlaboratory/openfold/blob/main/environment.yml). To install OpenFold for CUDA 12, please refer to the [Environment specific modifications](#Environment-specific-modifications) section.
+This package is currently supported for CUDA 12 and Pytorch 2. All dependencies are listed in the [`environment.yml`](https://github.com/aqlaboratory/openfold/blob/main/environment.yml).
 
 At this time, only Linux systems are supported.
 
@@ -53,12 +53,6 @@ Certain tests perform equivalence comparisons with the AlphaFold implementation.
 
 ## Environment specific modifications 
 
-### CUDA 12
-To use OpenFold on CUDA 12 environment rather than a CUDA 11 environment.
-	In step 1, use the branch [`pl_upgrades`](https://github.com/aqlaboratory/openfold/tree/pl_upgrades) rather than the main branch, i.e. replace the command in step 1 with `git clone -b pl_upgrades https://github.com/aqlaboratory/openfold.git`
-	and follow the rest of the steps of [Installation Guide](#Installation)
-
-
 ### MPI
 To use OpenFold with MPI support, you will need to add the package [`mpi4py`](https://pypi.org/project/mpi4py/). This can be done with pip in your OpenFold environment, e.g. `$ pip install mpi4py`. 
 
@@ -71,4 +65,4 @@ If you don't have access to `aws` on your system, you can use a different downlo
 
 ### Docker setup
 
-A [`Dockerfile`] is provided to build an OpenFold Docker image. Additional notes for setting up a docker container for OpenFold and running inference can be found [here](original_readme.md#building-and-using-the-docker-container).
+A [`Dockerfile`](https://github.com/aqlaboratory/openfold/blob/main/Dockerfile) is provided to build an OpenFold Docker image. Additional notes for setting up a docker container for OpenFold and running inference can be found [here](original_readme.md#building-and-using-the-docker-container).

environment.yml

@@ -3,36 +3,38 @@ channels:
   - conda-forge
   - bioconda
   - pytorch
+  - nvidia 
 dependencies:
-  - python=3.9
-  - libgcc=7.2
+  - cuda
+  - gcc=12.4
+  - python=3.10
   - setuptools=59.5.0
   - pip
-  - openmm=7.7
+  - openmm
   - pdbfixer
   - pytorch-lightning
   - biopython
   - numpy
   - pandas
-  - PyYAML==5.4.1
+  - PyYAML
   - requests
-  - scipy==1.7
-  - tqdm==4.62.2
-  - typing-extensions==4.0
+  - scipy
+  - tqdm
+  - typing-extensions
   - wandb
   - modelcif==0.7
   - awscli
   - ml-collections
   - aria2
-  - mkl=2024.0
+  - mkl
   - git
-  - bioconda::hmmer==3.3.2
-  - bioconda::hhsuite==3.3.0
-  - bioconda::kalign2==2.04
-  - bioconda::mmseqs2
-  - pytorch::pytorch=1.12.*
+  - bioconda::hmmer
+  - bioconda::hhsuite
+  - bioconda::kalign2
+  - pytorch::pytorch=2.5
+  - pytorch::pytorch-cuda=12.4
   - pip:
-      - deepspeed==0.12.4
+      - deepspeed==0.14.5
       - dm-tree==0.1.6
       - git+https://github.com/NVIDIA/dllogger.git
-      - git+https://github.com/Dao-AILab/flash-attention.git@5b838a8
+      - flash-attn 

notebooks/OpenFold.ipynb

@@ -1,5 +1,15 @@
 {
   "cells": [
+	{
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/aqlaboratory/OpenFold/blob/main/notebooks/OpenFold.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -107,11 +117,11 @@
         "\n",
         "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
         "\n",
-        "\n",
-        "os.system(\"wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh\")\n",
-        "os.system(\"bash Mambaforge-Linux-x86_64.sh -bfp /usr/local\")\n",
+        "os.system(\"wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh\")\n",
+        "os.system(\"bash Miniforge3-Linux-x86_64.sh -bfp /usr/local\")\n",
+        "os.environ[\"PATH\"] = \"/usr/local/bin:\" + os.environ[\"PATH\"]\n",
         "os.system(\"mamba config --set auto_update_conda false\")\n",
-        "os.system(f\"mamba install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0 openmm=7.7.0 python={python_version} pdbfixer biopython=1.83\")\n",
+        "os.system(f\"mamba install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0 openmm=8.2.0 python={python_version} pdbfixer biopython=1.83\")\n",
         "os.system(\"pip install -q torch ml_collections py3Dmol modelcif\")\n",
         "\n",
         "try:\n",
@@ -127,7 +137,7 @@
         "\n",
         "    %shell mkdir -p /content/openfold/openfold/resources\n",
         "\n",
-        "    commit = \"3bec3e9b2d1e8bdb83887899102eff7d42dc2ba9\"\n",
+        "    commit = \"1ffd197489aa5f35a5fbce1f00d7dd49bce1bd2f\"\n",
         "    os.system(f\"pip install -q git+https://github.com/aqlaboratory/openfold.git@{commit}\")\n",
         "\n",
         "    os.system(f\"cp -f -p /content/stereo_chemical_props.txt /usr/local/lib/python{python_version}/site-packages/openfold/resources/\")\n",
@@ -893,8 +903,7 @@
   "metadata": {
     "colab": {
       "provenance": [],
-      "gpuType": "T4",
-      "toc_visible": true
+      "gpuType": "T4"
     },
     "kernelspec": {
       "display_name": "Python 3",

openfold/config.py

@@ -660,7 +660,7 @@ def model_config(
         },
         "relax": {
             "max_iterations": 0,  # no max
-            "tolerance": 2.39,
+            "tolerance": 10.0,
             "stiffness": 10.0,
             "max_outer_iterations": 20,
             "exclude_residues": [],

openfold/model/primitives.py

@@ -28,7 +28,7 @@
 fa_is_installed = importlib.util.find_spec("flash_attn") is not None
 if fa_is_installed:
     from flash_attn.bert_padding import unpad_input
-    from flash_attn.flash_attn_interface import flash_attn_unpadded_kvpacked_func
+    from flash_attn.flash_attn_interface import flash_attn_varlen_kvpacked_func
 
 import torch
 import torch.nn as nn
@@ -808,10 +808,10 @@ def _flash_attn(q, k, v, kv_mask):
     # [B_flat, N, 2 * H * C]
     kv = kv.reshape(*kv.shape[:-3], -1) 
 
-    kv_unpad, _, kv_cu_seqlens, kv_max_s = unpad_input(kv, kv_mask)
+    kv_unpad, _, kv_cu_seqlens, kv_max_s, _ = unpad_input(kv, kv_mask)
     kv_unpad = kv_unpad.reshape(-1, *kv_shape[-3:])
 
-    out = flash_attn_unpadded_kvpacked_func(
+    out = flash_attn_varlen_kvpacked_func(
         q,
         kv_unpad,
         q_cu_seqlens,

openfold/np/relax/amber_minimize.py

@@ -34,6 +34,7 @@
 from openmm.app.internal.pdbstructure import PdbStructure
 
 ENERGY = unit.kilocalories_per_mole
+FORCE = unit.kilojoules_per_mole / unit.nanometer 
 LENGTH = unit.angstroms
 
 
@@ -439,7 +440,7 @@ def _run_one_iteration(
     exclude_residues = exclude_residues or []
 
     # Assign physical dimensions.
-    tolerance = tolerance * ENERGY
+    tolerance = tolerance * FORCE
     stiffness = stiffness * ENERGY / (LENGTH ** 2)
 
     start = time.perf_counter()

openfold/utils/superimposition.py

@@ -35,10 +35,10 @@ def _superimpose_np(reference, coords):
 
 
 def _superimpose_single(reference, coords):
-    reference_np = reference.detach().to(torch.float).cpu().numpy()    
-    coords_np = coords.detach().to(torch.float).cpu().numpy()
-    superimposed, rmsd = _superimpose_np(reference_np, coords_np)
-    return coords.new_tensor(superimposed), coords.new_tensor(rmsd)
+	reference_np = reference.detach().to(torch.float).cpu().numpy()  
+	coords_np = coords.detach().to(torch.float).cpu().numpy()
+	superimposed, rmsd = _superimpose_np(reference_np, coords_np)
+	return coords.new_tensor(superimposed), coords.new_tensor(rmsd)
 
 
 def superimpose(reference, coords, mask):

scripts/install_third_party_dependencies.sh

@@ -14,7 +14,7 @@ gunzip -c tests/test_data/sample_feats.pickle.gz > tests/test_data/sample_feats.
 python setup.py install
 
 echo "Download CUTLASS, required for Deepspeed Evoformer attention kernel"
-git clone https://github.com/NVIDIA/cutlass --depth 1
+git clone https://github.com/NVIDIA/cutlass --branch v3.6.0 --depth 1
 conda env config vars set CUTLASS_PATH=$PWD/cutlass
 
 # This setting is used to fix a worker assignment issue during data loading

setup.py

@@ -29,7 +29,7 @@
 ]
 
 extra_cuda_flags = [
-    '-std=c++14',
+    '-std=c++17',
     '-maxrregcount=50',
     '-U__CUDA_NO_HALF_OPERATORS__',
     '-U__CUDA_NO_HALF_CONVERSIONS__',
@@ -52,9 +52,9 @@ def get_cuda_bare_metal_version(cuda_dir):
         return raw_output, bare_metal_major, bare_metal_minor
 
 compute_capabilities = set([
-    (3, 7), # K80, e.g.
     (5, 2), # Titan X
     (6, 1), # GeForce 1000-series
+    (9, 0), # Hopper
 ])
 
 compute_capabilities.add((7, 0))
@@ -113,7 +113,7 @@ def get_cuda_bare_metal_version(cuda_dir):
 
 setup(
     name='openfold',
-    version='2.0.0',
+    version='2.2.0',
     description='A PyTorch reimplementation of DeepMind\'s AlphaFold 2',
     author='OpenFold Team',
     author_email='jennifer.wei@omsf.io',
@@ -130,7 +130,7 @@ def get_cuda_bare_metal_version(cuda_dir):
     classifiers=[
         'License :: OSI Approved :: Apache Software License',
         'Operating System :: POSIX :: Linux',
-        'Programming Language :: Python :: 3.9,'
+        'Programming Language :: Python :: 3.10,'
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
     ],
 )

tests/test_deepspeed_evo_attention.py

@@ -306,7 +306,6 @@ def test_compare_model(self):
         batch["residx_atom37_to_atom14"] = batch[
             "residx_atom37_to_atom14"
         ].long()
-        # print(batch["target_feat"].shape)
         batch["target_feat"] = torch.nn.functional.one_hot(batch["aatype"], consts.msa_logits - 1).to(torch.float32)
         batch["template_all_atom_mask"] = batch["template_all_atom_masks"]
         batch.update(
@@ -316,8 +315,9 @@ def test_compare_model(self):
         # Move the recycling dimension to the end
         move_dim = lambda t: t.permute(*range(len(t.shape))[1:], 0)
         batch = tensor_tree_map(move_dim, batch)
-        with torch.no_grad():
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        # Restrict this test to use only torch.float32 precision due to instability with torch.bfloat16
+        # https://github.com/aqlaboratory/openfold/issues/532
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float32):
                 model = compare_utils.get_global_pretrained_openfold()
                 model.globals.use_deepspeed_evo_attention = False
                 out_repro = model(batch)

tests/test_model.py

@@ -202,4 +202,4 @@ def run_alphafold(batch):
         out_repro = out_repro["sm"]["positions"][-1]
         out_repro = out_repro.squeeze(0)
 
-        self.assertTrue(torch.max(torch.abs(out_gt - out_repro)) < 1e-3)
+        compare_utils.assert_mean_abs_diff_small(out_gt, out_repro, 1e-3)