fixing arange not knowing half on cpu

facebookresearch · Jun 4, 2022 · 2050923 · 2050923
1 parent 0ef428a
commit 2050923
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/tests/test_rotary_embeddings.py b/tests/test_rotary_embeddings.py
@@ -68,8 +68,8 @@ def test_rotary_embeddings(device, dtype):
     assert k_rot.dtype == k.dtype
 
     # Check that the sequences now encode relative position information
-    att = torch.einsum("bhne,bhme->bhnm", q, k)
-    att_rot = torch.einsum("bhne,bhme->bhnm", q_rot, k_rot)
+    att = torch.einsum("bhne,bhme->bhnm", q, k).float()
+    att_rot = torch.einsum("bhne,bhme->bhnm", q_rot, k_rot).float()
 
     # - the attention for the same positions is not meaningfully changed
     assert torch.allclose(

diff --git a/xformers/components/positional_embedding/rotary.py b/xformers/components/positional_embedding/rotary.py
@@ -67,12 +67,14 @@ def _update_cos_sin_tables(self, x, seq_dimension=1):
             or self._cos_cached.dtype != x.dtype
         ):
             self._seq_len_cached = seq_len
-            t = torch.arange(x.shape[seq_dimension], device=x.device, dtype=x.dtype)
+            t = torch.arange(
+                x.shape[seq_dimension], device=x.device, dtype=torch.float32
+            )
             freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(x.dtype))
             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
 
-            self._cos_cached = emb.cos()[None, None, :, :]
-            self._sin_cached = emb.sin()[None, None, :, :]
+            self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype)
+            self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype)
 
         return self._cos_cached, self._sin_cached