Input gets lost when converting mBART decoder to onnx
Author: tobigue
Pull request: no

I'm trying to convert the mBART decoder to onnx and have the problem, that one of the inputs gets lost during the conversion, which leads to errors when trying to use the onnx model. (See code example below.)

I'm trying to understand why this is the case and how to circumvent this.

Thanks alot for any help!

## Environment info

- `transformers` version: 4.4.2
- Platform: Linux-5.4.0-65-generic-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

## Who can help

@mfuntowicz @patil-suraj @patrickvonplaten 

## Information

Model I am using (Bert, XLNet ...): mBART

The problem arises when using:
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] my own task or dataset: (give details below)

## To reproduce

If you run the code below, you should see the following print output:

```
['input_ids', 'encoder_attention_mask', 'encoder_hidden_states']
```

Now, if we uncomment the commented line in `DecoderWithLMhead.forward` and pass the `past_key_values` to the decoder and run the code again, the additional inputs will be added, but `encoder_hidden_states` is not present as an input any longer.

If we run `torch.onnx.export` with `verbose=True`, `encoder_hidden_states` seems not to be part of the graph. Is there a condition in the mBART decoder implementation that excludes `encoder_hidden_states` from the graph, when `past_key_values` is given to the decoder?

Code to reproduce the issue (adapted from [FastT5](https://github.com/Ki6an/fastT5/blob/master/fastT5/onnx_exporter.py)):
```python
import functools
import operator
import os
import tempfile

from transformers import AutoTokenizer, MBartForConditionalGeneration, AutoConfig
from onnxruntime import InferenceSession
import torch

model_or_model_path = 'facebook/mbart-large-cc25'
model = MBartForConditionalGeneration.from_pretrained(model_or_model_path)
model_config = AutoConfig.from_pretrained(model_or_model_path)

class DecoderWithLMhead(torch.nn.Module):
    def __init__(self, decoder, lm_head, config):
        super().__init__()
        self.decoder = decoder
        self.lm_head = lm_head
        self.config = config

    def forward(self, *inputs):
        input_ids, attention_mask, encoder_hidden_states = inputs[:3]
        list_pkv = inputs[3:]
        past_key_values = tuple(list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4))
        decoder_output = self.decoder(
            input_ids=input_ids,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=attention_mask,
            # past_key_values=past_key_values,
        )
        lm_head_out = self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5))
        return lm_head_out, decoder_output[1]

decoder_with_lm_head = DecoderWithLMhead(
    decoder=model.get_decoder(), 
    lm_head=model.get_output_embeddings(),
    config=model_config
)
    
batch_size = 5
sequence_length = 10

input_ids_dec = torch.ones((batch_size, 1), dtype=torch.int64)
attention_mask_dec = torch.ones((batch_size, sequence_length), dtype=torch.int64)
enc_out = torch.ones(
    (batch_size, sequence_length, model_config.d_model), dtype=torch.float32
)
head_dim = model_config.d_model // model_config.encoder_attention_heads
a = torch.ones((batch_size, model_config.decoder_attention_heads, sequence_length, head_dim), dtype=torch.float32)
attention_block = (a, a, a, a)
past_key_values = (attention_block,) * model_config.decoder_layers
flat_past_key_values = functools.reduce(operator.iconcat, past_key_values, [])
decoder_all_inputs = tuple(
    [input_ids_dec, attention_mask_dec, enc_out] + flat_past_key_values
)
num_of_inputs = 4 * model_config.decoder_layers

with torch.no_grad():
    decoder_inputs = [
        "input_ids",
        "encoder_attention_mask",
        "encoder_hidden_states",
    ]
    pkv_input_names = ["input_{}".format(i) for i in range(0, num_of_inputs)]      
    decoder_input_names = decoder_inputs + pkv_input_names
    decoder_output_names = ["logits", "output_past_key_values"]
    dyn_axis = {
        "input_ids": {0: "batch", 1: "sequence"},
        "encoder_attention_mask": {0: "batch", 1: "sequence"},
        "encoder_hidden_states": {0: "batch", 1: "sequence"},
        "logits": {0: "batch", 1: "sequence"},
        "output_past_key_values": {0: "batch", 1: "sequence"},
    }
    dyn_pkv = {
        "input_{}".format(i): {0: "batch", 1: "n_head", 2: "seq_length", 3: "d_kv"}
        for i in range(0, num_of_inputs)
    }
    dyn_axis_params = {**dyn_axis, **dyn_pkv}

    temp_dir = tempfile.TemporaryDirectory()
    onnx_output_path = os.path.join(temp_dir.name, "decoder.onnx")

    torch.onnx.export(
        decoder_with_lm_head,
        decoder_all_inputs,
        onnx_output_path,
        export_params=True,
        do_constant_folding=True,
        opset_version=12,
        input_names=decoder_input_names,
        output_names=decoder_output_names,
        dynamic_axes=dyn_axis_params,
        use_external_data_format=True,
    )
     
session = InferenceSession(onnx_output_path)
print(list(map(lambda x: x.name, session.get_inputs())))  # encoder_hidden_states should be in here
temp_dir.cleanup()
```

## Expected behavior

All inputs passed to the onnx export function are present in the created onnx model.


Comments:
It seems that the `encoder_hidden_states` are passed as `key_value_states` into the `MBartAttention` down the line and are not used in case `past_key_value` is given.
https://github.com/huggingface/transformers/blob/90ecc29656ce37fdbe7279cf586511ed678c0cb7/src/transformers/models/mbart/modeling_mbart.py#L183
In that case I guess it's expected that they are not in the graph of the decoder, so I'll see how I can work around that when converting to ONNX.