convert_graph_to_onnx.convert broken for gpt-neo-x.xB since 4.5.0.dev0
Author: oborchers
Pull request: no

## Environment info

- `transformers` version: 4.9.1
- Platform: Linux-4.15.0-151-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1+cu102 (True)
- Tensorflow version (GPU?): 2.4.2 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes (V100)
- Using distributed or parallel set-up in script?: No

### Who can help
This issue is a follow-up of #9803. People tagged in previous issues of the same kind:
@mfuntowicz @LysandreJik @patrickvonplaten 
@StellaAthena (because of EleutherAI)

## Information

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] 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)

When trying to run any of the ONNX export scripts on the `EleutherAI/gpt-neo-x.xb` models, they fail. The last version of `transformers` that I can trace this behavior back to is on commit [04ceee7d246dcb26e13fa6aa8e3b990d6a0bf289](https://github.com/huggingface/transformers/commit/04ceee7d246dcb26e13fa6aa8e3b990d6a0bf289). This is not the exact commit, but the last working one I know. This issue is also present in the current `4.9.1` tag, as well as the recently introduced custom configurations to export more easily.

## To reproduce

Running the following script with the current `4.9.1` tag fails (output below). Installing from the above mentioned commit results in a properly working export.

```
pip install -U git+git://github.com/huggingface/transformers.git@04ceee7d246dcb26e13fa6aa8e3b990d6a0bf289
```

Actual code to run:

```
from pathlib import Path
import torch
import transformers
from transformers import convert_graph_to_onnx
from transformers import pipeline

model_name = "EleutherAI/gpt-neo-1.3B"
model_pth = Path(f"gpt_neo/gpt_neo_13b.onnx")
model_pth.parent.mkdir(exist_ok=True, parents=True)

class GPTNeoSent(transformers.GPTNeoForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        self.sentence_embedding = torch.nn.Identity()

    def forward(self, input_ids, attention_mask):
        return self.sentence_embedding(
            super().forward(input_ids, attention_mask=attention_mask).logits
        )
model = GPTNeoSent(config=transformers.AutoConfig.from_pretrained(model_name)).from_pretrained(model_name)

nlp = pipeline(
    "feature-extraction",
    model=model,
    tokenizer=model_name,
)

inputs = nlp.tokenizer(["hello my friends!"], return_tensors="pt")

with torch.no_grad():
    (
        input_names,
        output_names,
        dynamic_axes,
        tokens,
    ) = convert_graph_to_onnx.infer_shapes(nlp, "pt")
    ordered_input_names, model_args = convert_graph_to_onnx.ensure_valid_input(
        nlp.model, tokens, input_names
    )

if not model_pth.exists():
    torch.onnx.export(
        model,
        (inputs["input_ids"], inputs["attention_mask"]),
        f=model_pth.as_posix(),
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        do_constant_folding=True,
        use_external_data_format=True, # Needed because of model size
        enable_onnx_checker=True,
        opset_version=13,
    )
```

Which runs into the following error:

```
Found input input_ids with shape: {0: 'batch', 1: 'sequence'}
Found input attention_mask with shape: {0: 'batch', 1: 'sequence'}
Found output output_0 with shape: {0: 'batch', 1: 'sequence'}
Ensuring inputs are in correct order
Generated inputs order: ['input_ids', 'attention_mask']
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py:779: 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!
  assert batch_size > 0, "batch_size has to be defined and > 0"
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py:149: 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!
  while seq_length % block_length != 0:
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py:220: UserWarning: ONNX export failed on Unfold because input size not accessible not supported
  warnings.warn("ONNX export failed on " + op + " because " + msg + " not supported")
Traceback (most recent call last):
  File "test.py", line 42, in <module>
    torch.onnx.export(
  File "/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/__init__.py", line 271, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py", line 88, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py", line 709, in _export
    proto, export_map = graph._export_onnx(
RuntimeError: ONNX export failed: Couldn't export operator aten::unfold

Defined at:
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py(189): _look_back
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py(216): create_local_attention_mask
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py(801): forward
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/nn/modules/module.py(860): _slow_forward
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/nn/modules/module.py(887): _call_impl
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/transformers/models/gpt_neo/modeling_gpt_neo.py(974): forward
test.py(18): forward
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/nn/modules/module.py(860): _slow_forward
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/nn/modules/module.py(887): _call_impl
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/jit/_trace.py(116): wrapper
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/jit/_trace.py(125): forward
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/nn/modules/module.py(889): _call_impl
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/jit/_trace.py(1139): _get_trace_graph
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py(380): _trace_and_get_graph_from_model
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py(420): _create_jit_graph
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py(457): _model_to_graph
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py(694): _export
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/utils.py(88): export
/home/oborchers/anaconda3/envs/dev/lib/python3.8/site-packages/torch/onnx/__init__.py(271): export
test.py(42): <module>
```

## Expected behavior

Export should work.


Comments:
Hi @oborchers, 

Thanks for raising the issue.

GPT-Neo is not supported by `convert_graph_to_onnx.py` and even if the model was potentially successfully exported in the past, I would not be surprised if some axis definition would be wrong.

With the new package `transformers.onnx` we are working on initial support for GPT-Neo, please see the PR [here](https://github.com/huggingface/transformers/pull/12911).

If you want to give it a try, we would love your feedback. 

Hi @mfuntowicz,

thanks for doing a great re-implementation of the ONNX export functions, those are of great help 👍 

I may have been a bit over eager in creating the issue, as well as changing state, as I am technically using a custom script as of now.

Went through a lot of hassle to actually re-create the onnx checkpoint to see if the most recent changes to ONNX do actually have any effect on the models performance, as back then original torch version was way faster than the exported one. See here: https://github.com/microsoft/onnxruntime/issues/7238

So, even if this was exportable properly on your side, it would be almost unusable due to being much slower than the pytorch version. At lest with the last exportable version (4.5.0.dev0). Let my try your code tomorrow to see if there are any differences in the results.

Thanks @oborchers for all the details.

We haven't run benchmark to properly say, we are just validating the outputs (with/without the past buffers) are matching the PyTorch outputs.

We would be very interested in supporting your efforts improving performance for GPT-Neo, so don't hesitate to ping us (@michaelbenayoun and myself).

Also, we can potential look at what the offline optimizations provided by Ort can bring here.

Thanks 🤗

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/master/CONTRIBUTING.md) are likely to be ignored.