llama fp16 torch.max bug fix
Author: prathikr
Pull request: yes

This PR explicitly sets the threshold tensor for `torch.max` to the same dtype as that of attn_weights to avoid accidental upcasting during mixed precision training. This unblocks ONNX Runtime integration because, without this fix, the torch onnx exporter receives mismatched types into the `torch.max` operation resulting in the following error:

```
Traceback (most recent call last):
  File "run_clm_ort.py", line 641, in <module>
    main()
  File "run_clm_ort.py", line 589, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/optimum/onnxruntime/trainer.py", line 454, in train
    return inner_training_loop(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/optimum/onnxruntime/trainer.py", line 749, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/trainer.py", line 2735, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/optimum/onnxruntime/trainer.py", line 365, in compute_loss
    return super().compute_loss(model_with_loss, inputs, return_outputs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/trainer.py", line 2767, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1675, in forward
    loss = self.module(*inputs, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_utils.py", line 375, in _forward
    return ortmodule._torch_module.forward(*inputs, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_utils.py", line 355, in _forward
    return torch_module_ort._execution_manager(torch_module_ort.is_training()).forward(*inputs, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 274, in forward
    self._fallback_manager.handle_exception(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_fallback.py", line 160, in handle_exception
    raise exception
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_training_manager.py", line 210, in forward
    self._initialize_graph_builder()
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/training/ortmodule/_graph_execution_manager.py", line 502, in _initialize_graph_builder
    self._graph_builder.initialize(self._onnx_models.exported_model.SerializeToString(), grad_builder_config)
RuntimeError: /onnxruntime_src/orttraining/orttraining/python/orttraining_pybind_state.cc:786 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 (Max) bound to different types (tensor(float16) and tensor(float) in node (/_original_module/base_model/model/layers.0/self_attn/Max).
```

Reproduction Intructions:
```Dockerfile
FROM mcr.microsoft.com/azureml/aifx/stable-ubuntu2004-cu117-py38-torch1131

# language-modeling dependencies taken from: https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/requirements.txt
RUN pip install accelerate datasets sentencepiece protobuf evaluate scikit-learn
# additional Hugging Face dependencies
RUN pip install optimum peft transformers

RUN git clone https://github.com/huggingface/optimum.git && \
    cd optimum/examples/onnxruntime/language-modelling && \
    python run_clm.py --model_name_or_path openlm-research/open_llama_7b_400bt_preview --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --do_train --output_dir output_dir --overwrite_output_dir --fp16 --deepspeed zero_stage_1.json --num_train_epochs 1 --logging_steps 1 --optim adamw_ort_fused
```

Who can review?
- text models: @ArthurZucker and @younesbelkada

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

@ArthurZucker I addressed your commment, please have another look. Thank you.

@amyeroberts I addressed your comment. Please merge asap, thank you.