Getting exception to trace t5 model in torchScript
Author: dhrubo-os
Pull request: no

### System Info

```
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large", torchscript=True)
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", torchscript=True)

tokenized_dict = tokenizer(
    ["please answer the following question: what is the boiling point of nitrogen",], ["-320.4F",], 
    return_tensors="pt"
)
input_tuple = (tokenized_dict['input_ids'], tokenized_dict['attention_mask'])

traced_model = torch.jit.trace(model, input_tuple)
torch.jit.save(traced_model, "flan-t5-large.pt")
```

I was trying to trace `google/flan-t5-large` model in torchScript. But I'm facing following exception:

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [29], in <cell line: 13>()
      7 tokenized_dict = tokenizer(
      8     ["please answer the following question: what is the boiling point of nitrogen",], ["-320.4F",], 
      9     return_tensors="pt"
     10 )
     11 input_tuple = (tokenized_dict['input_ids'], tokenized_dict['attention_mask'])
---> 13 traced_model = torch.jit.trace(model, input_tuple)
     14 torch.jit.save(traced_model, "flan-t5-large.pt")

File ~/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py:759, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    756     return func
    758 if isinstance(func, torch.nn.Module):
--> 759     return trace_module(
    760         func,
    761         {"forward": example_inputs},
    762         None,
    763         check_trace,
    764         wrap_check_inputs(check_inputs),
    765         check_tolerance,
    766         strict,
    767         _force_outplace,
    768         _module_class,
    769     )
    771 if (
    772     hasattr(func, "__self__")
    773     and isinstance(func.__self__, torch.nn.Module)
    774     and func.__name__ == "forward"
    775 ):
    776     return trace_module(
    777         func.__self__,
    778         {"forward": example_inputs},
   (...)
    785         _module_class,
    786     )

File ~/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py:976, in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    972     argument_names = get_callable_argument_names(func)
    974 example_inputs = make_tuple(example_inputs)
--> 976 module._c._create_method_from_trace(
    977     method_name,
    978     func,
    979     example_inputs,
    980     var_lookup_fn,
    981     strict,
    982     _force_outplace,
    983     argument_names,
    984 )
    985 check_trace_method = module._c._get_method(method_name)
    987 # Check the trace against new traces created from user-specified inputs

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/t5/modeling_t5.py:1660, in T5ForConditionalGeneration.forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1657         decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device)
   1659 # Decode
