FSDP2 not compatible with transformers >= 4.54.0 GenericForTokenClassification
Author: ETOgaosion
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.54.0
- Platform: Linux-5.10.135.bsk.6-amd64-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.34.3
- Safetensors version: 0.5.3
- Accelerate version: 1.9.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (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 L20

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

Python script `test_fsdp2.py`:

```py
import torch
from torch.distributed.fsdp import fully_shard
from transformers import AutoModelForTokenClassification

if not torch.distributed.is_initialized():
    torch.distributed.init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo", init_method="env://"
    )

    rank = torch.distributed.get_rank()

    if torch.cuda.is_available():
        torch.cuda.set_device(rank)


def load_valuehead_model(path):
    return AutoModelForTokenClassification.from_pretrained(path)

def apply_fsdp2(model, fsdp_kwargs):
    """model: AutoModelForCausalLM"""

    default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None)

    modules = []

    for idx, module in enumerate(modules):
        # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
        #     print(f"wrap module {module.__class__.__name__}")
        fully_shard(module, **fsdp_kwargs)

    # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
    #     print(f"wrap module {model.__class__.__name__}")
    fully_shard(model, **fsdp_kwargs)  # fsdp2 will not reshard_after_forward for root module

model = load_valuehead_model("Qwen/Qwen2-0.5B")
module = apply_fsdp2(model, {})

if torch.distributed.is_initialized():
    torch.distributed.destroy_process_group()
```

Execute:

```sh
torchrun --standalone --nnodes=1 --nproc-per-node=2 test_fsdp2.py
```

Bug:

```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/workspace/verl/test_fsdp2.py", line 37, in <module>
[rank1]:     module = apply_fsdp2(model, {})
[rank1]:   File "/workspace/verl/test_fsdp2.py", line 34, in apply_fsdp2
[rank1]:     fully_shard(model, **fsdp_kwargs)  # fsdp2 will not reshard_after_forward for root module
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/_composable/contract.py", line 150, in wrapper
[rank1]:     updated = func(inp_module, *args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_fully_shard/_fully_shard.py", line 241, in fully_shard
[rank1]:     module.__class__ = new_cls
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2041, in __setattr__
[rank1]:     super().__setattr__(name, value)
[rank1]: TypeError: __class__ assignment: 'FSDPQwen2ForTokenClassification' object layout differs from 'Qwen2ForTokenClassification'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/verl/test_fsdp2.py", line 37, in <module>
[rank0]:     module = apply_fsdp2(model, {})
[rank0]:   File "/workspace/verl/test_fsdp2.py", line 34, in apply_fsdp2
[rank0]:     fully_shard(model, **fsdp_kwargs)  # fsdp2 will not reshard_after_forward for root module
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/_composable/contract.py", line 150, in wrapper
[rank0]:     updated = func(inp_module, *args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_fully_shard/_fully_shard.py", line 241, in fully_shard
[rank0]:     module.__class__ = new_cls
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2041, in __setattr__
[rank0]:     super().__setattr__(name, value)
[rank0]: TypeError: __class__ assignment: 'FSDPQwen2ForTokenClassification' object layout differs from 'Qwen2ForTokenClassification'
```

### Expected behavior

Use transformers 4.53.2:

```
Successfully installed transformers-4.53.2
root@25ab4d5939b5:/workspace/verl# torchrun --standalone --nnodes=1 --nproc-per-node=2 test_fsdp2.py 
W0807 07:34:18.264000 3600 torch/distributed/run.py:766] 
W0807 07:34:18.264000 3600 torch/distributed/run.py:766] *****************************************
W0807 07:34:18.264000 3600 torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W0807 07:34:18.264000 3600 torch/distributed/run.py:766] *****************************************
Some weights of Qwen2ForTokenClassification were not initialized from the model checkpoint at Qwen/Qwen2.5-0.5B and are newly initialized: ['score.bias', 'score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Some weights of Qwen2ForTokenClassification were not initialized from the model checkpoint at Qwen/Qwen2.5-0.5B and are newly initialized: ['score.bias', 'score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

Comments:
Ah that's annoying, cc @Cyrilvallez we should patch this!

Thanks a lot for the report and simple repro @ETOgaosion, appreciated! See the PR I mentioned for the fix! 🤗