fix default mode missing additive mask option (#924)

seryilmaz · Sukru Eryilmaz · web-flow · commit 700d6825e205 · 2020-07-30T12:33:29.000-07:00
Co-authored-by: Sukru Eryilmaz &lt;seryilmaz@computelab-dgx1v-32.nvidia.com&gt;
diff --git a/apex/contrib/multihead_attn/self_multihead_attn_func.py b/apex/contrib/multihead_attn/self_multihead_attn_func.py
@@ -6,7 +6,7 @@ class SelfAttnFunc(torch.autograd.Function):
     def forward(ctx, use_time_mask, is_training, heads, scale, inputs,
                 input_weights, output_weights,
                 input_biases, output_biases,
-                mask, dropout_prob):
+                mask, is_additive_mask, dropout_prob):
         use_biases_t   = torch.tensor([input_biases is not None])
         heads_t        = torch.tensor([heads])
         scale_t        = torch.tensor([scale])
@@ -60,8 +60,11 @@ def forward(ctx, use_time_mask, is_training, heads, scale, inputs,
                 batches,seql_q,seql_k = matmul1_results.size()
                 seqs = int(batches / heads)
                 matmul1_results = matmul1_results.view(seqs, heads, seql_q, seql_k)
-                mask = mask.to(torch.bool)
-                matmul1_results = matmul1_results.masked_fill_(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+                if is_additive_mask:
+                    matmul1_results = matmul1_results + mask.unsqueeze(1).unsqueeze(2)
+                else:
+                    mask = mask.to(torch.bool)
+                    matmul1_results = matmul1_results.masked_fill_(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
                 matmul1_results = matmul1_results.view(seqs*heads, seql_q, seql_k)
 
         softmax_results = F.softmax(matmul1_results, dim=-1)
