Add onnx export cuda support
Author: JingyaHuang
Pull request: yes

# What does this PR do?

* Add CUDA support for `transformers.onnx.export_pytorch`.
* Add test for `transformers.onnx.export_pytorch` on CUDA.

# Context
While executing `optimum.ORTTrainer` with `--deepspeed` and `--fp16` enabled, the export to onnx will fail since all layers of the models are not implemented for half-precision. Need to trace on CUDA as workaround.

## Who can review?
@michaelbenayoun @lewtun 




Comments:
_The documentation is not available anymore as the PR was closed or merged._

> Great, thanks for adding this @JingyaHuang!
> 
> If I understand correctly, this enables tracing half-precision models?

Hi @michaelbenayoun , 
Yes, but only for PyTorch since `tf2onnx` has [specified the device to be CPU](https://github.com/onnx/tensorflow-onnx/blob/main/tf2onnx/convert.py#L609).

> Thanks for iterating on this @JingyaHuang !
> 
> I've left a few final nits, but this is looking really nice :)
> 
> Could you please confirm that the slow tests pass on both CPU and GPU devices?
> 
> ```
> RUN_SLOW=1 pytest tests/onnx/test_onnx_v2.py
> ```

Hi @lewtun , by running the slow tests on CPU and GPU, I got the following results. It seems that some models and tasks failed. Trying to find out the root of the problems now.
```
======================================================================================== short test summary info ========================================================================================
FAILED tests/onnx/test_onnx_v2.py::OnnxExportTestCaseV2::test_pytorch_export_12_bert_next_sentence_prediction - ValueError: next-sentence-prediction is not a supported task, supported tasks: dict_ke...
FAILED tests/onnx/test_onnx_v2.py::OnnxExportTestCaseV2::test_pytorch_export_71_mobilebert_next_sentence_prediction - ValueError: next-sentence-prediction is not a supported task, supported tasks: d...
FAILED tests/onnx/test_onnx_v2.py::OnnxExportTestCaseV2::test_pytorch_export_on_cuda_12_bert_next_sentence_prediction - ValueError: next-sentence-prediction is not a supported task, supported tasks:...
FAILED tests/onnx/test_onnx_v2.py::OnnxExportTestCaseV2::test_pytorch_export_on_cuda_20_big_bird_question_answering - AssertionError: big-bird, question-answering -> Expected all tensors to be on th...
FAILED tests/onnx/test_onnx_v2.py::OnnxExportTestCaseV2::test_pytorch_export_on_cuda_71_mobilebert_next_sentence_prediction - ValueError: next-sentence-prediction is not a supported task, supported ...
========================================================== 5 failed, 177 passed, 77 skipped, 43 deselected, 158 warnings in 2478.21s (0:41:18) ==========================================================
```

Oh yes, we recently reverted the next-sentence-prediction feature in #17276, so rebasing on `main` should fix those. The BigBird error looks more related to your PR, so let me know if you need some help debugging it :)

> Oh yes, we recently reverted the next-sentence-prediction feature in #17276, so rebasing on `main` should fix those. The BigBird error looks more related to your PR, so let me know if you need some help debugging it :)

Hi @lewtun , thanks for the details. After rebasing, all checks for bert passed. And the problem of big bird comes from a [bug in the modeling](https://github.com/huggingface/transformers/blob/main/src/transformers/models/big_bird/modeling_big_bird.py#L3102), that while creating the `token_type_ids`, the device is not specified lead to a mismatch of devices. I just fixed that. Now all checks of `pytorch_export` either on CPU or on CUDA passed.

> After rebasing, all checks for bert passed.

Cool! Just to double-check, did you run the tests:

* On a CPU machine (no GPU, CUDA installed)
* On a GPU machine

I'd like to be sure we don't accidentally break the test suite for developers coding on CPU machines :)