Pl upgrades #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

harshagrawal13 wants to merge 128 commits into softnanolab:pl_upgrades from aqlaboratory:pl_upgrades

.github/workflows/docker-image.yml

-Original file line number
+Diff line change
@@ Expand Up / @@ -11,7 +11,7 @@ jobs: @@
         runs-on: ubuntu-latest
         steps:
         - uses: actions/checkout@v4
-        - name: Cleanup
+        - name: Cleanup  # https://github.com/actions/virtual-environments/issues/2840
           run: sudo rm -rf /usr/share/dotnet && sudo rm -rf /opt/ghc && sudo rm -rf "/usr/local/share/boost" && sudo rm -rf "$AGENT_TOOLSDIRECTORY"
         - name: Build the Docker image
           run: docker build . --file Dockerfile --tag openfold:$(date +%s)

docs/source/Aux_seq_files.md

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -68,9 +68,9 @@ All together, the file directory would look like:
  
    		  └── 6kwc.cif

    	  └── alignment_db

    		  ├── alignment_db_0.db 

              ├── alignment_db_1.db

              ...

              ├── alignment_db_9.db

                      ├── alignment_db_1.db

                       ...

                      ├── alignment_db_9.db

    		  └── alignment_db.index 

    ```

docs/source/Inference.md

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -42,7 +42,7 @@ $ bash scripts/download_openfold_params.sh $PARAMS_DIR
  
    We recommend selecting `openfold/resources` as the params directory as this is the default directory used by the `run_pretrained_openfold.py` to locate parameters. 

    If you choose to use a different directory, you may make a symlink to the `openfold/resources` directory, or specify an alternate parameter path with the command line argument `--jax_path` for AlphaFold parameters or `--openfold_checkpoint_path` for OpenFold parameters. 

    If you choose to use a different directory, you may make a symlink to the `openfold/resources` directory, or specify an alternate parameter path with the command line argument `--jax_param_path` for AlphaFold parameters or `--openfold_checkpoint_path` for OpenFold parameters. 

    ### Model Inference 

    @@ -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 \

    @@ -138,6 +138,7 @@ Some commonly used command line flags are here. A full list of flags can be view
  
    - `--data_random_seed`: Specifies a random seed to use.

    - `--save_outputs`: Saves a copy of all outputs from the model, e.g. the output of the msa track, ptm heads.

    - `--experiment_config_json`: Specify configuration settings using a json file. For example, passing a json with `{globals.relax.max_iterations = 10}` specifies 10 as the maximum number of relaxation iterations. See for  [`openfold/config.py`](https://github.com/aqlaboratory/openfold/blob/main/openfold/config.py#L283) the full dictionary of configuration settings. Any parameters that are not manually set in these configuration settings will refer to the defaults specified by your `config_preset`.

    - `--use_custom_template`: Uses all .cif files in `template_mmcif_dir` as template input. Make sure the chains of interest have the identifier _A_ and have the same length as the input sequence. The same templates will be read for all sequences that are passed for inference.

    ### Advanced Options for Increasing Efficiency

    @@ -159,12 +160,12 @@ Note that chunking (as defined in section 1.11.8 of the AlphaFold 2 supplement)
  
    #### Long sequence inference 

    To minimize memory usage during inference on long sequences, consider the following changes:

    - As noted in the AlphaFold-Multimer paper, the AlphaFold/OpenFold template stack is a major memory bottleneck for inference on long sequences. OpenFold supports two mutually exclusive inference modes to address this issue. One, `average_templates` in the `template` section of the config, is similar to the solution offered by AlphaFold-Multimer, which is simply to average individual template representations. Our version is modified slightly to accommodate weights trained using the standard template algorithm. Using said weights, we notice no significant difference in performance between our averaged template embeddings and the standard ones. The second, `offload_templates`, temporarily offloads individual template embeddings into CPU memory. The former is an approximation while the latter is slightly slower; both are memory-efficient and allow the model to utilize arbitrarily many templates across sequence lengths. Both are disabled by default, and it is up to the user to determine which best suits their needs, if either.

    - Inference-time low-memory attention (LMA) can be enabled in the model config. This setting trades off speed for vastly improved memory usage. By default, LMA is run with query and key chunk sizes of 1024 and 4096, respectively. These represent a favorable tradeoff in most memory-constrained cases. Powerusers can choose to tweak these settings in `openfold/model/primitives.py`. For more information on the LMA algorithm, see the aforementioned Staats & Rabe preprint.

    - Disable `tune_chunk_size` for long sequences. Past a certain point, it only wastes time.

    - As a last resort, consider enabling `offload_inference`. This enables more extensive CPU offloading at various bottlenecks throughout the model.

    - As noted in the AlphaFold-Multimer paper, the AlphaFold/OpenFold template stack is a major memory bottleneck for inference on long sequences. OpenFold supports two mutually exclusive inference modes to address this issue. One, `average_templates` in the `template` section of the config, is similar to the solution offered by AlphaFold-Multimer, which is simply to average individual template representations. Our version is modified slightly to accommodate weights trained using the standard template algorithm. Using said weights, we notice no significant difference in performance between our averaged template embeddings and the standard ones. The second, `offload_templates`, temporarily offloads individual template embeddings into CPU memory. The former is an approximation while the latter is slightly slower; both are memory-efficient and allow the model to utilize arbitrarily many templates across sequence lengths. Both are disabled by default, and it is up to the user to determine which best suits their needs, if either.

    - Inference-time low-memory attention (LMA) can be enabled in the model config. This setting trades off speed for vastly improved memory usage. By default, LMA is run with query and key chunk sizes of 1024 and 4096, respectively. These represent a favorable tradeoff in most memory-constrained cases. Powerusers can choose to tweak these settings in `openfold/model/primitives.py`. For more information on the LMA algorithm, see the aforementioned Staats & Rabe preprint.

    - Disable `tune_chunk_size` for long sequences. Past a certain point, it only wastes time.

    - As a last resort, consider enabling `offload_inference`. This enables more extensive CPU offloading at various bottlenecks throughout the model.

    - Disable FlashAttention, which seems unstable on long sequences.

    Using the most conservative settings, we were able to run inference on a 4600-residue complex with a single A100. Compared to AlphaFold's own memory offloading mode, ours is considerably faster; the same complex takes the more efficent AlphaFold-Multimer more than double the time. Use the `long_sequence_inference` config option to enable all of these interventions at once. The `run_pretrained_openfold.py` script can enable this config option with the `--long_sequence_inference` command line option

    Using the most conservative settings, we were able to run inference on a 4600-residue complex with a single A100. Compared to AlphaFold's own memory offloading mode, ours is considerably faster; the same complex takes the more efficent AlphaFold-Multimer more than double the time. Use the `long_sequence_inference` config option to enable all of these interventions at once. The `run_pretrained_openfold.py` script can enable this config option with the `--long_sequence_inference` command line option

    Input FASTA files containing multiple sequences are treated as complexes. In this case, the inference script runs AlphaFold-Gap, a hack proposed [here](https://twitter.com/minkbaek/status/1417538291709071362?lang=en), using the specified stock AlphaFold/OpenFold parameters (NOT AlphaFold-Multimer).
      
    Input FASTA files containing multiple sequences are treated as complexes. In this case, the inference script runs AlphaFold-Gap, a hack proposed [here](https://twitter.com/minkbaek/status/1417538291709071362?lang=en), using the specified stock AlphaFold/OpenFold parameters (NOT AlphaFold-Multimer).

docs/source/installation.md → docs/source/Installation.md

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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)

    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.

    @@ -19,10 +19,17 @@ At this time, only Linux systems are supported.
  
          Mamba is recommended as the dependencies required by OpenFold are quite large and mamba can speed up the process.

        - Activate the environment, e.g `conda activate openfold_env`

    1. Run the setup script to configure kernels and folding resources.

    	> scripts/install_third_party_dependencies.sh`

    3. Prepend the conda environment to the $LD_LIBRARY_PATH., e.g. 

    		`export $LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH``. You may optionally set this as a conda environment variable according to the [conda docs](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#saving-environment-variables) to activate each time the environment is used.

    4. Download parameters. We recommend using a destination as `openfold/resources` as our unittests will look for the weights there.

    	> scripts/install_third_party_dependencies.sh

    1. Prepend the conda environment to the `$LD_LIBRARY_PATH` and `$LIBRARY_PATH`., e.g. 

    	```

    	export LIBRARY_PATH=$CONDA_PREFIX/lib:$LIBRARY_PATH

    	export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH

    	```

    	You may optionally set this as a conda environment variable according to the [conda docs](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#saving-environment-variables) to activate each time the environment is used.

    1. Download parameters. We recommend using a destination as `openfold/resources` as our unittests will look for the weights there.

    	-  For AlphaFold2 weights, use 

    		> ./scripts/download_alphafold_params.sh <dest>

    	 - For OpenFold weights, use : 

    @@ -46,10 +53,9 @@ 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 URL in step 1 with https://github.com/aqlaboratory/openfold/tree/pl_upgrades

    	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`. 

    ### Install OpenFold parameters without aws

    If you don't have access to `aws` on your system, you can use a different download source:

    @@ -59,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).

