attn softmax done in float32

Author: lucidrains

alphafold3_pytorch/attention.py

@@ -447,6 +447,8 @@ def forward(
         memory_kv: Float['2 h m d'] | None = None
     ) -> Float['b h i d']:
 
+        dtype = q.dtype
+
         is_windowed_attn_bias = None
 
         if exists(attn_bias):
@@ -506,7 +508,9 @@ def forward(
 
         # attention
 
-        attn = sim.softmax(dim = -1)
+        attn = sim.softmax(dim = -1, dtype = torch.float32)
+        attn = attn.to(dtype)
+
         attn = self.attn_dropout(attn)
 
         # aggregate values

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.66"
+version = "0.2.67"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }