torch.export & save doesn't work
Author: MrParosk
Pull request: no

### System Info

## Description

PyTorch 2.x supports torch.export.export to export a model into a graph representation and later serialize to with torch.export.save.

However, trying to do that with HF transformers models doesn't work (see reproduction section below).

I think the main issue is that HF transformers seems to return dataclasses (e.g. DetrObjectDetectionOutput) when doing a forward pass. However, these dataclasses cannot be serialized to json (and hence the errors).

Is there any plans to support this in the near future? If so, is it possible to contribute to it?

## System Info

- `transformers` version: 4.36.1
- Platform: Linux-6.2.0-39-generic-x86_64-with-glibc2.35
- Python version: 3.11.7
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu118 (True)
- 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?: no
- Using distributed or parallel set-up in script?: no

## Reproduction

```python
import torch
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)
dummy_input = torch.randint(0, 1000, size=(2, 512))

exported_model = torch.export.export(model, args=(dummy_input,))
torch.export.save(exported_model, "model.pt")
```

Gives:

```python
Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
SequenceClassifierOutput(loss=None, logits=tensor([[-0.1969,  0.6469, -0.4259,  0.0733, -0.4104],
        [-0.0467,  0.5932, -0.3885,  0.1355, -0.4127]],
       grad_fn=<AddmmBackward0>), hidden_states=None, attentions=None)
Traceback (most recent call last):
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/utils/_pytree.py", line 462, in _treespec_to_json
    serialized_context = json.dumps(spec.context)
                         ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/.pyenv/versions/3.11.7/lib/python3.11/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/.pyenv/versions/3.11.7/lib/python3.11/json/encoder.py", line 200, in encode
    chunks = self.iterencode(o, _one_shot=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/.pyenv/versions/3.11.7/lib/python3.11/json/encoder.py", line 258, in iterencode
    return _iterencode(o, 0)
           ^^^^^^^^^^^^^^^^^
  File "/home/erik/.pyenv/versions/3.11.7/lib/python3.11/json/encoder.py", line 180, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type type is not JSON serializable

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/erik/code/tutorials/huggingface/example.py", line 12, in <module>
    torch.export.save(exported_model, "model.pt")
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/export/__init__.py", line 1075, in save
    save(ep, f, extra_files=extra_files, opset_version=opset_version)
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/_export/__init__.py", line 473, in save
    serialized_program, serialized_state_dict = serialize(ep, opset_version)
                                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/_export/serde/serialize.py", line 1438, in serialize
    ExportedProgramSerializer(opset_version).serialize(exported_program)
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/_export/serde/serialize.py", line 810, in serialize
    ).serialize(exported_program.graph_module)
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/_export/serde/serialize.py", line 790, in serialize
    call_spec=serialize_call_spec(self.call_spec),
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/_export/serde/serialize.py", line 194, in serialize_call_spec
    out_spec=treespec_dumps(call_spec.out_spec, TREESPEC_VERSION) if call_spec.out_spec else "",
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/utils/_pytree.py", line 517, in treespec_dumps
    json_spec = _SUPPORTED_PROTOCOLS[protocol].treespec_to_json(treespec)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/erik/code/tutorials/huggingface/.venv/lib/python3.11/site-packages/torch/utils/_pytree.py", line 464, in _treespec_to_json
    raise TypeError(
TypeError: Unable to serialize context. Please make the context json dump-able, or register a custom serializer using _register_pytree_node.
```

Similar results are given when trying other models, like DETR:

```python
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
import requests

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)


processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")


inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

exported_model = torch.export.export(model, args=(inputs["pixel_values"], inputs["pixel_mask"]))
torch.export.save(exported_model, "model.pt")
```

### Who can help?

@ArthurZucker @amyeroberts

### Information

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

### Tasks

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

### Reproduction

See above

### Expected behavior

That we can export the models with torch.export.export & torch.export.save

Comments:
Yes, You are totally right, All hugging face models have outputs that are instances of subclasses of [ModelOutput](https://huggingface.co/docs/transformers/v4.36.1/en/main_classes/output#transformers.utils.ModelOutput) which as you mentioned aren't `JSON serializable`.

Thankfully, You can opt not to use them and just return tuples instead by setting the kwarg `return_dict=False` in the model loading method.

```python
import torch
from transformers import AutoModelForSequenceClassification

hf_model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5, return_dict=False)
dummy_input = torch.randint(0, 1000, size=(2, 512))

with torch.no_grad():
    print(hf_model(dummy_input))

# (tensor([[-0.1271, -0.2149, -0.0251,  0.3655,  0.2407],
#        [-0.1448, -0.1787, -0.0458,  0.3854,  0.2657]]),)

exported_model = torch.export.export(hf_model, args=(dummy_input,))
torch.export.save(exported_model, "model.pt")
pt_model = torch.export.load('model.pt')

with torch.no_grad():
    print(pt_model(dummy_input))

# (tensor([[-0.1271, -0.2149, -0.0251,  0.3655,  0.2407],
#        [-0.1448, -0.1787, -0.0458,  0.3854,  0.2657]]),)

```

Yes seems to be working, also works with the Trainer API, thanks a lot! @IbrahimAmin1