FSDP2 + LoRA `model.generate` raise `aten.embedding.default: got mixed torch.Tensor and DTensor` error
Author: Xiao-Chenguang
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.1
- Platform: Linux-5.15.0-136-generic-x86_64-with-glibc2.35
- Python version: 3.12.9
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.11.0
- Accelerate config:    not found
- DeepSpeed version: 0.18.2
- PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@3outeille @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

Use FSDP2 to load a Qwen2 model with LoRA adapter, model.generate failed with the `RuntimeError: aten.embedding.default: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!`.

The main script is:
```python
import os

import torch
from accelerate import Accelerator
from datasets import Dataset, load_dataset
from peft import LoraConfig, TaskType, get_peft_model
from torch.optim import AdamW
from transformers import AutoModelForCausalLM, AutoTokenizer

use_lora = True
temp_fix = False
model_name = "Qwen/Qwen2-0.5B-Instruct"

rank = int(os.environ.get("LOCAL_RANK", "0"))
torch.cuda.set_device(rank)

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")
assert isinstance(dataset, Dataset), f"Expected Dataset, got {type(dataset)}"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

if use_lora:
    lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj"],
        task_type=TaskType.CAUSAL_LM,
    )
    model = get_peft_model(model, lora_config)

if rank == 0:
    print(f"Rank [[{rank}]] model: {type(model)}")

accelerator = Accelerator()
optimizer = AdamW(model.parameters(), lr=1e-5)
model, optimizer, dataset = accelerator.prepare(model, optimizer, dataset)


if temp_fix:
    # fix `RuntimeError: aten.embedding.default: got mixed torch.Tensor and DTensor`
    # ==============================================================================
    from torch.distributed.fsdp import register_fsdp_forward_method

    register_fsdp_forward_method(model, "generate")
    # ==============================================================================

if rank == 0:
    print(f"Rank [[{rank}]] model: {model}")

for data in dataset:
    print(f"Rank: [[{rank}]] input data: {data['prompt'][0]['content']}")
    inputs = tokenizer(data["prompt"][0]["content"], return_tensors="pt").to(rank)
    response = model.generate(**inputs)
    print(
        f"Rank: [[{rank}]] generated response:",
        tokenizer.decode(response[0], skip_special_tokens=True),
    )
    break
```

fsdp config is
```yaml
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
  fsdp_version: 2
gpu_ids: '0,1'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

Launch script is
```bash
accelerate launch \
    --config_file "acc_fsdp.yaml \
    transformer-issue.py
```

Error message is
```log
[rank1]: Traceback (most recent call last):
[rank1]:   File "/project/grpo/transformer-issue.py", line 52, in <module>
[rank1]:     response = model.generate(**inputs)
[rank1]:                ^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/peft/peft_model.py", line 1973, in generate
[rank1]:     outputs = self.base_model.generate(*args, **kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2564, in generate
[rank1]:     result = decoding_method(
[rank1]:              ^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2784, in _sample
[rank1]:     outputs = self(**model_inputs, return_dict=True)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 918, in wrapper
[rank1]:     output = func(self, *args, **kwargs)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 449, in forward
[rank1]:     outputs: BaseModelOutputWithPast = self.model(
[rank1]:                                        ^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 1064, in wrapper
[rank1]:     outputs = func(self, *args, **kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 345, in forward
[rank1]:     inputs_embeds = self.embed_tokens(input_ids)
[rank1]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 192, in forward
[rank1]:     return F.embedding(
[rank1]:            ^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 2546, in embedding
[rank1]:     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/_compile.py", line 53, in inner
[rank1]:     return disable_fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/distributed/tensor/_api.py", line 350, in __torch_dispatch__
[rank1]:     return DTensor._op_dispatcher.dispatch(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/distributed/tensor/_dispatch.py", line 151, in dispatch
[rank1]:     op_info = self.unwrap_to_op_info(op_call, args, kwargs)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/distributed/tensor/_dispatch.py", line 350, in unwrap_to_op_info
[rank1]:     self._try_replicate_spec_for_scalar_tensor(
[rank1]:   File "/project/.venv/lib/python3.12/site-packages/torch/distributed/tensor/_dispatch.py", line 452, in _try_replicate_spec_for_scalar_tensor
[rank1]:     raise RuntimeError(
[rank1]: RuntimeError: aten.embedding.default: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!
```

The script include a temperory fix for it. Trun off LoRA or turn on the fix works as expected.

### Expected behavior

LLM generate a sequence under FSDP2 and LoRA.

Comments:
cc @3outeille 


`model.generate` is not pretty commen in traditional supervised training, while it maybe used in test and is common in LLM RL.
I got this error initially in trl GRPOTrainer with FSDP2 and PEFT as reported in huggingface/trl#4550.

Putting the fix into transformer Trainer maybe a better idea to fix this general issue for trl trainers like GRPO.

I tried to add those lines into transformer.Trainer `_inner_training_loopa` around line 2875, but not sure it's a general fix.

```python
        if self.is_fsdp_enabled:
            self.model = self.model_wrapped = model
            # fix `RuntimeError: aten.embedding.default: got mixed torch.Tensor and DTensor`
            # when use FSDP2 to train a LoRA model
            # ==============================================================================
            if self.accelerator.state.fsdp_plugin.fsdp_version == 2:
                from torch.distributed.fsdp import register_fsdp_forward_method

                register_fsdp_forward_method(self.model, "generate")
            # ==============================================================================
```