AutoModelForSequenceClassification + GPT2 + Accelerate + FSDP fails to load pretrained model
Author: jmif
Pull request: no
Labels: name: solved

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-6.2.0-1012-gcp-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: FSDP
	- mixed_precision: fp16
	- use_cpu: False
	- debug: False
	- num_processes: 2
	- machine_rank: 0
	- num_machines: 1
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- fsdp_config: {'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP', 'fsdp_backward_prefetch_policy': 'NO_PREFETCH', 'fsdp_forward_prefetch': False, 'fsdp_offload_params': False, 'fsdp_sharding_strategy': 1, 'fsdp_state_dict_type': 'FULL_STATE_DICT', 'fsdp_sync_module_states': True, 'fsdp_use_orig_params': False}
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: FSDP

### Who can help?

 @muellerz @pacman100

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

I'm trying to fine tune GPT2 using accelerate + HF training + FSDP with the intention of training larger GPT2 models later when this is working.  With this training script:

```
import logging
import tempfile

from datasets import load_dataset
from transformers import (
    set_seed,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
    AutoTokenizer,
    AutoModelForSequenceClassification,
)


logger = logging.getLogger()


def train(output_dir: str):
    set_seed(42)

    imdb = load_dataset("imdb")

    tokenizer = AutoTokenizer.from_pretrained("gpt2", max_length=1024)

    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForSequenceClassification.from_pretrained(
        "gpt2",
        num_labels=2,
        max_position_embeddings=1024,
        use_safetensors=True,
    )

    def tokenize(examples):
        return tokenizer(examples["text"], truncation=True)

    train_dataset = imdb["train"].map(tokenize, batched=True)

    training_args = TrainingArguments(
        output_dir=output_dir,
        eval_delay=0,
        evaluation_strategy="epoch",
        learning_rate=3e-5,
        weight_decay=3e-6,
        adam_beta1=0.9,
        adam_beta2=0.999,
        adam_epsilon=1e-8,
        num_train_epochs=10,
        lr_scheduler_type="linear",
        warmup_steps=10,
        log_level="info",
        save_strategy="epoch",
        seed=43,
        fp16=True,
        dataloader_drop_last=True,
    )

    collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="max_length")

    trainer = Trainer(
        model,
        training_args,
        tokenizer=tokenizer,
        data_collator=collator,
        train_dataset=train_dataset,
    )

    trainer.train()


def main():
    with tempfile.TemporaryDirectory() as d:
        train(d)


if __name__ == "__main__":
    main()
```

I get this error:

```
  File "repro.py", line 73, in main
    train(d)
  File "repro.py", line 27, in train
    model = AutoModelForSequenceClassification.from_pretrained(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 516, in from_pretrained
    return model_class.from_pretrained(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3091, in from_pretrained
    ) = cls._load_pretrained_model(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3492, in _load_pretrained_model
    set_module_tensor_to_device(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/utils/modeling.py", line 265, in set_module_tensor_to_device
    new_module = getattr(module, split)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1614, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
```

Here's how I run:

```
python -m accelerate.commands.launch --config_file config/accelerate.yaml epro.py
```

### Expected behavior

When I run this without accelerate (ie run the repro.py script directly) I don't run into the load error.

Comments:
Hello, this has been resolved by PR https://github.com/huggingface/transformers/pull/25820

Confirmed!