Merge pull request #395 from bigict/amp

refactor: disable autocast for headers

Author: chungongyu

profold2/model/alphafold2.py

@@ -13,7 +13,7 @@
 from einops import rearrange
 
 from profold2.common import residue_constants
-from profold2.model import commons, folding, functional
+from profold2.model import accelerator, commons, folding, functional
 from profold2.model.evoformer import Evoformer
 from profold2.model.head import HeaderBuilder
 from profold2.utils import env, exists, status
@@ -282,27 +282,29 @@ def forward(
     representations.update(msa=m, pair=x, single=s, frames=t, single_init=s)
 
     ret.headers = {}
-    for name, module, options in self.headers:
-      value = module(ret.headers, representations, batch)
-      if not exists(value):
-        continue
-      ret.headers[name] = value
-      if 'representations' in value:
-        representations.update(value['representations'])
-    if 'folding' in ret.headers:
-      batch = folding.multi_chain_permutation_alignment(ret.headers['folding'], batch)
-    if self.training and compute_loss:
+    with accelerator.autocast(enabled=False):
+      representations = functional.to(representations, dtype=torch.float)
       for name, module, options in self.headers:
-        if not hasattr(module, 'loss') or name not in ret.headers:
+        value = module(ret.headers, representations, batch)
+        if not exists(value):
           continue
-        loss = module.loss(ret.headers[name], batch)
-        if exists(loss):
-          ret.headers[name].update(loss)
-          lossw = loss['loss'] * options.get('weight', 1.0)
-          if exists(ret.loss):
-            ret.loss = commons.tensor_add(ret.loss, lossw)
-          else:
-            ret.loss = lossw
+        ret.headers[name] = value
+        if 'representations' in value:
+          representations.update(value['representations'])
+      if 'folding' in ret.headers:
+        batch = folding.multi_chain_permutation_alignment(ret.headers['folding'], batch)
+      if self.training and compute_loss:
+        for name, module, options in self.headers:
+          if not hasattr(module, 'loss') or name not in ret.headers:
+            continue
+          loss = module.loss(ret.headers[name], batch)
+          if exists(loss):
+            ret.headers[name].update(loss)
+            lossw = loss['loss'] * options.get('weight', 1.0)
+            if exists(ret.loss):
+              ret.loss = commons.tensor_add(ret.loss, lossw)
+            else:
+              ret.loss = lossw
 
     if return_recyclables:
       msa_first_row_repr, pairwise_repr = m[..., 0, :, :], representations['pair']

profold2/model/folding.py

@@ -142,36 +142,39 @@ def forward(self, representations, batch):
     single_repr_init = single_repr
     single_repr = self.single_repr_dim(single_repr)
 
-    # prepare float32 precision for equivariance
-    original_dtype = single_repr.dtype
-    single_repr, pairwise_repr = map(lambda t: t.float(), (single_repr, pairwise_repr))
+    # initial frames
+    if 'frames' in representations and exists(representations['frames']):
+      quaternions, translations = representations['frames']
+    else:
+      quaternions = torch.tensor([1., 0., 0., 0.], device=device)
+      quaternions = torch.tile(quaternions, b + (n, 1))
+      translations = torch.zeros(b + (n, 3), device=device)
 
     outputs = []
 
     # iterative refinement with equivariant transformer in high precision
-    with commons.torch_default_dtype(torch.float32):
-      # initial frames
-      if 'frames' in representations and exists(representations['frames']):
-        quaternions, translations = representations['frames']
-      else:
-        quaternions = torch.tensor([1., 0., 0., 0.], device=device)
-        quaternions = torch.tile(quaternions, b + (n, 1))
-        translations = torch.zeros(b + (n, 3), device=device)
+    with accelerator.autocast(enabled=False):
+      # prepare float32 precision for equivariance
+      single_repr, pairwise_repr = functional.to(
+          (single_repr, pairwise_repr), dtype=torch.float
+      )
+      quaternions, translations = functional.to(
+          (quaternions, translations), dtype=torch.float
+      )
       rotations = functional.quaternion_to_matrix(quaternions).detach()
 
       # go through the layers and apply invariant point attention and
       # feedforward
       for i in range(self.structure_module_depth):
         is_last = i == (self.structure_module_depth - 1)
 
-        with accelerator.autocast(enabled=False):
-          single_repr = self.ipa_block(
-              single_repr.float(),
-              mask=batch['mask'].bool(),
-              pairwise_repr=pairwise_repr.float(),
-              rotations=rotations.float(),
-              translations=translations.float()
-          )
+        single_repr = self.ipa_block(
+            single_repr,
+            mask=batch['mask'].bool(),
+            pairwise_repr=pairwise_repr,
+            rotations=rotations,
+            translations=translations
+        )
 
         # update quaternion and translation
         quaternion_update, translation_update = self.to_affine_update(single_repr
@@ -197,7 +200,6 @@ def forward(self, representations, batch):
               functional.l2_norm(angles)
           )
           coords = functional.rigids_to_positions(frames, batch['seq'])
-          coords.type(original_dtype)
           outputs.append(
               dict(
                   frames=functional.rigids_scale(
@@ -212,4 +214,5 @@ def forward(self, representations, batch):
 
 
 def multi_chain_permutation_alignment(value, batch):
-  return functional.multi_chain_permutation_alignment(value, batch)
+  with accelerator.autocast(enabled=False):
+    return functional.multi_chain_permutation_alignment(value, batch)

profold2/model/functional.py

@@ -8,7 +8,6 @@
 from einops import rearrange, repeat
 
 from profold2.common import residue_constants
-from profold2.model import accelerator
 from profold2.utils import default, exists
 
 
@@ -23,6 +22,23 @@ def apc(x, dim=None, epsilon=1e-8):
   return x - avg
 
 
+def to(x, *args, **kwargs):
+  """Performs Tensor dtype and/or device conversion. A torch.dtype and torch.device
+     are inferred from the arguments of to(*args, **kwargs).
+     See ```torch.Tensor.to```"""
+  if isinstance(x, dict):
+    for k, v in x.items():
+      x[k] = to(v, *args, **kwargs)
+  elif isinstance(x, list):
+    for i, v in enumerate(x):
+      x[i] = to(v, *args, **kwargs)
+  elif isinstance(x, tuple):
+    x = tuple(to(v, *args, **kwargs) for v in x)
+  elif isinstance(x, torch.Tensor):
+    x = x.to(*args, **kwargs)
+  return x
+
+
 def l2_norm(v, dim=-1, epsilon=1e-12):
   return v / torch.clamp(
       torch.linalg.norm(v, dim=dim, keepdim=True),  # pylint: disable=not-callable
@@ -1503,8 +1519,7 @@ def optimal_transform_create(pred_points, true_points, points_mask):
     with torch.no_grad():
       pred_ca = true_ca
 
-  with accelerator.autocast(enabled=False):
-    return kabsch_transform(pred_ca.float(), true_ca.float())
+  return kabsch_transform(pred_ca.float(), true_ca.float())
 
 
 def seq_crop_mask(fgt_seq_index, fgt_seq_color, seq_index, seq_color, seq_anchor):

profold2/model/head.py

@@ -282,7 +282,7 @@ def forward(self, headers, representations, batch):
           '... m j d,b i j q,q c d,d -> ... m i c', msa, eij, states, self.mask
       )
       ret.update(wab=states)
-      # eij = torch.einsum('b i j q,q c d -> b i j c d', eij, states)
+      # eij = torch.einsum('... i j q,q c d -> ... i j c d', eij, states)
     else:  # native potts model
       eij = rearrange(eij, '... i j (c d) -> ... i j c d', c=num_class, d=num_class)
       hi = torch.einsum('... m j d,b i j c d,d -> ... m i c', msa, eij, self.mask)
@@ -687,7 +687,7 @@ def yield_pred_angle_loss(pred_angles_list, true_angles):
 
     assert 'traj' in value
     # loss
-    loss = (fape_loss(value['traj'], batch) + torsion_angle_loss(value['traj'], batch))
+    loss = fape_loss(value['traj'], batch) + torsion_angle_loss(value['traj'], batch)
 
     return dict(loss=loss)
 
@@ -1591,9 +1591,10 @@ def forward(self, headers, representations, batch):
       )
       with torch.no_grad():
         # rotate / align
-        pred_points, true_points = functional.kabsch_align(
-            pred_points, true_points, mask=coord_mask
-        )
+        with accelerator.autocast(enabled=False):
+          pred_points, true_points = functional.kabsch_align(
+              pred_points.float(), true_points.float(), mask=coord_mask
+          )
         n = torch.sum(batch['mask'], dim=-1, keepdim=True)
         tms = functional.tmscore(pred_points, true_points, n=n)
       tms = torch.where(torch.any(coord_mask, dim=-1), tms, 0)