-> 1660 decoder_outputs = self.decoder(
   1661     input_ids=decoder_input_ids,
   1662     attention_mask=decoder_attention_mask,
   1663     inputs_embeds=decoder_inputs_embeds,
   1664     past_key_values=past_key_values,
   1665     encoder_hidden_states=hidden_states,
   1666     encoder_attention_mask=attention_mask,
   1667     head_mask=decoder_head_mask,
   1668     cross_attn_head_mask=cross_attn_head_mask,
   1669     use_cache=use_cache,
   1670     output_attentions=output_attentions,
   1671     output_hidden_states=output_hidden_states,
   1672     return_dict=return_dict,
   1673 )
   1675 sequence_output = decoder_outputs[0]
   1677 # Set device for model parallelism

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/t5/modeling_t5.py:949, in T5Stack.forward(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, inputs_embeds, head_mask, cross_attn_head_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    947 else:
    948     err_msg_prefix = "decoder_" if self.is_decoder else ""
--> 949     raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds")
    951 if inputs_embeds is None:
    952     assert self.embed_tokens is not None, "You have to initialize the model with valid token embeddings"

ValueError: You have to specify either decoder_input_ids or decoder_inputs_embeds
```


I also tried following way:

```
from transformers import T5ForConditionalGeneration
import torch

tokens_tensor = torch.ones(1, 10, dtype=torch.long)
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", torchscript=True)
model.eval()
scripted_model = torch.jit.trace(model, (tokens_tensor, tokens_tensor))
torch.jit.save(traced_model, "flan-t5-large.pt")
```


But this giving me following error:

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [34], in <cell line: 7>()
      5 model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", torchscript=True)
      6 model.eval()
----> 7 scripted_model = torch.jit.trace(model, (tokens_tensor, tokens_tensor))
      8 torch.jit.save(traced_model, "flan-t5-large.pt")

File ~/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py:759, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    756     return func
    758 if isinstance(func, torch.nn.Module):
--> 759     return trace_module(
    760         func,
    761         {"forward": example_inputs},
    762         None,
    763         check_trace,
    764         wrap_check_inputs(check_inputs),
    765         check_tolerance,
    766         strict,
    767         _force_outplace,
    768         _module_class,
    769     )
    771 if (
    772     hasattr(func, "__self__")
    773     and isinstance(func.__self__, torch.nn.Module)
    774     and func.__name__ == "forward"
    775 ):
    776     return trace_module(
    777         func.__self__,
    778         {"forward": example_inputs},
   (...)
    785         _module_class,
    786     )

File ~/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py:976, in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    972     argument_names = get_callable_argument_names(func)
    974 example_inputs = make_tuple(example_inputs)
--> 976 module._c._create_method_from_trace(
    977     method_name,
    978     func,
    979     example_inputs,
    980     var_lookup_fn,
    981     strict,
    982     _force_outplace,
    983     argument_names,
    984 )
    985 check_trace_method = module._c._get_method(method_name)
    987 # Check the trace against new traces created from user-specified inputs

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/t5/modeling_t5.py:1660, in T5ForConditionalGeneration.forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1657         decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device)
   1659 # Decode
-> 1660 decoder_outputs = self.decoder(
   1661     input_ids=decoder_input_ids,
   1662     attention_mask=decoder_attention_mask,
   1663     inputs_embeds=decoder_inputs_embeds,
   1664     past_key_values=past_key_values,
   1665     encoder_hidden_states=hidden_states,
   1666     encoder_attention_mask=attention_mask,
   1667     head_mask=decoder_head_mask,
   1668     cross_attn_head_mask=cross_attn_head_mask,
   1669     use_cache=use_cache,
   1670     output_attentions=output_attentions,
   1671     output_hidden_states=output_hidden_states,
   1672     return_dict=return_dict,
   1673 )
   1675 sequence_output = decoder_outputs[0]
   1677 # Set device for model parallelism

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/t5/modeling_t5.py:949, in T5Stack.forward(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, inputs_embeds, head_mask, cross_attn_head_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    947 else:
    948     err_msg_prefix = "decoder_" if self.is_decoder else ""
--> 949     raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds")
    951 if inputs_embeds is None:
    952     assert self.embed_tokens is not None, "You have to initialize the model with valid token embeddings"

ValueError: You have to specify either decoder_input_ids or decoder_inputs_embeds

```





How should I trace t5 model? Can you provide any example? Thanks


### Who can help?

@ArthurZucker @patric

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large", torchscript=True)
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", torchscript=True)

tokenized_dict = tokenizer(
    ["please answer the following question: what is the boiling point of nitrogen",], ["-320.4F",], 
    return_tensors="pt"
)
input_tuple = (tokenized_dict['input_ids'], tokenized_dict['attention_mask'])

traced_model = torch.jit.trace(model, input_tuple)
torch.jit.save(traced_model, "flan-t5-large.pt")
```

### Expected behavior

I should be able to trace the model in torchscript file.

Comments:
Hey! As the error mentions, you have to provided `decoder_input_ids` . 
The following works:
```python 
input_tuple = (tokenized_dict['input_ids'], tokenized_dict['attention_mask'], torch.Tensor([[2]]).long())

traced_model = torch.jit.trace(model, input_tuple)
torch.jit.save(traced_model, "flan-t5-small.pt")
```

Hi @ArthurZucker ,

Thanks for your reply. I'm able to trace the model now. But how can I load the model back and get the prediction from the traced model?

I tried to use this:

```
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

model = torch.jit.load("flan-t5-large.pt")

model.eval()


tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
t_input = "translate English to French: The universe is a dark forest."
token = tokenizer(t_input, return_tensors="pt")

tokens = model.generate(
    input_ids=token["input_ids"],
    attention_mask=token["attention_mask"],
    decoder_input_ids=token["input_ids"],
)

print(tokens)

output = tokenizer.decode(tokens[0].squeeze(), skip_special_tokens=True)
print(output)
```

Then I'm seeing error like:

```
Traceback (most recent call last):
  File "/Volumes/workplace/opensearch-py-ml/src/opensearch-py-ml/test1.py", line 13, in <module>
    tokens = model.generate(
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/torch/jit/_script.py", line 785, in __getattr__
    return super(RecursiveScriptModule, self).__getattr__(attr)
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/torch/jit/_script.py", line 502, in __getattr__
    return super(ScriptModule, self).__getattr__(attr)
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py", line 1269, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'RecursiveScriptModule' object has no attribute 'generate'
```

I tried with the `forward` method too:

```
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

model = torch.jit.load("flan-t5-large.pt")

model.eval()


tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
t_input = "translate English to French: The universe is a dark forest."
token = tokenizer(t_input, return_tensors="pt")

tokens = model.forward(
    input_ids=token["input_ids"],
    attention_mask=token["attention_mask"],
    decoder_input_ids=token["input_ids"],
)

print(tokens)

output = tokenizer.decode(tokens[0].squeeze(), skip_special_tokens=True)
print(output)
```

Then I'm facing error like:

```
Traceback (most recent call last):
  File "/Volumes/workplace/opensearch-py-ml/src/opensearch-py-ml/test1.py", line 21, in <module>
    output = tokenizer.decode(tokens[0].squeeze(), skip_special_tokens=True)
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/transformers/tokenization_utils_base.py", line 3471, in decode
    return self._decode(
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/transformers/tokenization_utils.py", line 931, in _decode
    filtered_tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens)
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/transformers/tokenization_utils.py", line 906, in convert_ids_to_tokens
    index = int(index)
TypeError: int() argument must be a string, a bytes-like object or a number, not 'list'
```



1. Not entirely sure that the way to script the generate function is correct as the error mentions. I suspect that only the forward path is supported. I am guessing that you should try something like `torch.jit.trace(model.generate,...)` but pinging @gante here as I am not very familiar with our current jit support
2. You are not decoding correctly. The output of the model are not individual tokens but `logits` which is a distribution of probability over what the next token is. This means that you first have to extract the argmax, and then decode the index.

Hey @dhrubo-os 👋 

`model.generate` is not fully exportable with `torch.jit`, but the model forward pass is. We have just added it to our examples, the workaround is to create an ad hoc model class with the jitted model + the generate function -- see [here](https://github.com/huggingface/transformers/blob/5fd4e3c87c685fba2dd9615be62131748a8b5ee3/examples/pytorch/text-generation/run_generation.py#L384)

I hope this helps 🤗 

@ArthurZucker Could you kindly give an example of how to correctly decode?

> You are not decoding correctly. The output of the model are not individual tokens but logits which is a distribution of probability over what the next token is. This means that you first have to extract the argmax, and then decode the index.

@gante Hi, I tried out the ad-hoc model class with the jitted model you mentioned [here](https://github.com/huggingface/transformers/blob/5fd4e3c87c685fba2dd9615be62131748a8b5ee3/examples/pytorch/text-generation/run_generation.py#L384) on a `AutoModelForSeq2SeqLM` for FLAN-T5. But I am getting an error. Is there any adjustment needed to the script you linked to be compatible with T5 models?

Env:
```
Optimum - 1.8.8
Transformers - 4.29.2
Torch - 1.13.1
onnxruntime-gpu - 1.15.1
onnx - 1.14.0
Linux
Python 3.8.16
```
Error:
```
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/dir/ai/models/py/summarization/modeling/flan/torchscript.py", line 79, in <module>
    outputs = fallback_model.generate(**tokenized_dict)
  File "/home/.virtualenvs/ai/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/.virtualenvs/ai/lib/python3.8/site-packages/transformers/generation/utils.py", line 1515, in generate
    return self.greedy_search(
  File "/home/.virtualenvs/ai/lib/python3.8/site-packages/transformers/generation/utils.py", line 2332, in greedy_search
    outputs = self(
  File "/home/dir/ai/models/py/summarization/modeling/flan/torchscript.py", line 26, in __call__
    outputs = self._optimized(*trace_graph_inputs)
  File "/home/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: forward() expected at most 4 argument(s) but received 5 argument(s). Declaration: forward(__torch__.transformers.models.t5.modeling_t5.___torch_mangle_2224.T5ForConditionalGeneration self, Tensor input_ids, Tensor attention_mask, Tensor decoder_input_ids) -> ((Tensor, ((Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor), (Tensor, Tensor, Tensor, Tensor)), Tensor))
```

Hi @shannonphu 👋 

Our support for torchscript is mainly hands-off, I'm afraid I don't have the bandwidth to dive deeper on bugs :) Looking at the trace, it seems like there is something wrong with the input preprocessing.

@shannonphu I'm running into something similar, what did you do to address the mismatch in number of arguments?

[Update]
Just making sure that I was clear, the class that is defined in the notebook ([readable format](https://github.com/pytorch/text/blob/bd3481896b7ac7cfcbeba43336158f841734aa70/notebooks/torchscriptable_t5_with_torchtext.ipynb)) allows to save and load a T5 model directly with `torch.jit.save` so no need for `torch.jit.trace`:
```
from torchtext.models import T5_LARGE_GENERATION

t5_large = get_jit_from_bundle(T5_LARGE_GENERATION)
model_filename = 'flan_t5_large_generation.pt'
torch.jit.save(t5_large, model_filename)
```

[Previous comment]
I have created a [notebook](https://github.com/rbahumi/text/blob/torchscriptable_t5/notebooks/torchscriptable_t5_with_torchtext.ipynb) that defines a Torchscriptable T5 class which works with torch.jit.script() directly instead of trace. 

I have a pending pull request here:
https://github.com/pytorch/text/pull/2122/commits/bd3481896b7ac7cfcbeba43336158f841734aa70

Hope that helps
@dhrubo-os @wumadeline  @gante @ArthurZucker @shannonphu