Problems with MTEncDecModel model in ONNX
State: closed | Pull request: no
Labels: name: bug

There are two issues I am having with exporting [MTEncDecModel](https://github.com/NVIDIA/NeMo/blob/r1.2.0/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py) to ONNX, they might be related in some way so going to post them together.

## 1. ONNX error in embedding in nemo version > 1.2:

If I try to run below example in nemo 1.5 I get the following error:

```
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : 
Non-zero status code returned while running Reshape node. Name:'Reshape_67' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:41
onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, std::vector<long int>&, bool)
gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. 
The input tensor cannot be reshaped to the requested shape. Input shape:{1,44,1024}, requested shape:{2,16,16,64}
```
However, it seems like this error is not present if using an older version of nemo.

## 2. Big performance discrepancy when using ONNX version of the model:
I notice a big performance drop when exporting the MTEncDecModel to ONNX. I am trying a minimal setup with greedy search and get about half the BLEU score in comparison to using MTEncDecModel directly with `TopKSequenceGenerator(top_k=1)`.

Here is a full script to get the translation with ONNX, you would need to additionally install  `onnxruntime-gpu`:
```python
import torch
import numpy as np
import onnxruntime
from nemo.collections.nlp.models import MTEncDecModel

# Load the model form pretrained:
model = MTEncDecModel.from_pretrained('nmt_en_de_transformer12x2')

# Export all model components to onnx:
model.encoder.export('encoder.onnx')
model.decoder.export('decoder.onnx')
model.log_softmax.export('classifier.onnx')

# Initialise all the onnx sessions:
encoder_session = onnxruntime.InferenceSession('encoder.onnx', providers=['CUDAExecutionProvider'])
decoder_session = onnxruntime.InferenceSession('decoder.onnx', providers=['CUDAExecutionProvider'])
classifier_session = onnxruntime.InferenceSession('classifier.onnx', providers=['CUDAExecutionProvider'])

# Preprocess the data using the original nemo model for simplicity:
TEXT = ['They are not even 100 metres apart: On Tuesday, the new B 33 pedestrian lights in Dorfparkplatz in Gutach became operational - within view of the existing Town Hall traffic lights.']
src_ids, src_mask = model.prepare_inference_batch(TEXT)
src_ids = src_ids.cpu().numpy()  # Convert to numpy for use with onnx
src_mask = src_mask.cpu().numpy().astype(int)

# Compute encoder hidden state:
encoder_input = {'input_ids': src_ids, 'encoder_mask': src_mask}
encoder_hidden_state = encoder_session.run(['last_hidden_states'], encoder_input)[0]

# Simple greedy search:
MAX_GENERATION_DELTA = 5
BOS = model.encoder_tokenizer.bos_id
EOS = model.encoder_tokenizer.eos_id
PAD = model.encoder_tokenizer.pad_id

def decode(tgt: np.array, embeding: np.array, src_mask: np.array) -> np.array:
      decoder_input = {
          'input_ids': tgt,
          'decoder_mask': (tgt != PAD).astype(np.int64),
          'encoder_mask': embeding,
          'encoder_embeddings': src_mask
      }
      decoder_hidden_state = decoder_session.run(['last_hidden_states'], decoder_input)[0]
      log_probs = classifier_session.run(['log_probs'], {'hidden_states': decoder_hidden_state})[0]
      return log_probs

max_out_len = encoder_hidden_state.shape[1] + MAX_GENERATION_DELTA
tgt=np.full(shape=encoder_hidden_state.shape[:-1], fill_value=0)
tgt[:, 0] = BOS

for i in range(1, max_out_len):
    log_probs = decode(tgt[:, :i], encoder_hidden_state, src_mask)
    next_tokens = log_probs[:, -1].argmax(axis=1) # NOTE: ONNX decoder returns multiple outputs which is different to pytorch version, so I get the last one (this could be where error is?)
    tgt[:, i] = next_tokens
    if ((tgt == EOS).sum(axis=1) > 0).all():
        break

tgt_torch = torch.from_numpy(tgt).to('cuda:0')
onnx_translation = model.ids_to_postprocessed_text(tgt_torch, model.decoder_tokenizer, model.target_processor)
```

I have run above against newstest2014 set and got a BLUE score of 13 vs 29 if I use `MTEncDecModel.translate` method. It seems like onnx model works well for shorter sentences, but for the longer ones it cuts off too soon. Could be because torch model uses decoder_mems, and onnx doesn't? 

Is there maybe a better way to set up onnx inference? 

**Environment overview (please complete the following information)**
- Environment Location: Docker on Kubernetes
- Method of NeMo install: using helm chart using nvcr.io/nvidia/nemo:1.2.0 / 1.5.1
- Additionally install `pip install onnxruntime-gpu==1.10.0`