RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16
Author: mosama1994
Pull request: no
Labels: name: PyTorch FSDP

name: PEFT

### System Info

transformers version = 4.40.0
python = 3.10.2

### Who can help?

_No response_

### Information

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

The code at this link: https://github.com/philschmid/deep-learning-pytorch-huggingface/blob/main/training/scripts/run_fsdp_qlora.py
is the one I am running to train llama3.
line 101:    torch_dtype = torch.bfloat16
line 102:    quant_storage_dtype = torch.bfloat16

When I use just float16, it runs fine. But when I use bfloat16 it gives me this error:
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16


### Expected behavior

Using bfloat16 for loading is causing the issue here in this code, not sure why. Please help ASAP.

Comments:
cc @younesbelkada @pacman100 

Hi @mosama1994 , use this context manager before your training block in the file that you mentioned: 
`with torch.cuda.amp.autocast():` or try moving the context accordingly as i don't have the traceback for  your error but it may solve your issue.
Here is the [reference](https://github.com/huggingface/peft/issues/1515) for the solution provided(solved issue on peft repo).

Cheers!

I had seen this already. How to use this? just wrap the model loading code only? Also, why is this issue still there, in the PEFT repo it is mentioned that this should have been solved a year ago.

> I had seen this already. How to use this? just wrap the model loading code only? 
>
@mosama1994 Basically context manager is used in with forward() and is closed before backward().

>Also, why is this issue still there, in the PEFT repo it is mentioned that this should have been solved a year ago.
>
I just checked and found that the solution wasn't merged with the main branch: look [here](https://github.com/younesbelkada/peft/blob/main/src/peft/tuners/lora.py) in the commit history so, that's why this error still exists.

So, if you want to use autocast then you have to do so in the trainer or it would be simple to use the PR solution which was mentioned in the previous message. Modify your `src/peft/tuners/lora.py` file with the [PR sol](https://github.com/huggingface/peft/pull/121/files) and check if it solves the thing



Hi @mosama1994 
Thanks for the issue ! The solution proposed by @RUFFY-369 is correct, you can also activate autocast through `bf16=True` argument in `TrainingArguments` that you can activate through `--bf16` command line argument if you use the script you shared

Hi @younesbelkada the solution does work. I wrapped the code in the autocast and it is working. I already have the bf16 as True in the TrainingArguments. That is why the issue is arising. When i do fp16 as True and bf16 as False, the code runs. However, with bf16 as True, the issue arises. The solution does work so I am closing but this should be mentioned in the documentation or rectified. Thanks