Cannot trace_module on models using model's generate function
Author: vikigenius
Pull request: no
Labels: name: wontfix

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.3.1
- Platform: Linux-5.8.14_1-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
@patrickvonplaten, @sshleifer
 -->

@patrickvonplaten, @sshleifer

## Information

Model I am using BART

The problem arises when using:
* [*] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [*] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. load any model that uses the generate function
2. try to trace it using trace_module

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

Can be easily reproduced with the following snippet:

```
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = 'sshleifer/bart-tiny-random'
tokenizer = AutoTokenizer.from_pretrained(model)
sqgen_model = AutoModelForSeq2SeqLM.from_pretrained(model, torchscript=True)
sqgen_model.eval()
dummy_input = ' '.join('dummy' for dummy in range(512))
batch = tokenizer(
    [dummy_input], return_tensors='pt', truncation=True, padding='longest',
)
with torch.no_grad():
    traced_model = torch.jit.trace_module(  # type: ignore
        sqgen_model,
        {
            'forward': (batch.input_ids, batch.attention_mask),
            'generate': (batch.input_ids, batch.attention_mask),
        },
    )

```

It throws an error:

```
  File "/home/void/.miniconda3/envs/lexml/src/transformers/src/transformers/generation_utils.py", line 288, in generate
    assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
AssertionError: `max_length` should be a strictly positive integer.
```

obviously because the generate function's second argument is supposed to be max_length and not attention_mask

## Expected behavior
Should be able to trace models that use the generate function.

<!-- A clear and concise description of what you would expect to happen. -->


Comments:
`generate` currently does not support `torch.jit.trace`. This is sadly also not on the short-term roadmap.

@patrickvonplaten, thanks for the response. In that case is there no way to trace the inference process of generative models provided here ? So any kind of inference of the form text -> text (for eg: summarization) cannot be exported to torchscript ? 

Yes, if you want to make it faster, reduce `num_beams`.

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.


Hi, are summarization kinds of models still not traceable ??

I am trying to deploy this onto AWS inferentia, whose prerequisite is that the model should be traceable !!

for example this sshleifer/distilbart-cnn-12-6

@sshleifer, @patrickvonplaten can you guys please help ??

Hey @DevBey,

Could you maybe open a new issue that states exactly what doesn't work in your case? This issue is quite old now and it would be nice to have a reproducible code snippet with the current `transformers` version.

Do we know if trace_module is supported by generate? If so, are there any examples?