diff --git a/docker/Dockerfile b/docker/Dockerfile index 81e54363..06346b18 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,13 +1,46 @@ -FROM hpcaitech/pytorch-cuda:1.12.0-11.3.0 +FROM image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.3.0-12.1.0 -RUN conda install openmm=7.7.0 pdbfixer -c conda-forge -y \ - && conda install hmmer==3.3.2 hhsuite=3.3.0 kalign2=2.04 -c bioconda -y -RUN pip install biopython==1.79 dm-tree==0.1.6 ml-collections==0.1.0 \ -scipy==1.7.1 ray pyarrow pandas einops +RUN apt-get update && \ + apt-get install -y gcc-9 g++-9 && \ + update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 100 \ + --slave /usr/bin/g++ g++ /usr/bin/g++-9 && \ + rm -rf /var/lib/apt/lists/* -RUN pip install colossalai +ENV CC=/usr/bin/gcc +ENV CXX=/usr/bin/g++ -Run git clone https://github.com/hpcaitech/FastFold.git \ - && cd ./FastFold \ - && python setup.py install +RUN conda update -n base -c defaults -y conda openssl ca-certificates && \ + conda config --remove-key channels || true && \ + conda config --add channels conda-forge && \ + conda config --add channels bioconda && \ + conda config --add channels defaults && \ + conda config --set channel_priority strict && \ + conda config --set show_channel_urls yes + +RUN conda install -n pytorch -y -c conda-forge -c bioconda \ + openmm=8.2.0 \ + pdbfixer \ + hmmer=3.3.2 \ + hhsuite=3.3.0 \ + scipy==1.11.4 \ + kalign2=2.04 && \ + conda clean -afy + +RUN /opt/conda/envs/pytorch/bin/pip install --no-cache-dir \ + biopython==1.79 \ + dm-tree==0.1.6 \ + ml-collections==0.1.0 \ + ray==2.47.1 \ + pyarrow \ + pandas \ + einops \ + colossalai==0.5.0 + + +RUN git clone --depth 1 https://github.com/hpcaitech/FastFold.git && \ + cd FastFold && \ + /opt/conda/envs/pytorch/bin/python setup.py install && + + +WORKDIR /workspace \ No newline at end of file diff --git a/fastfold/common/residue_constants.py b/fastfold/common/residue_constants.py index 01d98991..963957b2 100644 --- a/fastfold/common/residue_constants.py +++ b/fastfold/common/residue_constants.py @@ -1127,10 +1127,10 @@ def _make_rigid_transformation_4x4(ex, ey, translation): # and an array with (restype, atomtype, coord) for the atom positions # and compute affine transformation matrices (4,4) from one rigid group to the # previous group -restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=np.int) +restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int) restype_atom37_mask = np.zeros([21, 37], dtype=np.float32) restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32) -restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=np.int) +restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int) restype_atom14_mask = np.zeros([21, 14], dtype=np.float32) restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32) restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32) @@ -1286,7 +1286,7 @@ def make_atom14_dists_bounds( restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32) restype_atom14_ambiguous_atoms_swap_idx = np.tile( - np.arange(14, dtype=np.int), (21, 1) + np.arange(14, dtype=int), (21, 1) ) diff --git a/fastfold/data/data_pipeline.py b/fastfold/data/data_pipeline.py index ab6601e6..8fd1fa36 100644 --- a/fastfold/data/data_pipeline.py +++ b/fastfold/data/data_pipeline.py @@ -99,12 +99,12 @@ def make_sequence_features( ) features["between_segment_residues"] = np.zeros((num_res,), dtype=np.int32) features["domain_name"] = np.array( - [description.encode("utf-8")], dtype=np.object_ + [description.encode("utf-8")], dtype=object ) features["residue_index"] = np.array(range(num_res), dtype=np.int32) features["seq_length"] = np.array([num_res] * num_res, dtype=np.int32) features["sequence"] = np.array( - [sequence.encode("utf-8")], dtype=np.object_ + [sequence.encode("utf-8")], dtype=object ) return features @@ -137,7 +137,7 @@ def make_mmcif_features( ) mmcif_feats["release_date"] = np.array( - [mmcif_object.header["release_date"].encode("utf-8")], dtype=np.object_ + [mmcif_object.header["release_date"].encode("utf-8")], dtype=object ) mmcif_feats["is_distillation"] = np.array(0., dtype=np.float32) @@ -238,7 +238,7 @@ def make_msa_features(msas: Sequence[parsers.Msa]) -> FeatureDict: features["num_alignments"] = np.array( [num_alignments] * num_res, dtype=np.int32 ) - features["msa_species_identifiers"] = np.array(species_ids, dtype=np.object_) + features["msa_species_identifiers"] = np.array(species_ids, dtype=object) return features def run_msa_tool( @@ -681,7 +681,7 @@ def convert_monomer_features( ) -> FeatureDict: """Reshapes and modifies monomer features for multimer models.""" converted = {} - converted['auth_chain_id'] = np.asarray(chain_id, dtype=np.object_) + converted['auth_chain_id'] = np.asarray(chain_id, dtype=object) unnecessary_leading_dim_feats = { 'sequence', 'domain_name', 'num_alignments', 'seq_length' } diff --git a/fastfold/data/templates.py b/fastfold/data/templates.py index 3ed941c4..cb3de47c 100644 --- a/fastfold/data/templates.py +++ b/fastfold/data/templates.py @@ -85,8 +85,8 @@ class LengthError(PrefilterError): "template_aatype": np.int64, "template_all_atom_mask": np.float32, "template_all_atom_positions": np.float32, - "template_domain_names": np.object, - "template_sequence": np.object, + "template_domain_names": object, + "template_sequence": object, "template_sum_probs": np.float32, } @@ -1209,8 +1209,8 @@ def get_templates( "template_all_atom_positions": np.zeros( (1, num_res, residue_constants.atom_type_num, 3), np.float32 ), - "template_domain_names": np.array([''.encode()], dtype=np.object), - "template_sequence": np.array([''.encode()], dtype=np.object), + "template_domain_names": np.array([''.encode()], dtype=object), + "template_sequence": np.array([''.encode()], dtype=object), "template_sum_probs": np.array([0], dtype=np.float32), } diff --git a/fastfold/distributed/comm.py b/fastfold/distributed/comm.py index 3d667c9a..8ca2ef91 100644 --- a/fastfold/distributed/comm.py +++ b/fastfold/distributed/comm.py @@ -4,8 +4,8 @@ import torch.distributed as dist from torch import Tensor -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc +from colossalai.legacy.context import ParallelMode +from colossalai.legacy.core import global_context as gpc from .core import ensure_divisibility diff --git a/fastfold/distributed/comm_async.py b/fastfold/distributed/comm_async.py index ece417d3..d2e7c678 100644 --- a/fastfold/distributed/comm_async.py +++ b/fastfold/distributed/comm_async.py @@ -5,8 +5,8 @@ import torch.distributed as dist from torch import Tensor -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc +from colossalai.legacy.context import ParallelMode +from colossalai.legacy.core import global_context as gpc from .comm import _split, divide diff --git a/fastfold/distributed/core.py b/fastfold/distributed/core.py index ea9219f1..331b1b6c 100644 --- a/fastfold/distributed/core.py +++ b/fastfold/distributed/core.py @@ -1,3 +1,7 @@ +from colossalai.legacy import launch_from_torch as _ff_launch_from_torch +from colossalai.legacy import launch_from_torch as _ff_launch_from_torch +from colossalai.legacy import launch_from_torch as _ff_launch_from_torch +from colossalai.legacy import launch_from_torch as _ff_launch_from_torch import os import torch @@ -36,5 +40,5 @@ def init_dap(tensor_model_parallel_size_=None): set_missing_distributed_environ('MASTER_ADDR', "localhost") set_missing_distributed_environ('MASTER_PORT', 18417) - colossalai.launch_from_torch( + _ff_launch_from_torch( config={"parallel": dict(tensor=dict(size=tensor_model_parallel_size_))}) diff --git a/fastfold/model/fastnn/evoformer.py b/fastfold/model/fastnn/evoformer.py index 8af5d1dc..446c9f24 100644 --- a/fastfold/model/fastnn/evoformer.py +++ b/fastfold/model/fastnn/evoformer.py @@ -4,8 +4,8 @@ import torch import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc +from colossalai.legacy.context import ParallelMode +from colossalai.legacy.core import global_context as gpc from fastfold.model.fastnn import MSACore, OutProductMean, PairCore from fastfold.model.fastnn.ops import Linear diff --git a/fastfold/model/fastnn/msa.py b/fastfold/model/fastnn/msa.py index bfb68d06..50247188 100644 --- a/fastfold/model/fastnn/msa.py +++ b/fastfold/model/fastnn/msa.py @@ -18,8 +18,8 @@ import torch import torch.nn as nn import torch.nn.functional as F -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc +from colossalai.legacy.context import ParallelMode +from colossalai.legacy.core import global_context as gpc from fastfold.model.fastnn.kernel import LayerNorm, bias_dropout_add from fastfold.model.fastnn.ops import (ChunkMSARowAttentionWithPairBias, ChunkTransition, SelfAttention, GlobalAttention, Transition, diff --git a/fastfold/model/fastnn/template.py b/fastfold/model/fastnn/template.py index 8bcaee90..fa7b904a 100644 --- a/fastfold/model/fastnn/template.py +++ b/fastfold/model/fastnn/template.py @@ -18,8 +18,8 @@ import torch import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc +from colossalai.legacy.context import ParallelMode +from colossalai.legacy.core import global_context as gpc from fastfold.model.nn.primitives import Attention from fastfold.utils.checkpointing import checkpoint_blocks diff --git a/fastfold/utils/import_weights.py b/fastfold/utils/import_weights.py index 3e33a0b3..350d73ae 100644 --- a/fastfold/utils/import_weights.py +++ b/fastfold/utils/import_weights.py @@ -586,6 +586,7 @@ def FoldIterationParams(sm): def import_jax_weights_(model, npz_path, version="model_1"): + npz_path = '/' + npz_path data = np.load(npz_path) translations = get_translation_dict(model, version) diff --git a/setup.py b/setup.py index fd9bce06..afbb9ec7 100644 --- a/setup.py +++ b/setup.py @@ -111,7 +111,7 @@ def cuda_ext_helper(name, sources, extra_cuda_flags): cc_flag.append('arch=compute_80,code=sm_80') extra_cuda_flags = [ - '-std=c++14', '-maxrregcount=50', '-U__CUDA_NO_HALF_OPERATORS__', + '-std=c++17', '-maxrregcount=50', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr', '--expt-extended-lambda' ]