Use FP32 for multi-head attention softmax
@@ -129,7 +129,7 @@ class MultiheadAttention(nn.Module):
float('-inf'),
).type_as(attn_weights) # FP16 support: cast to float and back
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
- attn_weights = F.softmax(attn_weights, dim=-1)
+ attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights)
attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)
attn = torch.bmm(attn_weights, v)
