From 42da5f56b027c3b74b8dfef210e7bf2c867fa010 Mon Sep 17 00:00:00 2001 From: Molly Smith Date: Sat, 18 Feb 2023 01:38:20 +0000 Subject: [PATCH] Always convert input mask to half --- deepspeed/ops/transformer/inference/ds_attention.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py index 7a85190b71b1..f4ec14bf4990 100644 --- a/deepspeed/ops/transformer/inference/ds_attention.py +++ b/deepspeed/ops/transformer/inference/ds_attention.py @@ -244,8 +244,7 @@ def compute_attention(self, qkv_out, input_mask, layer_past, alibi): ) * self.num_attention_heads_per_partition if dist.is_initialized() else 0 attention_probs = self.softmax_func( attn_scores=attention_scores, - attn_mask=((1 - input_mask).half() * - minus_inf) if input_mask.dtype == torch.int64 else input_mask, + attn_mask=((1 - input_mask).half() * minus_inf), alibi=alibi, triangular=(self.config.triangular_masking and (attention_scores.shape[-2] > 1)),