ONNX export results for hidden states/attentions are incorrect if enabled
Author: minimaxir
Pull request: no

ONNX export results for hidden states/attentions are incorrect if enabled (e.g. via `config.output_attentions = True` and `config.output_hidden_states = True` before loading the model). This is likely due to the returning-tuples nature of the corresponding `forward` functions which ONNX may not be able to trace correctly; wrapping said functions and using `torch.stack` to handle said tuples is a workaround to the issue.

## Environment info
- `transformers` version: 4.18.0.dev0
- Platform: macOS-12.3-x86_64-i386-64bit
- Python version: 3.9.7
- Huggingface_hub version: 0.1.1
- PyTorch version (GPU?): 1.11.0 (False)

### Who can help
@lewtun 

## Information

Taking from the [new export tutorial](https://huggingface.co/docs/transformers/serialization):

```python3
from transformers import AutoConfig
from typing import Mapping, OrderedDict
from transformers.onnx import OnnxConfig
import torch
from pathlib import Path
from transformers.onnx import export
from transformers import AutoTokenizer, AutoModel


model_ckpt = "distilbert-base-uncased"
config = AutoConfig.from_pretrained(model_ckpt)

config.output_attentions = True
config.output_hidden_states = True

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
base_model = AutoModel.from_pretrained(model_ckpt, config=config)

pt_inputs = tokenizer(["Using DistilBERT with ONNX Runtime is so very cool!"]*2, return_tensors="pt")
pt_outputs = base_model(**pt_inputs)

# Expected results:
# torch.Size([2, 16, 768])
# torch.Size([2, 7, 16, 768])
# torch.Size([2, 6, 12, 16, 16])
print(pt_outputs['last_hidden_state'].shape)
for key in ["hidden_states", "attentions"]:
    print(torch.transpose(torch.stack(list(pt_outputs[key])), 0, 1).shape) # transposing to set the batch dimension first

class DistilBertOnnxConfig(OnnxConfig):
    @property
    def inputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict(
            [
                ("input_ids", {0: "batch", 1: "sequence"}),
                ("attention_mask", {0: "batch", 1: "sequence"}),
            ]
        )

    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict(
            [
                ("last_hidden_state", {0: "batch", 1: "sequence_length"}),
                ("hidden_states", {0: "batch", 1: "sequence_length"}),
                ("attentions", {0: "batch", 1: "sequence_length"},),
            ]
        )

onnx_path = Path("onnx/model.onnx")
onnx_config = DistilBertOnnxConfig(config)

onnx_inputs, onnx_outputs = export(tokenizer, base_model, onnx_config, onnx_config.default_onnx_opset, onnx_path)

session = InferenceSession(str(onnx_path))
onnx_inputs = tokenizer(["Using DistilBERT with ONNX Runtime is so very cool!"]*2, return_tensors="np")
onnx_outputs = session.run(output_names=["last_hidden_state", "hidden_states", "attentions"], input_feed=dict(onnx_inputs))

# Incorrect results; using the same shape as the first `last_hidden_state`:
# (2, 16, 768)
# (2, 16, 768)
# (2, 16, 768)
for i in range(len(outputs)):
    print(outputs[i].shape)
```

The ugly workaround:

```python3
from transformers import AutoConfig
from typing import Mapping, OrderedDict
from transformers.onnx import OnnxConfig
import torch
from pathlib import Path
from transformers.onnx import export
from transformers import AutoTokenizer, AutoModel


model_ckpt = "distilbert-base-uncased"
config = AutoConfig.from_pretrained(model_ckpt)

config.output_attentions = True
config.output_hidden_states = True

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
base_model = AutoModel.from_pretrained(model_ckpt, config=config)

from transformers.modeling_utils import PreTrainedModel

class ExportModel(PreTrainedModel):
    def __init__(self):
        super().__init__(config)
        self.model = base_model

    def forward(self, input_ids=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,):

        out = self.model(input_ids, attention_mask)
        return {
            "last_hidden_state": out["last_hidden_state"],
            "hidden_states": torch.transpose(torch.stack(list(out["hidden_states"])), 1, 0),
            "attentions": torch.transpose(torch.stack(list(out["attentions"])), 1, 0)
        }

    def call(self, input_ids=None,
        attention_mask=None,):
        self.forward(input_ids,
        attention_mask,)

class DistilBertOnnxConfig(OnnxConfig):
    @property
    def inputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict(
            [
                ("input_ids", {0: "batch", 1: "sequence"}),
                ("attention_mask", {0: "batch", 1: "sequence"}),
            ]
        )

    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        return OrderedDict(
            [
                ("last_hidden_state", {0: "batch", 1: "sequence_length"}),
                ("hidden_states", {0: "batch", 2: "sequence_length"}),
                ("attentions",{0: "batch", 3: "sequence_length", 4: "sequence_length"},),
            ]
        )

new_model = ExportModel()

onnx_path = Path("onnx/model.onnx")
onnx_config = DistilBertOnnxConfig(config)

onnx_inputs, onnx_outputs = export(tokenizer, new_model, onnx_config, onnx_config.default_onnx_opset, onnx_path)

session = InferenceSession(str(onnx_path))
onnx_inputs = tokenizer(["Using DistilBERT with ONNX Runtime is so very cool!"]*2, return_tensors="np")
onnx_outputs = session.run(output_names=["last_hidden_state", "hidden_states", "attentions"], input_feed=dict(onnx_inputs))

# correct results
# (2, 16, 768)
# (2, 7, 16, 768)
# (2, 6, 12, 16, 16)
for i in range(len(outputs)):
    print(outputs[i].shape)
```

Comments:
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.

Thanks for reporting this bug @minimaxir ! I'll need to have a closer look at what's going wrong in the tracing, but I think your hunch about tuples in the forward pass is the source of the problem.

cc @michaelbenayoun 

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.