set all fusedrope inputs to bf16

[skaulintel]

When passing --bf16 to train/inference cmd, ensure all input tensors are cast to bf16.

Change fused rope input type to bf16 for llama models

[mandy-li]

@skaulintel , what is the reason to hard cast data type to bf16? what about --bf16 is not used in the model's command?

[ghost]

@skaulintel , what is the reason to hard cast data type to bf16? what about --bf16 is not used in the model's command?

+1

[skaulintel]

@mandy-li In order to calculate RoPE in bf16, we have to cast all inputs to bf16. From my analysis, I see that when we pass --bf16 True to the training script when running llama-7b (examples/language-modeling/run_lora_clm.py) I see the following dtypes passed to apply_customized_rope

q: torch.float32, k: torch.bfloat16, cos: torch.float32, sin: torch.float32

Since query is passed in as float32, we won't compute that in bf16. but for the key states we will compute them in bf16. @kbinias please review.

[astachowiczhabana]

huggingface#1026