Converting a tensor to a Python boolean might cause the trace to be incorrect when converting gpt2 to onnx format
Author: HaoboGu
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.18.0
- Platform: Darwin-21.5.0-x86_64-i386-64bit
- Python version: 3.7.9
- Huggingface_hub version: 0.7.0
- PyTorch version (GPU?): 1.11.0 (False)
- 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?

@patil-suraj 
@patrickvonplaten 
@LysandreJik 

### Information

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

### Tasks

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

### Reproduction

I'm trying to convert transformers GPT2LMHeadModel to onnx format using `torch.onnx.export`

I got the following warning: 
```
/Users/Project/.venv/lib/python3.7/site-packages/transformers/models/gpt2/modeling_gpt2.py:797: 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 batch_size <= 0:
```

I tried to print `batch_size` here, and I got `tensor(3)`(which should be `3` here?)


### Expected behavior

```shell
No warnings when converting
```


Comments:
@lewtun for ONNX here

@HaoboGu before we can answer the issue could you please add a complete reproducible code snippet here?

@patrickvonplaten 

Sure

```python
from transformers import GPT2Config
from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
import torch


def convert_gpt2_model_to_onnx() -> str:
    """
    convert pytorch model to onnx format
    """
    # Load model and set the model to eval mode
    model: GPT2LMHeadModel = GPT2LMHeadModel.from_pretrained('sshleifer/tiny-gpt2')
    model.eval()

    # batch_size, input_ids_length and past_sequence_length are dynamic axes
    # We have to initialize a random input(the value doesn't matter) for the model, because the converting requires execution of the model
    # See: https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html
    batch_size = 3
    input_ids_length = 2
    past_sequence_length = 1
    config: GPT2Config = model.config
    num_attention_heads = config.n_head
    hidden_size = config.n_embd
    num_layer = config.n_layer
    vocab_size = config.vocab_size
    config.is_decoder = True

    # past(`past_key_values` in model) is a list, its length is num_layer.
    # each element in the list is a tuple(key, value), and key/value's shape is past_shape,
    # aka [batch_size, n_heads, past_sequence_length, embd_size_each_head]
    past_shape = [batch_size, num_attention_heads,
                  past_sequence_length, int(hidden_size/num_attention_heads)]
    past = [(torch.rand(past_shape, dtype=torch.float32, device='cpu'), torch.rand(past_shape, dtype=torch.float32, device='cpu'))
            for _ in range(num_layer)]

    # input_ids is a [batch_length, input_ids_length] tensor
    input_ids = torch.randint(
        low=0,
        high=vocab_size - 1,
        size=(batch_size, input_ids_length),
        dtype=torch.long,
        device='cpu',
    )

    # attention_mask is a 0/1 tensor of [batch_size, past_sequence_length + input_ids_length]
    attention_mask = torch.ones(
        [batch_size, past_sequence_length + input_ids_length]).to(torch.long)

    # token_type_ids is not needed in our case, its size is [batch_size, input_ids_length]
    token_type_ids = torch.zeros([batch_size, input_ids_length]).to(torch.long)

    # position_ids, size is [batch_size, input_ids_length]
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(position_ids < 0, 0)
    position_ids = position_ids[:, past_sequence_length:].to(torch.long)

    # Run the model and get output from the model
    output = model(input_ids, past_key_values=past, attention_mask=attention_mask,
                   position_ids=position_ids, return_dict=True, use_cache=True)

    # Set output names
    output_names = ['logits']
    for i in range(num_layer):
        output_names.append("present_key" + str(i))
        output_names.append("present_value" + str(i))

    # Set input_names
    input_names = ['input_ids']
    for i in range(num_layer):
        input_names.append("past_key" + str(i))
        input_names.append("past_value" + str(i))
    input_names += ['attention_mask', 'position_ids']

    # Set dynamic axes
    dynamic_axes = {}
    dynamic_axes['input_ids'] = {0: 'batch_size', 1: 'input_ids_length'}
    dynamic_axes['attention_mask'] = {0: 'batch_size', 1: 'total_length'}
    dynamic_axes['position_ids'] = {0: 'batch_size', 1: 'input_ids_length'}
    dynamic_axes['logits'] = {0: 'batch_size', 1: 'input_ids_length'}
    for i in range(num_layer):
        dynamic_axes['past_key' +
                     str(i)] = {0: 'batch_size', 2: 'past_sequence_length'}
        dynamic_axes['past_value' +
                     str(i)] = {0: 'batch_size', 2: 'past_sequence_length'}
        dynamic_axes['present_key' +
                     str(i)] = {0: 'batch_size', 2: 'total_length'}
        dynamic_axes['present_value' +
                     str(i)] = {0: 'batch_size', 2: 'total_length'}

    # The first input is required, and other inputs can be passed to torch.onnx.export using dict
    inputs = (input_ids, {
        'attention_mask': attention_mask,
        'position_ids': position_ids,
        'past_key_values': past,
    })

    # Do export using torch.onnx.export
    exported_model = "converted_model.onnx"
    torch.onnx.export(
        model,
        args=inputs,
        f=exported_model,
        export_params=True,
        verbose=False,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=11,
    )
    return exported_model

convert_gpt2_model_to_onnx()
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@patrickvonplaten any ideas?

Hi @HaoboGu thanks for sharing the code snippet! 

Just so I understand a bit better - is there a specific reason why you're trying to avoid the warnings?  

@lewtun I just worry about the warning may lead to unpredictable results when I use the model.

Hey @HaoboGu, I don't think the warning would lead to unpredictable results (or can you clarify a bit more) :-) Can we maybe just leave it? 

@patrickvonplaten yeah, no problem if it's expected