ONNX export vmap functorch"
Author: Worke1221
Pull request: no
Labels: name: Feature request

### Feature request

## Environment info

- `transformers` version: 4.44.2 (or your version)
- Platform: Linux (Ubuntu 22.04)
- Python version: 3.10.12
- PyTorch version (GPU?): 2.7.0+cu126
- Using GPU in script?: Yes (but export fails on CPU too)
- ONNX Runtime version: 1.18.0 (if applicable)

## Information

Model I am using: **Qwen/Qwen2-0.5B**

The problem arises when using:
- [x] the official example scripts: (give details below)
- [x] my own modified scripts: (see below)

The tasks I am working on:
- [x] ONNX export of causal language model

## To reproduce

Steps to reproduce the behavior:

1. Run this minimal export script:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen2-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Dummy input
input_ids = torch.randint(0, 1000, (1, 512))
attention_mask = torch.ones_like(input_ids)

# Temporarily disable cache for export
model.config.use_cache = False

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    "qwen2.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "seq"},
        "attention_mask": {0: "batch", 1: "seq"},
        "logits": {0: "batch", 1: "seq"}
    },
    opset_version=14
)

### Motivation

## Environment info

- `transformers` version: 4.44.2 (or your version)
- Platform: Linux (Ubuntu 22.04)
- Python version: 3.10.12
- PyTorch version (GPU?): 2.7.0+cu126
- Using GPU in script?: Yes (but export fails on CPU too)
- ONNX Runtime version: 1.18.0 (if applicable)

## Information

Model I am using: **Qwen/Qwen2-0.5B**

The problem arises when using:
- [x] the official example scripts: (give details below)
- [x] my own modified scripts: (see below)

The tasks I am working on:
- [x] ONNX export of causal language model

## To reproduce

Steps to reproduce the behavior:

1. Run this minimal export script:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen2-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Dummy input
input_ids = torch.randint(0, 1000, (1, 512))
attention_mask = torch.ones_like(input_ids)

# Temporarily disable cache for export
model.config.use_cache = False

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    "qwen2.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "seq"},
        "attention_mask": {0: "batch", 1: "seq"},
        "logits": {0: "batch", 1: "seq"}
    },
    opset_version=14
)

### Your contribution

## Environment info

- `transformers` version: 4.44.2 (or your version)
- Platform: Linux (Ubuntu 22.04)
- Python version: 3.10.12
- PyTorch version (GPU?): 2.7.0+cu126
- Using GPU in script?: Yes (but export fails on CPU too)
- ONNX Runtime version: 1.18.0 (if applicable)

## Information

Model I am using: **Qwen/Qwen2-0.5B**

The problem arises when using:
- [x] the official example scripts: (give details below)
- [x] my own modified scripts: (see below)

The tasks I am working on:
- [x] ONNX export of causal language model

## To reproduce

Steps to reproduce the behavior:

1. Run this minimal export script:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen2-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Dummy input
input_ids = torch.randint(0, 1000, (1, 512))
attention_mask = torch.ones_like(input_ids)

# Temporarily disable cache for export
model.config.use_cache = False

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    "qwen2.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "seq"},
        "attention_mask": {0: "batch", 1: "seq"},
        "logits": {0: "batch", 1: "seq"}
    },
    opset_version=14
)

Comments:
Hey 👋 
this is kind of expected atm, there is an extra wrapper module at https://github.com/huggingface/transformers/blob/9c20660138830ca362533551ca978c27b48283a1/src/transformers/integrations/executorch.py#L193
with example usage in the tests, e.g. https://github.com/huggingface/transformers/blob/9c20660138830ca362533551ca978c27b48283a1/tests/models/qwen2/test_modeling_qwen2.py#L200

It's still far from perfect tho and we also want to default to non-vmap functions in the future which may alleviate these issues as well! cc @Cyrilvallez