Similar usage of `past_key_values` in CausalLM and Seq2SeqLM
Author: forest1988
Pull request: no
Labels: name: Good First Issue

# 🚀 Feature request

It seems GPT-2 and BartDecoder has a different style of `past_key_values`.

In GPT-2, `past_key_values` is explained as below:
(the explanation is from https://huggingface.co/transformers/model_doc/gpt2.html#gpt2model)

```
(parameters)
past_key_values (List[torch.FloatTensor] of length config.n_layers) – Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see past_key_values output below). Can be used to speed up sequential decoding. The input_ids which have their past given to this model should not be passed as input_ids as they have already been computed.


(returns)
past_key_values (List[torch.FloatTensor], optional, returned when use_cache=True is passed or when config.use_cache=True) – List of torch.FloatTensor of length config.n_layers, with each tensor of shape (2, batch_size, num_heads, sequence_length, embed_size_per_head)).

Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see past_key_values input) to speed up sequential decoding.
```

In BartDecoder and its inner BartDecoderLayer, `past_key_values` is explained and treated as below:
(the explanation is from https://huggingface.co/transformers/model_doc/bart.html#bartmodel)

``` 
past_key_values (Tuple[Tuple[torch.Tensor]] of length config.n_layers with each tuple having 2 tuples each of which has 2 tensors of shape (batch_size, num_heads, sequence_length - 1, embed_size_per_head)) –

Contains precomputed key and value hidden-states of the attention blocks. Can be used to speed up decoding.
```

`v4.1.1 modeling_bart`
https://github.com/huggingface/transformers/blob/v4.1.1/src/transformers/models/bart/modeling_bart.py

``` python
# in BartDecoder

        for idx, decoder_layer in enumerate(self.layers):
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            dropout_probability = random.uniform(0, 1)
            if self.training and (dropout_probability < self.layerdrop):
                continue

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            hidden_states, layer_self_attn, present_key_value, layer_cross_attn = decoder_layer(
                hidden_states,
                attention_mask=combined_attention_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=encoder_attention_mask,
                past_key_value=past_key_value,
                output_attentions=output_attentions,
            )
```

``` python
# in BartDecoderLayer

        # Self Attention
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        # add present self-attn cache to positions 1,2 of present_key_value tuple
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )
        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        if not self.normalize_before:
            hidden_states = self.self_attn_layer_norm(hidden_states)

        # Cross-Attention Block
        cross_attn_present_key_value = None
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states
            if self.normalize_before:
                hidden_states = self.encoder_attn_layer_norm(hidden_states)

            # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
                hidden_states=hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                past_key_value=cross_attn_past_key_value,
                output_attentions=output_attentions,
            )
            hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
            hidden_states = residual + hidden_states
            if not self.normalize_before:
                hidden_states = self.encoder_attn_layer_norm(hidden_states)

            # add cross-attn to positions 3,4 of present_key_value tuple
            present_key_value = present_key_value + cross_attn_present_key_value
```

## Motivation

It seems that one of the aims of the refactoring of Bart by @patrickvonplaten https://github.com/huggingface/transformers/pull/8900 is "Allow to use BartEncoder and BartDecoder separately from the BartModel".

I appreciate this very much and would love to treat `BartDecoder` as well as `gpt2`, but I feel that the difference in the handling of `past_key_values` is a barrier.

In `gpt2`, `past_key_value` in `past_key_values` is `torch.tensor` with each tensor of shape (2, batch_size, num_heads, sequence_length, embed_size_per_head)).
However, in `Bart`, `past_key_value` in `past_key_values` is `Tuple[torch.Tensor]` and `self_atten` part is not a tensor but a "2 tensors of shape (batch_size, num_heads, sequence_length - 1, embed_size_per_head))".

If we want to handle `self_attn_past_key_value` in `Bart` like that in `gpt2`, is it the right way to concatename the 2 tensors in `past_key_value`?
Or, is there the other correct way to treat it?

Thank you in advance.

Comments:
Hey @forest1988,

In order to fully use `BartDecoder` separately from `BartModel` as a `BartForCausalLM` model, we're still waiting on this PR: #9128.
And again, you're very correct in your assessment that the behavior between `BartDecoder` and `GPT2` is not fully aligned here. IMO, we should change GPT2's cache format from a single tensor of `[2, batch_size, ...,]` to `tuple([batch_size, ...])` => If you're keen feel free to open a PR for it! Actually, we could make this a "Good second issue" here. 

So to answer your question, no you should not contact the 2 tensors in `self_attn_past_key_value` to a single tensor, but we should rather change the code in GPT2 slightly to also have a tuple of 2 tensors instead of one tensor. 
In GPT2, we create a new tensor at each iteration when using `use_cache` here: https://github.com/huggingface/transformers/blob/d944966b19a4d6860bddc7cdc1ba928ca8a0da91/src/transformers/models/gpt2/modeling_gpt2.py#L235 => this is a bit unnecessary IMO. When the inputs are getting longer allocating new memory for `key` and `value` can actually lead to a small slow-down IMO. If instead we would just use a tuple => `present = (key, value)` we would not allocate new memory. 

So 1) As soon as #9128 is merged you can use `BartForCausalLM` the same way as `GPT2` without having to change anything.
2) Let's see if someone is interested in tackling this "inconsistency" issue in GPT2. This "First good issue" should replace this line: https://github.com/huggingface/transformers/blob/d944966b19a4d6860bddc7cdc1ba928ca8a0da91/src/transformers/models/gpt2/modeling_gpt2.py#L235
with 
```python
present = (key.transpose(-2, -1), value))
```
(I think it should actually be that simple)


