* Force cast to fp32 to avoid atten layer overflow

maosatgithub · Dec 7, 2022 · 6e92cda · 6e92cda
1 parent d7440ac
commit 6e92cda
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py
@@ -167,9 +167,13 @@ def forward(self, x, context=None, mask=None):
 
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
 
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        # force cast to fp32 to avoid overflowing
+        with torch.autocast(enabled=False, device_type = 'cuda'):
+            q, k = q.float(), k.float()
+            sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
         del q, k
-
+    
         if exists(mask):
             mask = rearrange(mask, 'b ... -> b (...)')
             max_neg_value = -torch.finfo(sim.dtype).max