Can we convert dynamic DNN model to TorchScript?
Author: ranggihwang
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-5.4.0-42-generic-x86_64-with-glibc2.10
- Python version: 3.8.13
- Huggingface_hub version: 0.13.4
- PyTorch version (GPU?): 2.0.1+cu117 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker
@sgugger
@gante

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hi all.

I'm trying to convert SwitchTransformer model to TorchScript.
(SwitchTransformer model is MoE DNN based on Google T5 model.)

When converting both T5 and SwitchTransforemer, there's no error for T5 but I got following error for SwitchTransformer.

```
/root/HuggingFace/.HF/lib/python3.8/site-packages/transformers/modeling_utils.py:776: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if causal_mask.shape[1] < attention_mask.shape[1]:
Traceback (most recent call last):
  File "example.py", line 423, in <module>
    traced_model = torch.jit.trace(model, (input_ids, attention_mask, decoder_input_ids))
  File "/root/HuggingFace/.HF/lib/python3.8/site-packages/torch/jit/_trace.py", line 794, in trace
    return trace_module(
  File "/root/HuggingFace/.HF/lib/python3.8/site-packages/torch/jit/_trace.py", line 1056, in trace_module
    module._c._create_method_from_trace(
RuntimeError: Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
```

I think it is because of the dynamic characteristics of SwitchTransformer.

This is the code for T5.
```python
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

tokenizer = T5Tokenizer.from_pretrained('t5-small')
model = T5ForConditionalGeneration.from_pretrained('t5-small', torchscript = True)
input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
attention_mask = input_ids.ne(model.config.pad_token_id).long()
decoder_input_ids = tokenizer('<pad> <extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids

traced_model = torch.jit.trace(model, (input_ids, attention_mask, decoder_input_ids))
torch.jit.save(traced_model, "traced_t5.pt")
```

And this is the code for SwitchTransformer.
```python
from transformers import AutoTokenizer, SwitchTransformersForConditionalGeneration
from transformers import AutoTokenizer, SwitchTransformersConfig
import torch

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "google/switch-base-8", resume_download=True)
model = SwitchTransformersForConditionalGeneration.from_pretrained(
    "google/switch-base-8",
    resume_download=True, torch_dtype=torch.bfloat16,
    torchscript=True,
)

input_text = "A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>."
output_text = "<pad> <extra_id_0> man<extra_id_1> beer<extra_id_2> a<extra_id_3> salt<extra_id_4>.
"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
decoder_input_ids = tokenizer(output_text, return_tensors="pt", padding=True).input_ids

attention_mask = input_ids.ne(model.config.pad_token_id).long()

# model.eval()

traced_model = torch.jit.trace(model, (input_ids, attention_mask, decoder_input_ids))
```

### Expected behavior

TorchScript version of SwitchTransformer.

Comments:
Hey! Thanks for reporting, I think it might also be because our Switch returns None values sometimes, checking
! 

Thanks a lot for providing the full reproduction script 🤗 
Actually, I don't think that MoE models can be torch scripted as the path taken by the inputs will be different for each tokens (because of the routing mechanism). 
However, there was a problem of returning `None` values, and not returning the router probs if labels were `None`. Fixing in #24300

Thanks, so, do you mean that the MoE model cannot be torchscripted as it has a dynamic workflow depending on its input, not because returning `None` value?

And for the `None` value issue, which one of these is the newest correct version? They are slightly different.
https://github.com/huggingface/transformers/pull/24300/files/3b899c180d4a38a06d34dcb1687626594f0497a0
https://github.com/huggingface/transformers/commit/ba3fb4b8d72b9202423cda01896349a883480d2e#diff-897fe3777ef1c9d71d6268fac217b0e163f2e20a3a5e4fabfe5a3675bc9202c7

The return value was indeed an issue, which prevent starting the tracing. But now that `None` are not returned anymore, the model still cannot be traced because of the dynamic workflow yes. 
The correct commit is the one that was merged to main! 

Thank a lot!
But then can I just use scripting for torchscript?
As far as I know, there're two optimization schemes for torchscript, tracing and scripting.
So can I just adopt only scripting selectively?

Yep! I think scripting is what you should be using for dynamic workflows! 

Well, I found something interesting.

As shown below, scripting for the T5 model also does not work.
But as shown above, tracing worked.

How does this happen?

```python
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

tokenizer = T5Tokenizer.from_pretrained('t5-small')
model = T5ForConditionalGeneration.from_pretrained('t5-small', torchscript = True)
input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
attention_mask = input_ids.ne(model.config.pad_token_id).long()
decoder_input_ids = tokenizer('<pad> <extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids

# traced_model = torch.jit.trace(model, (input_ids, attention_mask, decoder_input_ids))
scripted_model = torch.jit.script(model)
# torch.jit.save(traced_model, "traced_t5.pt")
```