Fix ORTTrainer failure on gpt2 fp16 training
Author: JingyaHuang
Pull request: yes

# What does this PR do?

Fixes [#11279 of onnxruntime](https://github.com/microsoft/onnxruntime/issues/11279)

## Context

Optimum users reported that the mixed-precision training on gpt2 with `optimum.onnxruntime.ORTTrainer` is broken since transformers>4.16.0. After investigation, the break comes from the [removal of `float()`](https://github.com/huggingface/transformers/pull/14321/files#r912791693) in gpt2 modeling from PR #14321.

__Reproduction__

Run optimum onnxruntime training example [run_glue.py](https://github.com/huggingface/optimum/blob/main/examples/onnxruntime/training/text-classification/run_glue.py) with:
```bash
python run_glue.py \
    --model_name_or_path gpt2 \
    --task_name sst2 \
    --do_train \
    --do_eval \
    --fp16 \
    --output_dir /tmp/ort-gpt2-sst2/
```
 
__Error Message__
```
RuntimeError: /onnxruntime_src/orttraining/orttraining/python/orttraining_pybind_state.cc:752 
onnxruntime::python::addObjectMethodsForTraining(pybind11::module&, 
onnxruntime::python::ExecutionProviderRegistrationFn)::<lambda(onnxruntime::training::OrtModuleGraphBuilder*, const 
pybind11::bytes&, const onnxruntime::training::OrtModuleGraphBuilderConfiguration&)> [ONNXRuntimeError] : 1 : FAIL : 
Type Error: Type parameter (T) of Optype (Where) bound to different types (tensor(float) and tensor(float16) in node 
(Where_201).
```
As mentioned in the error message, the forward with onnxruntime InferenceSession will fail on a node __Where__ in the graph, which corresponds to the [Where op](https://github.com/michaelbenayoun/transformers/blob/9ef5813d73b2af238e8820e117dd32415ce4c173/src/transformers/models/gpt2/modeling_gpt2.py#L206) in gpt2 modeling.

And the problem comes from the fact that after removing `float()`, during fp16 training, the inputs of Where have different dtype (one in fp32 and one in fp16), which violates the definition in ONNX and leads to the failure.

## Who can review?

@michaelbenayoun @patrickvonplaten, @LysandreJik

## Fix

Ensure `attn_weights` and `value` has the same type in exported ONNX IR.



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

Hello @JingyaHuang By looking this 2 lines

https://github.com/huggingface/transformers/blob/04ffba9a9a4d452b328eb7e57f7c87f2ed78bd10/src/transformers/models/decision_transformer/modeling_decision_transformer.py#L196-L197

`mask_value` is already of type `attn_weights.dtype`. Does the issue only occur when we use ONNX? (i.e. if we run the model in Python with FP16, does it work?). This issue seems strange. Do you happen to know which argument gets fp32 and which one gets fp16?



Hi @ydshieh, yes this issue only occurs with ONNX. When exporting the ONNX IR `mask_value` is exported as a constant initializer(min of fp32) with dtype float32. Thus during the mixed-precision training with onnxruntime, the `attn_weights` will be in dtype fp16 and `mask_value` as a constant always fp32 -> two inputs with different dtype -> training failed. 
Here is the ONNX IR which illustrates what happened with Where op:
![image](https://user-images.githubusercontent.com/44135271/177803009-ca35b67b-6614-4d83-a34a-33252e6e305d.png)

__[EDIT]__ Here I made a mistake, according to the training graph, actually `mask_value` was successfully cast to fp16, but `attn_weights` not. Check the local exported IR below.

And if we run the model with PyTorch backend, there is no problem of the tricky tracing or op definition, it should work fine.

@JingyaHuang Thank you!

Hi @ydshieh, I've just double-checked the debug exported training onnx graph. Actually the `mask_value` has been cast to fp16 before `Where` node, and it was `attn_weights` which was fp32, and the fix inserts another cast op to cast `attention_mask` from fp32 to fp16.

The IR corresponding to this line
https://github.com/huggingface/transformers/blob/2544c1434f8d831daff3fe6a925dced67bc70c64/src/transformers/models/decision_transformer/modeling_decision_transformer.py#L181


The IR before fix:
<img width="180" alt="image" src="https://user-images.githubusercontent.com/44135271/177832731-570a6ae4-7e84-4b57-ba11-b2cd29194ed9.png">
The IR after fix: 
<img width="180" alt="image" src="https://user-images.githubusercontent.com/44135271/177832878-cc9480a8-fffb-4e24-9648-12a9dd28fb7a.png">

So this is exactly what we want for fp16 training.

Gently pinging @patrickvonplaten and @LysandreJik for a review. 

I got the same error. But I use the torch.fx and amp to train the GPT2 model. I fix this error with the method is add `attn_weights.to(attn_weights.dtype)` in `torch.where`. I don't know why this way can fix it, but it does.

Hi @TXacs , which version you tried? Could you try to install the latest version on `main`: 

```bash
pip install git+https://github.com/huggingface/accelerate
```
and see if you still have the issue (without your fix). Thanks!