fp16 fixes

openfold/model/heads.py

@@ -137,7 +137,7 @@ def __init__(self, c_z, no_bins, **kwargs):
 
         self.linear = Linear(self.c_z, self.no_bins, init="final")
 
-    def forward(self, z):  # [*, N, N, C_z]
+    def _forward(self, z):  # [*, N, N, C_z]
         """
         Args:
             z:
@@ -149,8 +149,16 @@ def forward(self, z):  # [*, N, N, C_z]
         logits = self.linear(z)
         logits = logits + logits.transpose(-2, -3)
         return logits
-
-
+
+    def forward(self, z):
+
+        float16_enabled = (torch.get_autocast_gpu_dtype() == torch.float16)
+        if float16_enabled and torch.is_autocast_enabled():
+            with torch.cuda.amp.autocast(enabled=False):
+                return self._forward(z.float())
+        else:
+            return self._forward(z)
+
 class TMScoreHead(nn.Module):
     """
     For use in computation of TM-score, subsection 1.9.7

openfold/model/outer_product_mean.py

@@ -93,7 +93,7 @@ def _chunk(self,
 
         return outer
 
-    def forward(self, 
+    def _forward(self, 
         m: torch.Tensor, 
         mask: Optional[torch.Tensor] = None,
         chunk_size: Optional[int] = None,
@@ -143,3 +143,18 @@ def forward(self,
             outer = outer / norm
 
         return outer
+
+    def forward(self,
+                m: torch.Tensor,
+                mask: Optional[torch.Tensor] = None,
+                chunk_size: Optional[int] = None,
+                inplace_safe: bool = False,
+    ) -> torch.Tensor:
+
+        float16_enabled = (torch.get_autocast_gpu_dtype() == torch.float16)
+        if float16_enabled and torch.is_autocast_enabled():
+            with torch.cuda.amp.autocast(enabled=False):
+                return self._forward(m.float(), mask, chunk_size, inplace_safe)
+        else:
+            return self._forward(m, mask, chunk_size, inplace_safe)
+

openfold/model/primitives.py

@@ -479,6 +479,9 @@ def forward(
         q, k, v = self._prep_qkv(q_x, kv_x)
 
         # [*, Q, H, C_hidden]
+        float16_enabled = (torch.get_autocast_gpu_dtype() == torch.float16)
+        if float16_enabled:
+            use_memory_efficient_kernel = False
         if(use_memory_efficient_kernel):
             if(len(biases) > 2):
                 raise ValueError(

openfold/model/structure_module.py

@@ -312,10 +312,18 @@ def forward(
             z[0] = z[0].cpu()
 
         # [*, H, N_res, N_res]
-        a = torch.matmul(
-            permute_final_dims(q, (1, 0, 2)),  # [*, H, N_res, C_hidden]
-            permute_final_dims(k, (1, 2, 0)),  # [*, H, C_hidden, N_res]
-        )
+        float16_enabled = (torch.get_autocast_gpu_dtype() == torch.float16)
+        if float16_enabled and torch.is_autocast_enabled():
+            with torch.cuda.amp.autocast(enabled=False):
+                a = torch.matmul(
+                    permute_final_dims(q.float(), (1, 0, 2)),  # [*, H, N_res, C_hidden]
+                    permute_final_dims(k.float(), (1, 2, 0)),  # [*, H, C_hidden, N_res]
+                )
+        else:
+            a = torch.matmul(
+                permute_final_dims(q, (1, 0, 2)),  # [*, H, N_res, C_hidden]
+                permute_final_dims(k, (1, 2, 0)),  # [*, H, C_hidden, N_res]
+            )
         a *= math.sqrt(1.0 / (3 * self.c_hidden))
         a += (math.sqrt(1.0 / 3) * permute_final_dims(b, (2, 0, 1)))
 

openfold/model/triangular_multiplicative_update.py

@@ -391,7 +391,12 @@ def forward(self,
         b = mask
         b = b * self.sigmoid(self.linear_b_g(z))
         b = b * self.linear_b_p(z)
-        x = self._combine_projections(a, b)
+        float16_enabled = (torch.get_autocast_gpu_dtype() == torch.float16)
+        if float16_enabled and torch.is_autocast_enabled():
+            with torch.cuda.amp.autocast(enabled=False):
+                x = self._combine_projections(a.float(), b.float())
+        else:
+            x = self._combine_projections(a, b)
         del a, b
         x = self.layer_norm_out(x)
         x = self.linear_z(x)