MLlama does not work with FSDP mix-precision
Author: YangFei1990
Pull request: no
Labels: name: bug

name: PyTorch FSDP

name: Vision

### System Info

transformers==4.45.2

### Who can help?

@ArthurZucker

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

MLlama is setting `self.dtype` from its own parameters' dytpe and casting some input/intermediate tensor into the `self.dtype`. Here are references
- https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1471
- https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1504
However running with FSDP's mix-precision FSDP expects model to hold fp32 master params and it will cast the params to low precision during training runtime. However MLlama will expect its `self.dtype` to be fp32 while it should actually be low precision. 

So will run into below error
```
F.conv2d(input, weight, bias, self.stride
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (CUDABFloat16Type) should be the same
[12:56](https://amzn-aws.slack.com/archives/D01D5TA374N/p1729108566881649)
```

### Expected behavior

MLlama should work with FSDP mix-precision

Comments:
cc @muellerzr and @SunMarc as well

Hey @YangFei1990 , thanks for the report. Can you check if you have the same issue with basic mixed-precision training ? I think that the modeling code needs potentially to be changed a bit. In other vision model, we would do : 

In instructblip model

```python
target_dtype = self.patch_embedding.weight.dtype
patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]
```

or do nothing to the dtype of pixel_values (e.g idefics 2)

```python
patch_embeds = self.patch_embedding(pixel_values)
```

Would you like to open a PR to propose a fix ? Thanks ! 


Hi @SunMarc, not sure what you meant basic mixed-precision training. For Pytorch, both autocast and FSDP requires model weights to stay in fp32. I think 2 should be working to let FSDP handle the cast automatically. I'm not sure how 1 will work with FSDP.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> Hi @SunMarc, not sure what you meant basic mixed-precision training. For Pytorch, both autocast and FSDP requires model weights to stay in fp32. I think 2 should be working to let FSDP handle the cast automatically. I'm not sure how 1 will work with FSDP.

Enable mixed-precision training but remove FSDP 

@SunMarc Removing FSDP is not a feasible solution for 90B on 80GB HBM. 

Could you then try adding the modification I specified [here](https://github.com/huggingface/transformers/issues/34207#issuecomment-2422356109) and see if it works ? 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.