How to go about utilizing MBART for conditional generation with beam search in ONNXRuntime with TensorRT/CUDA
Author: JeroendenBoef
Pull request: no

Hi HuggingFace team,

Last December I looked into exporting `MBartForConditionalGeneration` from `"facebook/mbart-large-50-many-to-one-mmt"` for the purpose of multilingual machine translation. Originally I followed the approach as described in this [BART + beam search example](https://github.com/huggingface/transformers/tree/master/examples/research_projects/onnx/summarization), extending the example to support MBART and overriding the max 2GB model size. While this approach worked for `CPUExecutionProvider` in ORT sessions, it did not actually improve runtime, nor did it work for `TensorRT` or `CUDA` execution providers (out of cuda memory and dynamic shape inference failure).

Today I saw [this issue ](https://github.com/huggingface/transformers/issues/15716) and exported `MBartForConditionalGeneration` with `python -m transformers.onnx --model=facebook/mbart-large-50-many-to-one-mmt --feature seq2seq-lm-with-past --atol=5e-5 onnx/`. While this worked for exporting to ONNX (passing all validation checks), I couldn't run an actual ORT session due to input dimensionality mismatch (past keys encoder/decoder missing for `seq2seq-lm-with-past`, `decoder_inputs_ids` and `decoder_attention_mask` missing for `seq2seq-lm`).

I could use some clarification as to whether this is the implementation I am looking for (does the latter ONNX export support `.generate()` through beam search or should I refocus my attempts at the BART + beam search modification). In case the newer command line ONNX export implementation is what I require, which feature head would be the correct head for the ConditionalGeneration many-to-one-mmt MBART head (`seq2seq-lm` or `seq2seq-lm-with-past`) and where can I find the additional inputs that I need for the model to run `.generate()` in an ORT session? The BART beam search implementation I mentioned earlier required `input_ids`, `attention_mask`, `num_beams`, `max_length` and `decoder_start_token_id`. The required inputs for the newer conversion are a bit more confusing to me.

I assume @lewtun would be the person to ask for help here but I appreciate any pointers!

## Environment info
- `transformers` version: 4.16.2
- Platform: Linux-5.4.0-96-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.10.1+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?:  yes
- Using distributed or parallel set-up in script?: no


Comments:
I was also adjusting the BART export to work for mBART in a translation setting. There you need to be able to set the `decoder_start_token_id` dynamically. I noticed that the ONNX conversion example you link to (thank you!) is not complete. It is missing the `dynamic_axes` configuration for the `attention_mask`. It needs to be adjusted to:
```
        torch.onnx.export(
            bart_script_model,
            (
                inputs["input_ids"],
                inputs["attention_mask"],
                num_beams,
                max_length,
                decoder_start_token_id,
            ),
            onnx_file_path,
            opset_version=14,
            input_names=[
                "input_ids",
                "attention_mask",
                "num_beams",
                "max_length",
                "decoder_start_token_id",
            ],
            output_names=["output_ids"],
            dynamic_axes={
                "input_ids": {0: "batch", 1: "seq"},
                "attention_mask": {0: "batch", 1: "seq"},
                "output_ids": {0: "batch", 1: "seq_out"},
            },
            example_outputs=output_ids,
        )
```

Hope this helps!

Thanks for jumping in with the tip @HaukurPall! Did this approach allow you to not only export the mBART model to ONNX but also run it (with increased speed) on CUDA/TensorRT execution providers? Moreover, I assume you made a custom version of `convert.py`. where you overrode the `torch.onnx.export()` call with the snippet you posted above, is this correct? Thanks in advance!

Thanks for opening an issue @JeroendenBoef!

Pinging @lewtun, @mfuntowicz, should this issue be moved to optimum?

Thanks for the ping! Yes, I think it would make sense to move this issue to the `optimum` repo :)

@JeroendenBoef we currently have a PR in `optimum` that will enable simple inference / text-generation for ONNX Runtime:  https://github.com/huggingface/optimum/pull/113

Once that is completed, I think it should address most of the points raised in this issue!

Hey @JeroendenBoef. No, I was not able to get the inference working efficiently on the CUDA execution providers. I even attempted to use the IOBindings (as suggested by the ONNX team) but was not successful. I have put this endeavour aside until there is better support for autoregressive inference.

If you still want to try this, there is a different approach to exporting the models presented in https://github.com/Ki6an/fastT5/. This is for T5, but some of the code has been adjusted to work for mBART, see issue: https://github.com/Ki6an/fastT5/issues/7. I did not try to running that model on CUDA as it would require some work getting the IOBindings correct/efficient.

Thanks for the reply and the pointer to the new PR on `optimum` @lewtun. I will close this issue for now and keep an eye out for the new developments regarding seq2seq models on `optimum`. I feel like [this](https://github.com/huggingface/optimum/issues/55) issue already covers the core of this issue, save for maybe the problems with actually achieving improved inference speed with the exported model, so unless this is preferred for documentation, I would not open a new issue on `optimum`.

Thanks for the detailed response @HaukurPall, this saves me some headaches and time :). I was already afraid there would not be an improved performance but now I have confirmation that I should also postpone my efforts on this until there is a better approach in place for ORT seq2seq models.