Fix tf.concatenate + test past_key_values for TF models
Author: ydshieh
Pull request: yes

# What does this PR do?

- Fix `tf.concatenate` (#15477)
    -  it should be `tf.concat`
- Add testing about `past_key_values` for TF models
  Currently, only added to 
    - `testing_modeling_tf_bert.py`
    - TF model testing template

  The changes include:
    - add tests for: causal lm **base model** / causal lm model / decoder (i.e. with cross attention) 
    - rename some methods / add docstrings, to make them more clear (without reading their codes)
- Fix `extended_attention_mask` when `past_key_values` is passed:
    - it should be the mask for the present inputs (i.e. shouldn't contain the past part) 
    (detected when the above tests are added)
    (this is already done for `TFT5Model`, see [here](https://github.com/huggingface/transformers/blob/7f921bcf47f2d6875feb6c1232dc6cfbeb4668e8/src/transformers/models/t5/modeling_tf_t5.py#L732))


TF: @gante  @Rocketknight1 
Also cc @sgugger for the changes in the templates

Comments:
_The documentation is not available anymore as the PR was closed or merged._

> This all looks good to me, thanks a lot for fixing! Pinging @patrickvonplaten for the `past_key_values` change, to make sure it's all okay.
> 
> Re-the test. Is BERT the only test file to adapt?

Hi, I don't understand this question, but let me try to answer
  - the test file for TF-BERT is already changed.
  - we need to add the same tests to other TF models, like `roberta`, `electra` (the encoder-like models that can also be used as decoder)

(Let me know if I misunderstand your question, @sgugger )

Yes, that was my question. Will you add those other test files in the same PR?

> Yes, that was my question. Will you add those other test files in the same PR?

Sure, I intend to do so (after an approval for the change in TF-BERT test).
I will wait a bit - to have reviews from other people - and add to other test files :-)

@patrickvonplaten 

I have a few questions while looking other models like TF-T5: (see the end of the actual question)

https://github.com/huggingface/transformers/blob/d1fcc90abf34cc498c8a65a717ad0d9354ceca97/src/transformers/models/t5/modeling_tf_t5.py#L717-L734

- Line 717-718 implies that if **3D attention mask** is passed, then it means `from_seq_length - cached_seq_length` (because there is no slice)
- Line 732 implies if 2D mask + is_decoder + pass_key_values, attn mask means `from_seq_length`, and we slice it to become `from_seq_length - cached_seq_length`.

Furthermore, from 

https://github.com/huggingface/transformers/blob/d1fcc90abf34cc498c8a65a717ad0d9354ceca97/src/transformers/models/t5/modeling_tf_t5.py#L388-L393

 - `TFT5Attention` requires `mask` to be `from_seq_length - cached_seq_length`.

**My Question**: depending **3D** or **2D** format, `attention_mask` in `TFT5MainLayer` has different meaning (in terms of the query length). It could make sense - but only if we document it clearly.

Or, is this rather a mistake, and even for 3D case, it should actually mean to have the full sequence length, and we need to slice it to the present sequence only (so a fix is required here) ?
 

I think it expects the user to provide the correct shape in 3D case, i.e already sliced to only the present sequence length. Just want to have a double check from you.

One point about the `extended_attention_mask` - when we call `generate()` with the new changes we **do** expect that the `attention_mask` can be used to mask positions in the `past_key_values`. This is something I only realized when I started working on `generate()`! It probably only applies to causal LMs, but the fact that we seem to be truncating the attention mask in some cases is concerning, as it will cause that method to break. Are we sure that that truncation is only happening in models that don't have a `generate()` method?

> One point about the `extended_attention_mask` - when we call `generate()` with the new changes we **do** expect that the `attention_mask` can be used to mask positions in the `past_key_values`. This is something I only realized when I started working on `generate()`! It probably only applies to causal LMs, but the fact that we seem to be truncating the attention mask in some cases is concerning, as it will cause that method to break. Are we sure that that truncation is only happening in models that don't have a `generate()` method?


Hi, @Rocketknight1 

- 1. the truncation `extended_attention_mask[:, :, -seq_length:, :]` - I only see it when a model is used as a decoder, which has `generate` in general. This happens when the provided attention mask is 2D (with full past+present length).
- 2. For 3D attention mask, from the code, no such truncation - so I think it expects the one only contains for the `present` inputs without `past`. (and this is what I want to have a confirmation from @patrickvonplaten)
- 3. For `generate`, do you ever see in the code dealing with a case that uses 3D attention mask? If only 2D case occurrs, I think the code for dealing with attention mask in the models is fine, from point 1.

Add the tests to other TF models:

- roberta (2 failed - need to investigate) [Fixed](https://github.com/huggingface/transformers/pull/15774/files#r814551298)
- rembert
- electra (no `TFElectraForCausalLM` at this moment)

> One point about the `extended_attention_mask` - when we call `generate()` with the new changes we **do** expect that the `attention_mask` can be used to mask positions in the `past_key_values`. This is something I only realized when I started working on `generate()`! It probably only applies to causal LMs, but the fact that we seem to be truncating the attention mask in some cases is concerning, as it will cause that method to break. Are we sure that that truncation is only happening in models that don't have a `generate()` method?

That's not true really I think. When `past_key_values` is used there is no need to mask them with `attention_mask`.
The reason is that `past_key_values` are **only** used for causal LM models so that future tokens coming after `past_key_values` can never influence previous tokens (that's the whole reason why we are able to cache past key values states). This means that when `attention_mask` is used in combination with `past_key_values`, the `attention_mask` can only influence the `input_ids` that are forwarded with it - never previous cached `past_key_values`. 

Now when talking about XLA-compatible generation note that even there this should hold true because all "padded"/"non-reached" vectors of `past_key_values` should be padded anyways with -10,000.0

@ydshieh, this PR looks very nice and it's an important one. Let's try to keep the discussion and the code focus on just this issue here which is:

**Correct `past_key_values` concatenation and `attention_mask`-slicing behavior if BERT-like models are used for generation**

Given that literally the syntax is not even correct, there were 0 tests, 0 usage, so it's great that we are adding tests here. That's it - nothing more.

`past_key_values` is one of the most complicated topics, so I'd be very happy if we could move further discussions / changes to issues.



> > One point about the `extended_attention_mask` - when we call `generate()` with the new changes we **do** expect that the `attention_mask` can be used to mask positions in the `past_key_values`. This is something I only realized when I started working on `generate()`! It probably only applies to causal LMs, but the fact that we seem to be truncating the attention mask in some cases is concerning, as it will cause that method to break. Are we sure that that truncation is only happening in models that don't have a `generate()` method?
> 
> That's not true really I think. When `past_key_values` is used there is no need to mask them with `attention_mask`. The reason is that `past_key_values` are **only** used for causal LM models so that future tokens coming after `past_key_values` can never influence previous tokens (that's the whole reason why we are able to cache past key values states). This means that when `attention_mask` is used in combination with `past_key_values`, the `attention_mask` can only influence the `input_ids` that are forwarded with it - never previous cached `past_key_values`.
> 
> Now when talking about XLA-compatible generation note that even there this should hold true because all "padded"/"non-reached" vectors of `past_key_values` should be padded anyways with -10,000.0

After some discussion with @Rocketknight1 offline, there is actually one use case where it's very important to make sure that the `attention_mask` when forwarding `past_key_values` states. So for completeness and correctness I want to outline this here real quick. When doing batched generation where one input is longer than another one we usually make sure to do **left-padding** meaning that the input_ids and attention_mask would look like this:

```python
input_ids = [["hello", "this", "is", "very", "long"], ["<PAD>", "<PAD>", "This", "is", "short"]]
attention_mask = [[1,1,1,1,1], [0, 0, 1, 1, 1]]
```

In this case, the second `attention_mask`'s 0s will mask cached `past_key_values` for all future query vectors to come. But apart from this there should really be no use case where the `attention_mask` should influence `past_key_values`.


@patrickvonplaten 

Could you check this part, please:

https://github.com/huggingface/transformers/pull/15774/files#r814551298

The computation of `position_ids` in Roberta model (regarding to `padding_idx`) failed some testing for `past` -> I need to ensure `padding_idx` not in the initial input ids.

(Finally, the tests are all green 🚀)

Although currently there is no `TFElectraForCausalLM`, but we can still add tests like 

  - create_and_check_causal_lm_model_past
  - create_and_check_causal_lm_model_past_with_attn_mask
  - create_and_check_causal_lm_model_past_large_inputs
  - create_and_check_decoder_model_past_large_inputs

by using `TFElectraModel` (as decoder and/or add cross attention) - although the method names here are a bit misleading in this case.

Let me add them before merge! 

==> Done. Think everything is fine from my side, waiting for reviews when you have time :-)

@ydshieh I have a couple of questions I left as comments, but overall this seems like a good PR and I'm happy to merge after those are addressed!

> Just finished reviewing and this looks good! In my earlier comment, I hadn't looked at these changes thoroughly and didn't realize that the attention mask crop only applied when cross-attention was being performed for encoder-decoder models - this was my bad, I'm sorry!

No worry, I am a fan of TFBoys :-)

I think it is used when `past_key_values` is provided, and this can happen in 2 cases:
  - causal LM (without cross-attention)
  - decoder in an encoder-decoder setting (with cross-attention)

Hope I understand your words correctly.