Repository: huggingface/optimum

Issue:
Support for GPU optimization of xlm-roberta model type
### System Info

```shell
optimum==1.3.0
python==3.8.13
Ubuntu 18.04
```


### Who can help?

@philschmid 

### Information

- [X] The official example scripts
- [x] 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

Hi,

Thanks for this very nice tool for simple and rapid ONNX optimization of Transformer models.

I would like to GPU optimize this multilingual cross-encoder model : `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` using the code provided in a blog post by @philschmid:

```
from optimum.onnxruntime import ORTOptimizer
from optimum.onnxruntime.configuration import OptimizationConfig

# create ORTOptimizer and define optimization configuration
optimizer = ORTOptimizer.from_pretrained(model_id, feature=model.pipeline_task)
optimization_config = OptimizationConfig(optimization_level=99,     # enable all optimizations
                                        optimize_for_gpu=True,
                                        fp16=True
                                        )

# apply the optimization configuration to the model
optimizer.export(
    onnx_model_path=onnx_path / "model.onnx",
    onnx_optimized_model_output_path=onnx_path / "model-optimized.onnx",
    optimization_config=optimization_config,
)
```

But got the following error:
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb Cellule 15 in <cell line: 12>()
      [6](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=5) optimization_config = OptimizationConfig(optimization_level=99,     # enable all optimizations
      [7](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=6)                                         optimize_for_gpu=True,
      [8](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=7)                                         fp16=True
      [9](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=8)                                         )
     [11](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=10) # apply the optimization configuration to the model
---> [12](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=11) optimizer.export(
     [13](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=12)     onnx_model_path=onnx_path / "model.onnx",
     [14](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=13)     onnx_optimized_model_output_path=onnx_path / "model-optimized.onnx",
     [15](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=14)     optimization_config=optimization_config,
     [16](vscode-notebook-cell:/home/matthieu/Code/Python/ONNX-Export-Optimum-mmarco-mMiniLMv2-L12-H384-v1.ipynb#ch0000014?line=15) )

File ~/anaconda3/envs/haystack-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/optimization.py:125, in ORTOptimizer.export(self, onnx_model_path, onnx_optimized_model_output_path, optimization_config, use_external_data_format)
    122 if not onnx_model_path.exists():
    123     export(self.preprocessor, self.model, self._onnx_config, self.opset, onnx_model_path)
--> 125 ORTConfigManager.check_supported_model_or_raise(self._model_type)
    126 num_heads = getattr(self.model.config, ORTConfigManager.get_num_heads_name(self._model_type))
    127 hidden_size = getattr(self.model.config, ORTConfigManager.get_hidden_size_name(self._model_type))

File ~/anaconda3/envs/haystack-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/utils.py:110, in ORTConfigManager.check_supported_model_or_raise(cls, model_type)
    107 @classmethod
    108 def check_supported_model_or_raise(cls, model_type: str) -> bool:
    109     if model_type not in cls._conf:
...
    111             f"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. "
    112             f"If you want to support {model_type} please propose a PR or open up an issue."
    113         )

KeyError: "xlm-roberta model type is not supported yet. Only ['bert', 'albert', 'camembert', 'codegen', 'distilbert', 'deberta', 'deberta-v2', 'electra', 'roberta', 'bart', 'gpt2', 'gpt_neo'] are supported. If you want to support xlm-roberta please propose a PR or open up an issue."
```

Would it be possible to add support for this model type?

Thanks!

### Expected behavior

Would it be possible to add support for `xlm-roberta` model type?


Discussion:
Hi @Matthieu-Tinycoaching, `xlm-roberta` is not yet supported by `ORTOptimizer`, and [this PR](https://github.com/huggingface/optimum/pull/312/files) will enable it once it makes its way to the main. 

Resolving patch:
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -70,6 +70,7 @@ class ORTConfigManager:
         "bart": ("encoder_attention_heads", "d_model", "bart"),
         "gpt2": ("n_head", "n_embd", "gpt2"),
         "gpt_neo": ("num_heads", "hidden_size", "gpt2"),
+        "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"),
     }
 
     @classmethod


Test patch:
diff --git a/tests/onnxruntime/test_optimization.py b/tests/onnxruntime/test_optimization.py
--- a/tests/onnxruntime/test_optimization.py
+++ b/tests/onnxruntime/test_optimization.py
@@ -30,6 +30,7 @@
     ElectraForSequenceClassification,
     GPT2ForSequenceClassification,
     RobertaForSequenceClassification,
+    XLMRobertaForSequenceClassification,
 )
 
 import onnx
@@ -67,6 +68,7 @@ class ORTOptimizerTest(unittest.TestCase):
         (GPT2ForSequenceClassification, "hf-internal-testing/tiny-random-gpt2"),
         (RobertaForSequenceClassification, "hf-internal-testing/tiny-random-roberta"),
         (ElectraForSequenceClassification, "hf-internal-testing/tiny-random-electra"),
+        (XLMRobertaForSequenceClassification, "hf-internal-testing/tiny-xlm-roberta"),
     )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID)