support running in multimer mode

fastfold/config.py

@@ -575,7 +575,7 @@ multimer_model_config_update = {
         "tm": {
             "c_z": c_z,
             "no_bins": aux_distogram_bins,
-            "enabled": tm_enabled,
+            "enabled": True,
         },
         "masked_msa": {
             "c_m": c_m,

fastfold/data/templates.py

@@ -881,18 +881,20 @@ def _process_single_hit(
     ) as e:
         # These 3 errors indicate missing mmCIF experimental data rather than a
         # problem with the template search, so turn them into warnings.
-        warning = (
-            "%s_%s (sum_probs: %.2f, rank: %d): feature extracting errors: "
-            "%s, mmCIF parsing errors: %s"
-            % (
-                hit_pdb_code,
-                hit_chain_id,
-                hit.sum_probs,
-                hit.index,
-                str(e),
-                parsing_result.errors,
-            )
-        )
+        
+        # warning = (
+        #     "%s_%s (sum_probs: %.2f, rank: %d): feature extracting errors: "
+        #     "%s, mmCIF parsing errors: %s"
+        #     % (
+        #         hit_pdb_code,
+        #         hit_chain_id,
+        #         hit.sum_probs,
+        #         hit.index,
+        #         str(e),
+        #         parsing_result.errors,
+        #     )
+        # )
+        warning=None
         if strict_error_check:
             return SingleHitResult(features=None, error=warning, warning=None)
         else:

fastfold/model/fastnn/evoformer.py

@@ -9,7 +9,7 @@ from colossalai.core import global_context as gpc
 
 from fastfold.model.fastnn import MSACore, OutProductMean, PairCore
 from fastfold.model.fastnn.ops import Linear
-from fastfold.distributed.comm import gather, scatter
+from fastfold.distributed.comm import gather, scatter, col_to_row
 from fastfold.distributed.comm_async import All_to_All_Async, All_to_All_Async_Opp
 from fastfold.utils.checkpointing import checkpoint_blocks
 
@@ -49,7 +49,10 @@ class Evoformer(nn.Module):
             m = torch.nn.functional.pad(m, (0, 0, 0, padding_size))
             z = torch.nn.functional.pad(z, (0, 0, 0, padding_size, 0, padding_size))
 
-            m = scatter(m, dim=1)
+            if self.is_multimer:
+                m = scatter(m, dim=2)
+            else:
+                m = scatter(m, dim=1)
             z = scatter(z, dim=1)
 
         msa_mask = msa_mask.unsqueeze(0)
@@ -76,7 +79,10 @@ class Evoformer(nn.Module):
             m = m.squeeze(0)
             z = z.squeeze(0)
 
-            m = gather(m, dim=0)
+            if self.is_multimer:
+                m = gather(m, dim=1)
+            else:
+                m = gather(m, dim=0)
             z = gather(z, dim=0)
 
             m = m[:, :-padding_size, :]
@@ -107,8 +113,10 @@ class Evoformer(nn.Module):
             z[0] = torch.nn.functional.pad(z[0], (0, 0, 0, padding_size, 0, padding_size))
             torch.cuda.empty_cache()
             
-
-            m[0] = scatter(m[0], dim=1, drop_unused=True)
+            if self.is_multimer:
+                m[0] = scatter(m[0], dim=2)
+            else:
+                m[0] = scatter(m[0], dim=1, drop_unused=True)
             torch.cuda.empty_cache()
             z[0] = scatter(z[0], dim=1, drop_unused=True)
         torch.cuda.empty_cache()
@@ -126,15 +134,8 @@ class Evoformer(nn.Module):
             z = self.pair.inplace(z, pair_mask)
             m[0] = All_to_All_Async_Opp.apply(m[0], work, 1, 2)
         else:
-            # z = self.communication.inplace(m[0], msa_mask, z)
-            # z_ori = z[0].clone()
-            # m[0], work = All_to_All_Async.apply(m[0], 1, 2)
-            # z = self.pair_stack.inplace(z, pair_mask)
-            # m[0] = All_to_All_Async_Opp.apply(m[0], work, 1, 2)
-            # m[0] = self.msa_stack(m[0], z_ori, msa_mask)
             z = self.communication.inplace(m[0], msa_mask, z)
-            m[0], work = All_to_All_Async.apply(m[0], 1, 2)
-            m[0] = All_to_All_Async_Opp.apply(m[0], work, 1, 2)
+            m[0] = col_to_row(m[0])
             m[0] = self.msa(m[0], z[0], msa_mask)
             z = self.pair.inplace(z, pair_mask)
 
@@ -143,7 +144,10 @@ class Evoformer(nn.Module):
             z[0] = z[0].squeeze(0)
             torch.cuda.empty_cache()
 
-            m[0] = gather(m[0], dim=0, chunks=4)
+            if self.is_multimer:
+                m[0] = gather(m[0], dim=1)
+            else:
+                m[0] = gather(m[0], dim=0, chunks=4)
             z[0] = gather(z[0], dim=0, chunks=4)
 
             m[0] = m[0][:, :-padding_size, :]

fastfold/model/hub/alphafold.py

@@ -360,7 +360,7 @@ class AlphaFold(nn.Module):
                     pair_mask=pair_mask.to(dtype=z[0].dtype),
                     _mask_trans=self.config._mask_trans,
                 )[0]
-        del extra_msa_feat, extra_msa_fn
+            del extra_msa_feat, extra_msa_fn
         
         torch.cuda.empty_cache()
 

fastfold/model/nn/structure_module.py

@@ -530,7 +530,7 @@ class InvariantPointAttention(nn.Module):
         # [*, N_res, C_s]
         if self.is_multimer:
             s = self.linear_out(
-                torch.cat((o, *o_pt, o_pt_norm, o_pair), dim=-1).to(dtype=z.dtype)
+                torch.cat((o, *o_pt, o_pt_norm, o_pair), dim=-1).to(dtype=z[0].dtype)
             )
         else:
             s = self.linear_out(
@@ -874,7 +874,8 @@ class StructureModule(nn.Module):
 
     def _forward_multimer(
         self,
-        evoformer_output_dict,
+        s: torch.Tensor,
+        z: torch.Tensor,
         aatype: torch.Tensor,
         mask: Optional[torch.Tensor] = None,
     ) -> Dict[str, Any]:
@@ -898,6 +899,7 @@ class StructureModule(nn.Module):
             s.device,
         )
         outputs = []
+        z = [z]
         for i in range(self.no_blocks):
             # [*, N, C_s]
             s = s + self.ipa(s, z, rigids, mask)
@@ -960,7 +962,7 @@ class StructureModule(nn.Module):
             A dictionary of outputs
         """
         if self.is_multimer:
-            outputs = self._forward_multimer(evoformer_output_dict, aatype, mask)
+            outputs = self._forward_multimer(evoformer_output_dict["single"], evoformer_output_dict["pair"], aatype, mask)
         else:
             outputs = self._forward_monomer(evoformer_output_dict, aatype, mask)
 

fastfold/utils/import_weights.py

@@ -126,11 +126,9 @@ def assign(translation_dict, orig_weights):
                 print(ref[0].shape)
                 print(weights[0].shape)
                 raise
-  
-    
-def import_jax_weights_(model, npz_path, version="model_1", is_multimer: bool = False):
-    data = np.load(npz_path)
-    # translations = get_translation_dict(model, is_multimer=("multimer" in version))
+
+def get_translation_dict(model, version):
+    is_multimer = "multimer" in version
     #######################
     # Some templates
     #######################
@@ -540,16 +538,14 @@ def import_jax_weights_(model, npz_path, version="model_1", is_multimer: bool =
             },
         }
 
-    # return translations
-
     no_templ = [
-            "model_3",
-            "model_4",
-            "model_5",
-            "model_3_ptm",
-            "model_4_ptm",
-            "model_5_ptm",
-        ]
+        "model_3",
+        "model_4",
+        "model_5",
+        "model_3_ptm",
+        "model_4_ptm",
+        "model_5_ptm",
+    ]
     if version in no_templ:
         evo_dict = translations["evoformer"]
         keys = list(evo_dict.keys())
@@ -557,10 +553,19 @@ def import_jax_weights_(model, npz_path, version="model_1", is_multimer: bool =
             if "template_" in k:
                 evo_dict.pop(k)
 
-    if "_ptm" in version:
+    if "_ptm" in version or is_multimer:
         translations["predicted_aligned_error_head"] = {
             "logits": LinearParams(model.aux_heads.tm.linear)
         }
+        
+    return translations
+
+
+
+def import_jax_weights_(model, npz_path, version="model_1"):
+    data = np.load(npz_path)
+
+    translations = get_translation_dict(model, version)
 
     # Flatten keys and insert missing key prefixes
     flat = _process_translations_dict(translations)
@@ -578,5 +583,3 @@ def import_jax_weights_(model, npz_path, version="model_1", is_multimer: bool =
 
     # Set weights
     assign(flat, data)
-
-

fastfold/workflow/template/fastfold_data_workflow.py

@@ -119,7 +119,8 @@ class FastFoldDataWorkFlow:
 
 
     def run(self, fasta_path: str, alignment_dir: str=None, storage_dir: str=None) -> None:
-        storage_dir = "file:///tmp/ray/" + os.getlogin() + "/workflow_data"
+        timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) 
+        storage_dir = "file:///tmp/ray/" + str(timestamp) + "/workflow_data"
         if storage_dir is not None:
             if not os.path.exists(storage_dir):
                 os.makedirs(storage_dir[7:], exist_ok=True)

fastfold/workflow/template/fastfold_multimer_data_workflow.py

@@ -137,7 +137,8 @@ class FastFoldMultimerDataWorkFlow:
 
 
     def run(self, fasta_path: str, alignment_dir: str=None, storage_dir: str=None) -> None:
-        storage_dir = "file:///tmp/ray/" + os.getlogin() + "/workflow_data"
+        timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) 
+        storage_dir = "file:///tmp/ray/" + str(timestamp) + "/workflow_data"
         if storage_dir is not None:
             if not os.path.exists(storage_dir):
                 os.makedirs(storage_dir[7:], exist_ok=True)

inference.py

@@ -312,21 +312,31 @@ def inference_multimer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
-    amber_relaxer = relax.AmberRelaxation(
-        use_gpu=True,
-        **config.relax,
-    )
+    if(args.relaxation):
+        amber_relaxer = relax.AmberRelaxation(
+            use_gpu=True,
+            **config.relax,
+        )
 
-    # Relax the prediction.
-    t = time.perf_counter()
-    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
-    print(f"Relaxation time: {time.perf_counter() - t}")
+        # Relax the prediction.
+        t = time.perf_counter()
+        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
+        print(f"Relaxation time: {time.perf_counter() - t}")
+
+        # Save the relaxed PDB.
+        relaxed_output_path = os.path.join(args.output_dir,
+                                            f'{tag}_{args.model_name}_relaxed.pdb')
+        with open(relaxed_output_path, 'w') as f:
+            f.write(relaxed_pdb_str)
+        
+    if(args.save_outputs):
+            output_dict_path = os.path.join(
+                args.output_dir, f'{tag}_{args.model_name}_output_dict.pkl'
+            )
+            with open(output_dict_path, "wb") as fp:
+                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
 
-    # Save the relaxed PDB.
-    relaxed_output_path = os.path.join(args.output_dir,
-                                        f'{tag}_{args.model_name}_relaxed.pdb')
-    with open(relaxed_output_path, 'w') as f:
-        f.write(relaxed_pdb_str)
+            logger.info(f"Model output written to {output_dict_path}...")
 
 
 def inference_monomer_model(args):
@@ -454,21 +464,22 @@ def inference_monomer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
-    amber_relaxer = relax.AmberRelaxation(
-        use_gpu=True,
-        **config.relax,
-    )
-
-    # Relax the prediction.
-    t = time.perf_counter()
-    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
-    print(f"Relaxation time: {time.perf_counter() - t}")
+    if(args.relaxation):
+        amber_relaxer = relax.AmberRelaxation(
+            use_gpu=True,
+            **config.relax,
+        )
 
-    # Save the relaxed PDB.
-    relaxed_output_path = os.path.join(args.output_dir,
-                                        f'{tag}_{args.model_name}_relaxed.pdb')
-    with open(relaxed_output_path, 'w') as f:
-        f.write(relaxed_pdb_str)
+        # Relax the prediction.
+        t = time.perf_counter()
+        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
+        print(f"Relaxation time: {time.perf_counter() - t}")
+
+        # Save the relaxed PDB.
+        relaxed_output_path = os.path.join(args.output_dir,
+                                            f'{tag}_{args.model_name}_relaxed.pdb')
+        with open(relaxed_output_path, 'w') as f:
+            f.write(relaxed_pdb_str)
     
     if(args.save_outputs):
             output_dict_path = os.path.join(
@@ -512,6 +523,9 @@ if __name__ == "__main__":
                         help="""Path to model parameters. If None, parameters are selected
              automatically according to the model name from 
              ./data/params""")
+    parser.add_argument(
+        "--relaxation", action="store_false", default=False,
+    )
     parser.add_argument(
         "--save_outputs", action="store_true", default=False,
         help="Whether to save all model outputs, including embeddings, etc."