bug in trainer with accelerate prepare of GPT2LMHeadModel using fp16
Author: StevenSong
Pull request: no

### System Info

```
- `transformers` version: 4.30.2
- Platform: Linux-4.15.0-192-generic-x86_64-with-glibc2.27
- Python version: 3.11.3
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1 (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?: yes, model parallelism
```

### Who can help?

@sgugger ~@ pacma~ oops

### 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

```
import os
import sys
import numpy as np
from itertools import chain

import torch
from datasets import load_dataset
from transformers import (
    GPT2TokenizerFast,
    GPT2LMHeadModel,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
    set_seed,
)

seed = 42
torch.manual_seed(seed)
set_seed(seed)
np.random.seed(seed)

tok = GPT2TokenizerFast.from_pretrained("gpt2")
tok.pad_token = tok.eos_token
tok.pad_token_id = tok.eos_token_id

test_size = 0.1
_chunk_size = 256
text_col = "text"

num_workers = min(os.cpu_count(), 2)

max_seq_length = min(_chunk_size, tok.model_max_length)

ds = load_dataset("wikitext", "wikitext-2-v1")

tokenized_ds = ds.map(
    lambda x: tok(x["text"], padding=True, pad_to_multiple_of=max_seq_length),
    remove_columns=[text_col],
    batched=True,
    num_proc=num_workers,
)

def chunk_text(examples, max_seq_length):
    concatenated = {k: list(chain(*examples[k])) for k in examples.keys()}
    tot_len = len(concatenated[list(examples.keys())[0]])
    if tot_len >= max_seq_length:
        tot_len = (
            tot_len // max_seq_length
        ) * max_seq_length
    result = {
        k: [t[i : i + max_seq_length] for i in range(0, tot_len, max_seq_length)]
        for k, t in concatenated.items()
    }
    return result

chunked_ds = tokenized_ds.map(
    lambda x: chunk_text(x, max_seq_length), batched=True, num_proc=num_workers
)

model = GPT2LMHeadModel.from_pretrained(
    "gpt2",
    device_map="auto",
)

data_collator = DataCollatorForLanguageModeling(tok, mlm=False)

args = TrainingArguments(
    output_dir="delete-me",
    per_device_train_batch_size=6,
    logging_steps=500,
    gradient_accumulation_steps=1,
    gradient_checkpointing=False,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=50,
    lr_scheduler_type="cosine",
    learning_rate=5e-6,
    save_steps=10_000,
    fp16=True,  # fp16 bug with GPT2 models in huggingface?
    dataloader_pin_memory=True,
    dataloader_num_workers=2,
    optim="adafactor",
)

trainer = Trainer(
    model=model,
    tokenizer=tok,
    args=args,
    data_collator=data_collator,
    train_dataset=chunked_ds["train"],
)

trainer.train()

trainer.save_model("temp")

```

### Expected behavior