Hi @patrickvonplaten, 

Thank you for answering this issue!
I'm sorry I haven't checked the PR https://github.com/huggingface/transformers/pull/9128 before creating this issue. I'll check it!

And, thanks for telling me your opinion about the need to change GPT2's cache format from a single tensor to a tuple of 2 tensors.
I'd love to open a PR, but I'm afraid I don't have enough time now.
I will work on it as soon as I find the time, but of course, if someone else who is interested in the same issue would like to work on it, I would appreciate that!

I now understand that I should not contact the 2 tensors in `self_attn_past_key_value` to a single tensor, but rather the code in GPT2 should be changed.

I'm looking forward to seeing the PR https://github.com/huggingface/transformers/pull/9128 is merged.
Also, I would like to think about how to avoid contacting the 2 tensors in `self_attn_past_key_value` for what I am currently working on.

Thank you so much!


Hi,

I've started fixing the issue but failed in some tests.

```
========================================================================================= short test summary info ==========================================================================================
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_sample_generate - AttributeError: 'tuple' object has no attribute 'index_select'
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_generate - AttributeError: 'tuple' object has no attribute 'index_select'
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_gpt2_gradient_checkpointing - TypeError: CheckpointFunctionBackward.forward: expected Tensor or tuple of Tensor (got tuple) for return value 1
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_group_beam_search_generate - AttributeError: 'tuple' object has no attribute 'index_select'
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_model_parallel_beam_search - AttributeError: 'tuple' object has no attribute 'index_select'
FAILED tests/test_trainer_distributed.py::TestTrainerDistributed::test_trainer - OSError: [Errno 12] Cannot allocate memory
================================================================== 6 failed, 4626 passed, 834 skipped, 642 warnings in 1766.16s (0:29:26) ==================================================================
```

I think this may be related to the `index_select` used in [generation_utils.py](
https://github.com/huggingface/transformers/blob/143289dcf759a663c03317e30167e89ee6d86588/src/transformers/generation_utils.py).

I will continue to look into this when I have more time.

Hi @patrickvonplaten,

Now that I have time, I'm thinking about how to fix this issue so that the testing part works well.

I think where to modify in `generation_utils.py` is here:

https://github.com/huggingface/transformers/blob/02e05fb0a532e572b56ba75dad6ba3db625bbdeb/src/transformers/generation_utils.py#L506-L517

`past` is taken from `past_key_values` or other output variations in:
https://github.com/huggingface/transformers/blob/02e05fb0a532e572b56ba75dad6ba3db625bbdeb/src/transformers/generation_utils.py#L477-L489

Then, it is treated like this:
https://github.com/huggingface/transformers/blob/02e05fb0a532e572b56ba75dad6ba3db625bbdeb/src/transformers/generation_utils.py#L1654-L1655

Can I modify #L506-L517 so that `past` is replaced from `Tuple[torch.Tensor]` to `Tuple[Tuple[torch.Tensor]]`,
or should I consider other output variations, `output.mem` and `outputs.past_buckets_states`?

Thank you.


Hey @forest1988,

Thanks for you in-detail code snippet. I think the easiest solution would be to open a PR showcasing the required changes :-) I think you're right `Tuple[torch.Tensor]` should indeed be `Tuple[Tuple[torch.Tensor]]`.

Would you be interested in opening a PR so that we can add the fixes there? 

Hi @patrickvonplaten,

Thank you for your comment! 
After doing some additions to change `Tuple[torch.Tensor]` to `Tuple[Tuple[torch.Tensor]]`, I would like to open a PR and ask you all to add fixes.
I'll open the PR in a few days!

I'm sorry to keep you waiting. I have opened a PR #9596 for this issue.
I marked the PR as WIP because it has not yet been resolved.
I will continue to look into this issue myself in the future and any advice would be greatly appreciated.