Merge pull request #295 from nikitos9000/improve_fp16_stability

Improve TriangularMultiplicativeUpdate stability in fp16 mode
aqlaboratory · Apr 10, 2023 · 208cce6 · 208cce6
2 parents ee5d2c3 + 6625e8d
commit 208cce6
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/openfold/model/triangular_multiplicative_update.py b/openfold/model/triangular_multiplicative_update.py
@@ -392,8 +392,13 @@ def forward(self,
         b = mask
         b = b * self.sigmoid(self.linear_b_g(z))
         b = b * self.linear_b_p(z)
-
-        if(is_fp16_enabled()): 
+
+        # Prevents overflow of torch.matmul in combine projections in
+        # reduced-precision modes
+        a = a / a.std()
+	b = b / b.std()
+
+        if(is_fp16_enabled()):
             with torch.cuda.amp.autocast(enabled=False):
                 x = self._combine_projections(a.float(), b.float())
         else: