Repository: huggingface/transformers

Issue:
[Feature Request] Add 3D attention mask for T5Model
## Environment info

- `transformers` version: 4.2.1
- Platform: Linux
- Python version: 3.6.9
- PyTorch version (GPU?): 1.7.0+cuda11.0

### Who can help
T5: @patrickvonplaten

## Information
The `get_extended_attention_mask()` does not sufficiently address the case with 3D attention mask. The problem emerges for T5Model as `input_ids` and `decoder_input_ids` are of different length and the `attention_mask` is of shape [Batch_size, Seq_length, Seq_length]. The decoder uses `attention_mask` directly for `encoder_attention_mask` in cross-attention, which is of incorrect shape and the error message does give any information about why it happens. 

## To reproduce
As described above. I can add code later if needed. 

## Expected behavior
I propose to add a sanity check for attention masks or improve the `get_extended_attention_mask()` method. 



Discussion:
Hello @yongyi-wu, 

yes you're right. T5 is not yet fully compatible with 3D attention_mask input. Currently, I won't find enough time to work on adding this feature, but I'll post it under "Community projects" in case someone from the community is interested in giving it a shot.

Also feel free to open a PR yourself, if you want to try :-) 
Hi, I am new here but would like to give this a shot. Because it is my first issue I could use some direction if on how to tackle this, if this is okay with you. 

- Would you prefer the sanity check or an improved `get_extended_attention_mask()` method? 
- Do you know of an already existing implementations with 3D attention mask as a reference. 
- Where would you like to see the solution implemented. 

Resolving patch:
diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py
--- a/src/transformers/models/t5/modeling_t5.py
+++ b/src/transformers/models/t5/modeling_t5.py
@@ -904,7 +904,13 @@ def forward(
         # ourselves in which case we just need to make it broadcastable to all heads.
         extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, inputs_embeds.device)
 
-        if self.is_decoder and encoder_attention_mask is not None:
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device)
             encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
         else:
             encoder_extended_attention_mask = None


Test patch:
diff --git a/tests/test_modeling_t5.py b/tests/test_modeling_t5.py
--- a/tests/test_modeling_t5.py
+++ b/tests/test_modeling_t5.py
@@ -530,6 +530,34 @@ def test_decoder_model_past_with_attn_mask(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_attention_mask_past(*config_and_inputs)
 
+    def test_decoder_model_past_with_3d_attn_mask(self):
+        (
+            config,
+            input_ids,
+            decoder_input_ids,
+            attention_mask,
+            decoder_attention_mask,
+            lm_labels,
+        ) = self.model_tester.prepare_config_and_inputs()
+
+        attention_mask = ids_tensor(
+            [self.model_tester.batch_size, self.model_tester.encoder_seq_length, self.model_tester.encoder_seq_length],
+            vocab_size=2,
+        )
+        decoder_attention_mask = ids_tensor(
+            [self.model_tester.batch_size, self.model_tester.decoder_seq_length, self.model_tester.decoder_seq_length],
+            vocab_size=2,
+        )
+
+        self.model_tester.create_and_check_decoder_model_attention_mask_past(
+            config,
+            input_ids,
+            decoder_input_ids,
+            attention_mask,
+            decoder_attention_mask,
+            lm_labels,
+        )
+
     def test_decoder_model_past_with_large_inputs(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)