

TITLE
Similar usage of `past_key_values` in CausalLM and Seq2SeqLM

COMMENTS
7

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 🚀 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.