Models fine-tuned with gradient checkpointing (=True) fails to export to ONXX
Author: samru-rai
Pull request: no
Labels: name: wontfix

## Environment info
- `transformers` version: 3.5.0
- Platform: Linux-5.4.0-52-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.9
- PyTorch version (GPU?): 1.7.0 (True)
- Tensorflow version (GPU?): 2.3.1 (False)
- Using GPU in script?: False
- Using distributed or parallel set-up in script?: False

### Who can help
 
Hi @LysandreJik and @patrickvonplaten, I hope I have tagged the right person. If not please untag yourself and tag the right person.

## Information

Model I am using (Bert):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

It's a super simple script that uses gradient check-pointing with BERT.

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

A dummy data 

## To reproduce

I have made a reproducible google collab here -> https://colab.research.google.com/drive/1tUpIzbugZ4xPz6eAOJtZT-fGww9LemwN?usp=sharing

Steps to reproduce the behavior:

1. Open the notebook
2. Runtime->Run all

Error thrown:
```python
RuntimeError                              Traceback (most recent call last)

<ipython-input-5-2702a59a9c3e> in <module>()
      7         tokenizer=tokenizer,  # <-- CHANGED: add tokenizer
      8         output=Path("onnx/bert-base-cased.onnx"),
----> 9         opset=11)
     10 
     11 # Tensorflow

4 frames

/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, opset_version, _retain_param_name, do_constant_folding, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, enable_onnx_checker, use_external_data_format, onnx_shape_inference, use_new_jit_passes)
    648                     params_dict, opset_version, dynamic_axes, defer_weight_export,
    649                     operator_export_type, strip_doc_string, val_keep_init_as_ip, custom_opsets,
--> 650                     val_add_node_names, val_use_external_data_format, model_file_location)
    651             else:
    652                 proto, export_map = graph._export_onnx(

RuntimeError: ONNX export failed: Couldn't export Python operator CheckpointFunction


```




Comments:
Indeed, I see why this would fail. I don't have access to your notebook, but as a temporary workaround you could do:

```py
model.save_pretrained("here")
model = ModelClass.from_pretrained("here", gradient_checkpointing=False)
```

You should be able to convert that model to ONNX then.

I made it public now, my bad. Another work around I found was to edit the `config.json` by setting   `"gradient_checkpointing": false`. I did this because the convert script looks for the model in a path and not in memory.

Yes, this works too!

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.