[Bart] FP16 Support

src/transformers/modeling_bart.py

@@ -640,9 +640,9 @@ def forward(
             reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool)
             attn_weights = attn_weights.masked_fill(reshaped, float("-inf"))
             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-        attn_weights_float = F.softmax(attn_weights, dim=-1, dtype=torch.float32)
-        attn_weights = attn_weights_float.type_as(attn_weights)
-        attn_probs = F.dropout(attn_weights_float, p=self.dropout, training=self.training,)
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_probs = F.dropout(attn_weights, p=self.dropout, training=self.training,)
+
         assert v is not None
         attn_output = torch.bmm(attn_probs, v)
         assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
@@ -696,7 +696,7 @@ def _cat_prev_key_padding_mask(
         elif prev_key_padding_mask is not None:
             filler = torch.zeros(batch_size, src_len - prev_key_padding_mask.size(1))
             if prev_key_padding_mask.is_cuda:
-                filler = filler.cuda()
+                filler = filler.to(prev_key_padding_mask.device)
             new_key_padding_mask = torch.cat([prev_key_padding_mask.float(), filler.float()], dim=1)
         elif key_padding_mask is not None:
             filler = torch.zeros(batch_size, src_len - key_padding_mask.size(1))

tests/test_modeling_bart.py

@@ -294,6 +294,13 @@ def test_tokenization(self):
             bart_toks = tokenizer.encode(ex, return_tensors="pt")
             _assert_tensors_equal(desired_result.long(), bart_toks, prefix=ex)
 
+    @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
+    def test_generate_fp16(self):
+        config, input_ids, batch_size = self._get_config_and_data(output_past=True)
+        attention_mask = input_ids.ne(1)
+        lm_model = BartForMaskedLM(config).eval().to(torch_device).half()
+        lm_model.generate(input_ids, attention_mask)
+
 
 def _assert_tensors_equal(a, b, atol=1e-12, prefix=""):
     """If tensors not close, or a and b arent both tensors, raise a nice Assertion error."""