docs/source/Multimer_Inference.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -72,8 +72,7 @@ python3 run_pretrained_openfold.py \ @@
         --output_dir ./
     ```
-    Note that template searching in the multimer pipeline
-    uses HMMSearch with the PDB SeqRes database, replacing HHSearch and PDB70 used in the monomer pipeline.
-    As with monomer inference, if you've already computed alignments for the query, you can use
-    the `--use_precomputed_alignments` option.
+    **Notes:**
+    - Template searching in the multimer pipeline uses HMMSearch with the PDB SeqRes database, replacing HHSearch and PDB70 used in the monomer pipeline.
+    - As with monomer inference, if you've already computed alignments for the query, you can use the `--use_precomputed_alignments` option.
+    - At this time, only AlphaFold parameter weights are available for multimer mode.

docs/source/index.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -5,7 +5,7 @@ @@
     :align: center
     :alt: Comparison of OpenFold and AlphaFold2 predictions to the experimental structure of PDB 7KDX, chain B._
     ```
-    Welcome to the Documentation for OpenFold, the fully open source, trainable, PyTorch-based reproduction of DeepMind's
+    Welcome to the Documentation for [OpenFold](https://github.com/aqlaboratory/openfold), the fully open source, trainable, PyTorch-based reproduction of DeepMind's
     [AlphaFold 2](https://github.com/deepmind/alphafold).
     Here, you will find guides and documentation for:
@@ Expand Down Expand Up / @@ -115,4 +115,4 @@ Aux_seq_files.md @@
     OpenFold_Parameters.md
     FAQ.md
     original_readme.md
-    ```
+    ```

environment.yml

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -8,34 +8,33 @@ dependencies:
  
      - cuda

      - gcc=12.4

      - python=3.10

      - libgcc=7.2

      - setuptools=59.5.0

      - pip

      - openmm=7.7

      - openmm

      - pdbfixer

      - pytorch-lightning

      - biopython

      - numpy<2.0.0

      - numpy

      - pandas

      - PyYAML==5.4.1

      - PyYAML

      - requests

      - scipy

      - tqdm==4.62.2

      - tqdm

      - typing-extensions

      - wandb

      - modelcif==0.7

      - awscli

      - ml-collections

      - mkl=2022.1

      - aria2

      - mkl

      - git

      - bioconda::hmmer

      - bioconda::hhsuite

      - bioconda::kalign2

      - pytorch::pytorch=2.1

      - pytorch::pytorch-cuda=12.1

      - 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

          - flash-attn

notebooks/OpenFold.ipynb

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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.79\")\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 = \"a96ffd67f8c96f8c4decc3abdd2cffbb57fc5764\"\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",

    @@ -907,4 +916,4 @@
  
      },

      "nbformat": 4,

      "nbformat_minor": 0

    }
      
    }

