ONNX T5 with Beam Search
Author: amanpreet692
Pull request: no
Labels: name: wontfix

Hey guys,
I didn't know where this belonged so opening up a generic issue.
I was working on integrating the ONNX T5 code by @abelriboulot with the HuggingFace Beam Search decoding code since I already had a decently performing T5 model for summarization and wanted to improve performance on CPU while maintaining the inference accuracy.
It works for the most part, but is slower as the HF code uses cached past state values to speed up the decoding. I got around this issue by creating two decoders with lm-head, one which doesn't take in past values for the initial decoding and another for subsequent steps where past values are considered. This is a bit complicated as the past values have to be flattened out to pass through the ONNX graph which I did and it works for getting back the output.
But for passing the input parameters, I get the following error:
**RUNTIME_EXCEPTION : Non-zero status code returned while running Mul node. Name:'Mul_48' Status Message: /Users/runner/work/1/s/onnxruntime/core/providers/cpu/math/element_wise_ops.h:479 void onnxruntime::BroadcastIterator::Init(int64_t, int64_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 2 by 3**

I feel like I am close to the solution which could essentially be added to the repo but this error is tripping me up :(
Any help whatsoever will be appreciated.
Thanks
@mfuntowicz @abelriboulot @patrickvonplaten @patil-suraj @sshleifer 

ONNX Export code:
```python
    past_state_input_pre = torch.rand((1,12,1,64))
    past_state_input_post = torch.rand((1, 12, 10, 64))
    past_key_value_states = [(past_state_input_pre, past_state_input_pre, past_state_input_post, past_state_input_post) for i in range(12)]

    past_val_outputs = {'past_states_op_'+str(i): {0:'batch', 2: 'sequence'} for i in range(48)}
    past_val_inputs = {'past_states_ip' + str(i): {0: 'batch', 2: 'sequence'} for i in range(48)}
    dynamix_axes_dict = {
                              'input_ids': {0:'batch', 1: 'sequence'},
                              'encoder_hidden_states': {0:'batch', 1: 'sequence'}
                            }
    dynamix_axes_dict.update(past_val_inputs)
    dynamix_axes_dict.update({'hidden_states': {0:'batch', 1: 'sequence'}})
    dynamix_axes_dict.update(past_val_outputs)
    output_names_list = ['hidden_states'] + ['past_states_op_' + str(i) for i in range(48)]
    input_names_list = ['input_ids', 'encoder_hidden_states'] + ['past_states_ip' + str(i) for i in range(48)]
    # Exports to ONNX
    _ = torch.onnx.export(
                            decoder_with_lm_head,
                            (torch.tensor([[42]]), simplified_encoder(input_ids), past_key_value_states),
                                   f"{output_prefix}-decoder-with-lm-head.onnx",
                                   export_params=True,
                            opset_version=12,
                            input_names=input_names_list,
                            output_names=output_names_list,
                            dynamic_axes= dynamix_axes_dict)

Comments:
Hi @amanpreet692 ,
I'm not sure what this error means, but I've `T5` `onnx` version ready which is compatible with `generate` method.
To be able to use cache I exported the `encoder` and `lm_head` to `onnx` and kept the `decoder` in `torch`. This is bit hacky but still gives 1.4-1.6x speed-up for beam search, I'll be sharing it soon.

Yep, Even I was able to do that but since majority of the time is taken while decoding I wanted to convert decoder as well! Will keep trying for now.

@patil-suraj A question, somehow converting both lm-head and encoder is giving me worse result as compared to only converting the encoder. Did you try any additional optimizations like quantization?

No, I didn't try quantization with T5, so far I'm getting good enough speed-up and results are same as that of torch.

Not related to your `onnx` question, but you could also distill the models, to get additional speed-ups with minimal perf drop. Sam has just relased amazing s2s distillation [paper](https://arxiv.org/pdf/2010.13002.pdf). See if that helps you with speeding-up inference.

Hey @amanpreet692!
Thanks a lot for looking at making the ONNX version compatible with beam search. Could you send over your full script to make it easier to debug? Happy to hop on a call this week and hear a bit more what you have in mind. The two decoders solution sounds interesting!

I've posted the script on the [forum ](https://discuss.huggingface.co/t/speeding-up-t5-inference/1841).

@abelriboulot Thanks a lot for getting back :)
Here are the scripts for my work (The first two are changes on top of your code and the third is my custom model with 2 decoders):
1) [huggingface_utilities.py](https://gist.github.com/amanpreet692/41dba767220b5b1a6417066197781328) : Additional changes to include past states as input and output and convert 3 components (2 decoders, 1 encoder) into onnx format.
2) [models.py](https://gist.github.com/amanpreet692/d36af959e0d8d9cf84b19ff26d9b19d8) : Smallish change to include a new class CombinedDecoderNoPast 
3) [t5_onnx_model.py](https://gist.github.com/amanpreet692/a8bf2d45a8f368830f3838790461d26b) : Complete T5 model that works with beam search, major changes in decoder processing.

Just an update: I was able to resolve to above issue but started getting a new shape issue for buffers, have raised an issue on onnx repo as well: [ONNX Issue](https://github.com/microsoft/onnxruntime/issues/5646)

Any pointers for debugging would be great, and sure it would be awesome if we can get on a call and work on this!! 
Will keep trying on my own till then.

@patil-suraj Good job! I looked at your code and I had tried something very similar. Although am still skeptical as I was getting worse performance with converting both encoder and lm-head rather than only encoder. Will look at your results again.

Thanks again @abelriboulot !

as long as we pass same arguments to `generate`  then we should get same results, I didn't observe any loss in accuracy.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.