Cannot trace_module on models using model's generate function

Environment info

• 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

@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

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.

[patrickvonplaten]

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

[patrickvonplaten]

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.