fix triton matmul

ai-forever · Nov 17, 2021 · e186fcd · e186fcd
1 parent 50eee08
commit e186fcd
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/mpu/transformer.py b/src/mpu/transformer.py
@@ -125,9 +125,9 @@ def forward(self, hidden_states, ltor_mask):
 
         if self.use_deepspeed_sparse:
             context_layer = self.sparse_self_attention(
-                query_layer,
-                key_layer,
-                value_layer,
+                query_layer.float(),
+                key_layer.float(),
+                value_layer.float(),
                 attn_mask=ltor_mask)
         else:
             # Raw attention scores. [b, np, s, s]