Mllama training via FSDP device and dtype misassignment
Author: blbadger
Pull request: no
Labels: name: bug

### System Info

 - `transformers` version: 4.48.1
- Platform: Linux-5.15.0-1073-azure-x86_64-with-glibc2.35
- Python version: 3.11.0rc1
- Huggingface_hub version: 0.27.1
- Safetensors version: 0.5.2
- Accelerate version: 1.3.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- Tensorflow version (GPU?): 2.15.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Yes
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-40GB

### Who can help?

@amyeroberts, @qubvel

### 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

Shell command used to launch FSDP:
```bash
accelerate launch --config_file "configs/a100_config.yaml" train.py \
--seed 100 \
--model_name_or_path "/path/to/model" \
--dataset_path "/path/to/dataset" \
--add_special_tokens False \
--append_concat_token False \
--max_seq_len 1024 \
--num_train_epochs 15 \
--logging_steps 50 \
--log_level "info" \
--logging_strategy "steps" \
--evaluation_strategy "epoch" \
--save_strategy "epoch" \
--bf16 True\
--fp16 False \
--packing False \
--learning_rate 7e-5 \
--lr_scheduler_type "linear" \
--weight_decay 0.0 \
--warmup_ratio 0.0 \
--max_grad_norm 1.0 \
--output_dir "/path/to/output" \
--per_device_train_batch_size 1 \
--gradient_checkpointing True \
--use_reentrant True \
--dataset_text_field "content" \
--use_flash_attn False \
--use_peft_lora False \
--report_to "none"
```

`train.py` uses a standard `SFTTrainer()` trainer object called via `trainer.train()` after loading the processor and model, nothing fancy there. 

```python
 def main(model_args, data_args, training_args):
    ... # dataset, model, processor initialization

    sft_training_args = SFTConfig(
        output_dir=training_args.output_dir,
        gradient_checkpointing=training_args.gradient_checkpointing,
        bf16=training_args.bf16,
        remove_unused_columns=False,
        report_to=training_args.report_to,
        num_train_epochs=training_args.num_train_epochs,
        logging_steps=training_args.logging_steps,
        evaluation_strategy = training_args.evaluation_strategy,
        save_strategy = training_args.save_strategy,
        max_seq_length=data_args.max_seq_length
    )

    trainer = SFTTrainer(
        model=model,
        peft_config=peft_config,
        args=sft_training_args,
        data_collator=collate_fn,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=processor.tokenizer,
    )
    trainer.train()
)
```

The FSDP config `a100_config.yaml` is as follows:
```yaml
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: 'bf16'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

note that we can specify `mixed_precision: 'no'` for the config and get the same error as well.

### Expected behavior

FSDP (with offloaded parameters) training halts during the first forward pass. The resulting error,

```
 [rank4]:   File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 549, in _conv_forward
[rank4]:     return F.conv2d(
[rank4]:            ^^^^^^^^^
[rank4]: RuntimeError: Input type (torch.FloatTensor) and weight type (CUDABFloat16Type) should be the same or input should be a MKLDNN tensor and weight is a dense tensor
```

is caused by `modeling_mllama.py` line 1541,

```python
patch_embeds = self.patch_embedding(pixel_values.to(self.dtype).to(self.device))
```

AFAIK we cannot generally use `.to()` assignments referencing the model's device or datatype at the start of the forward pass when training via FSDP with offloaded params, as the parameters may not be located in the device or with the datatype that is used during the forward pass.

The simplest solution here is to simply omit both assignments: if we change line 1541 to 

```python
patch_embeds = self.patch_embedding(pixel_values)
```

the problem is eliminated. One would need to change the documentation for Llama 3.2 Vision, as inference involves sending processed inputs to the model's device but not dtype. For example, the official inference code snippet [here](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision)

```python
inputs = processor(image, prompt, return_tensors="pt").to(model.device)
```

needs to become

```python
inputs = processor(image, prompt, return_tensors="pt").to(model.dtype).to(model.device)
```

If this is an acceptable approach I would be happy to put in the PR, just let me know!


Comments:
Hi @blbadger, thanks for opening an issue and proposing a solution! As far as I know, other models in Transformers have a similar pattern. Although the fix might be viable for FSDP, it could be breaking for existing scenarios. I might not have enough experience with that, so let's get some opinions from @SunMarc or @muellerzr.

I think in other vision model we do `self.patch_embedding.weight.device`, e.g. in CLIP

Thanks @zucchini-nlp. I was eager to suggest the same solutions, but according to the issue 

> the parameters may not be located in the device or with the datatype that is used during the forward pass.

@blbadger can you please check if proposed solution fix the issue? If so, we would appreciate a PR 🤗 

Hey @blbadger, there is a similar issue [here](https://github.com/huggingface/transformers/issues/34207#issuecomment-2422356109) where I proposed a potential fix. Can you try it ? 

 @zucchini-nlp @qubvel   Good idea to use the patch embedding's param dtype and weight as that should (IIRC) be dispatched by the time we reference dtype or device as the reference is in that module's forward call.  @SunMarc  I can confirm the following fixes the bug

```python
target_dtype = self.patch_embedding.weight.dtype
target_device = self.patch_embedding.weight.device
patch_embeds = self.patch_embedding(pixel_values.to(target_dtype).to(target_device)))  
```

but as you point out, we are able to omit the `target_device` as it is not strictly necessary for FSDP dispatch or training (tested this too and can confirm no misassignment).  Happy to put in a PR for this if y'all want!

Btw, regarding device movements, I believe it was causing fullgraph compile breaks in some other issues so maybe we should omit it, @qubvel ?

I'm not sure about the graph breaks, but device movement seems unnecessary here. The only concern, removing it might be a breaking change, so I would prefer to leave it unless we get an issue re compilation problem or smth.

@blbadger feel free to ping me in PR, thanks for raising the issue 🤗 

@qubvel Sounds good to me, I will put in the PR with both device and dtype assignment and tag you there.

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.

Resolved