admlakd

stanford-crfm · Apr 26, 2024 · 1ad2a42 · 1ad2a42
1 parent 324f881
commit 1ad2a42
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/levanter/models/attention.py b/src/levanter/models/attention.py
@@ -755,6 +755,10 @@ def wrap_flash_attention(q, k, v):
             mask=kernel_mask, head_shards=1, q_seq_shards=1, block_sizes=block_sizes
         )
 
+        # try upcasting to float32 to see if it fixes crash?
+        q = q.astype(jnp.float32)
+        k = k.astype(jnp.float32)
+        v = v.astype(jnp.float32)
         print(q.dtype, k.dtype, v.dtype)
         return jax.vmap(splash_kernel)(q, k, v, segment_ids=None)