Seems like there were some changes to trainer between v4.29.2 and v4.30.0 to utilize accelerate to prepare the model ([here's the git blame](https://github.com/huggingface/transformers/blame/fe861e578f50dc9c06de33cd361d2f625017e624/src/transformers/trainer.py#L1751-L1752)). With a GPT2LMHeadModel using fp16 precision for training, these changes to trainer lead to the following error from the above script:

```
Traceback (most recent call last):
  File "[...]/min-reproducible.py", line 93, in <module>
    trainer.train()
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/trainer.py", line 1645, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/trainer.py", line 1756, in _inner_training_loop
    model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/accelerate/accelerator.py", line 1182, in prepare
    result = tuple(
             ^^^^^^
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/accelerate/accelerator.py", line 1183, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/accelerate/accelerator.py", line 1022, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/miniconda3/envs/llm/lib/python3.11/site-packages/accelerate/accelerator.py", line 1308, in prepare_model
    model.forward = MethodType(torch.cuda.amp.autocast(dtype=torch.float16)(model.forward.__func__), model)
                                                                            ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'function' object has no attribute '__func__'. Did you mean: '__doc__'?
```

Seems like the `model.forward` object is a `function` rather than a `method` so `__func__` isn't defined. `model` is an instance of GPT2LMHeadModel so I would've expected `model.forward` to be a method on the instance but maybe it's modified somewhere else. ~Overall, I'm not sure if this is a bug of trainer or accelerate or the model.~ Seems like actually this might be an issue on `accelerate` as the folks in the linked issue below are running into it when manually preparing the model (as opposed to letting trainer prepare as I did) - I can reopen this issue in the `accelerate` repo if that's better?

Interestingly, if not using fp16, it runs fine. Ideally, I'd be able to use fp16 with a GPT2LMHeadModel using the trainer.

Seems like someone else has also run into this issue using a LLaMA model: https://github.com/OpenAccess-AI-Collective/axolotl/issues/195#issuecomment-1589657199

Would appreciate any help/fix!

Comments:
cc @pacman100 and @muellerzr 

Having the same issue with QLoRA style PEFT (also when setting FP16 to off).

Hello @StevenSong, 

Thank you for the minimal reproducible example. The culprit in the above example is `device_map="auto",`. The below code runs fine:

```diff
import os
import sys
import numpy as np
from itertools import chain

import torch
from datasets import load_dataset
from transformers import (
    GPT2TokenizerFast,
    GPT2LMHeadModel,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
    set_seed,
)

seed = 42
torch.manual_seed(seed)
set_seed(seed)
np.random.seed(seed)

tok = GPT2TokenizerFast.from_pretrained("gpt2")
tok.pad_token = tok.eos_token
tok.pad_token_id = tok.eos_token_id

test_size = 0.1
_chunk_size = 256
text_col = "text"

num_workers = min(os.cpu_count(), 2)

max_seq_length = min(_chunk_size, tok.model_max_length)

ds = load_dataset("wikitext", "wikitext-2-v1")

tokenized_ds = ds.map(
    lambda x: tok(x["text"], padding=True, pad_to_multiple_of=max_seq_length),
    remove_columns=[text_col],
    batched=True,
    num_proc=num_workers,
)

def chunk_text(examples, max_seq_length):
    concatenated = {k: list(chain(*examples[k])) for k in examples.keys()}
    tot_len = len(concatenated[list(examples.keys())[0]])
    if tot_len >= max_seq_length:
        tot_len = (
            tot_len // max_seq_length
        ) * max_seq_length
    result = {
        k: [t[i : i + max_seq_length] for i in range(0, tot_len, max_seq_length)]
        for k, t in concatenated.items()
    }
    return result

chunked_ds = tokenized_ds.map(
    lambda x: chunk_text(x, max_seq_length), batched=True, num_proc=num_workers
)

model = GPT2LMHeadModel.from_pretrained(
    "gpt2",
-    device_map="auto",
)

data_collator = DataCollatorForLanguageModeling(tok, mlm=False)

args = TrainingArguments(
    output_dir="delete-me",
    per_device_train_batch_size=6,
    logging_steps=500,
    gradient_accumulation_steps=1,
    gradient_checkpointing=False,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=50,
    lr_scheduler_type="cosine",
    learning_rate=5e-6,
    save_steps=10_000,
    fp16=True,  # fp16 bug with GPT2 models in huggingface?
    dataloader_pin_memory=True,
    dataloader_num_workers=2,
    optim="adafactor",
)

trainer = Trainer(
    model=model,
    tokenizer=tok,
    args=args,
    data_collator=data_collator,
    train_dataset=chunked_ds["train"],
)

trainer.train()

trainer.save_model("temp")

```

Hello @sgugger, seems like `device_map` changes the `model.forward` to `function` rather than preserving it as a `method`.

Hello @imarquart, please open a new issue on PEFT wrt issue you are facing with a minimal reproducible example.

The above PR should fix this

Thank you!