openfold/config.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -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": [],
@@ Expand Down @@

openfold/data/data_pipeline.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -23,8 +23,19 @@ @@
     from typing import Mapping, Optional, Sequence, Any, MutableMapping, Union
     import numpy as np
     import torch
-    from openfold.data import templates, parsers, mmcif_parsing, msa_identifiers, msa_pairing, feature_processing_multimer
-    from openfold.data.templates import get_custom_template_features, empty_template_feats
+    from openfold.data import (
+        templates,
+        parsers,
+        mmcif_parsing,
+        msa_identifiers,
+        msa_pairing,
+        feature_processing_multimer,
+    )
+    from openfold.data.templates import (
+        get_custom_template_features,
+        empty_template_feats,
+        CustomHitFeaturizer,
+    )
     from openfold.data.tools import jackhmmer, hhblits, hhsearch, hmmsearch
     from openfold.np import residue_constants, protein
@@ Expand All / @@ -38,7 +49,9 @@ def make_template_features( @@
         template_featurizer: Any,
     ) -> FeatureDict:
         hits_cat = sum(hits.values(), [])
-        if(len(hits_cat) == 0 or template_featurizer is None):
+        if template_featurizer is None or (
+            len(hits_cat) == 0 and not isinstance(template_featurizer, CustomHitFeaturizer)
+        ):
             template_features = empty_template_feats(len(input_sequence))
         else:
             templates_result = template_featurizer.get_templates(
@@ Expand Down @@

openfold/data/mmcif_parsing.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -283,7 +283,7 @@ def parse( @@
                 author_chain = mmcif_to_author_chain_id[chain_id]
                 seq = []
                 for monomer in seq_info:
-                    code = PDBData.protein_letters_3to1.get(monomer.id, "X")
+                    code = PDBData.protein_letters_3to1_extended.get(monomer.id, "X")
                     seq.append(code if len(code) == 1 else "X")
                 seq = "".join(seq)
                 author_chain_to_sequence[author_chain] = seq
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Pl upgrades #1

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Pl upgrades #1

Are you sure you want to change the base?

Uh oh!

Pl upgrades #1

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!