Cannot enable gradient checkpointing for pre_trained models.
Author: VINUK0
Pull request: no

### System Info

### Environment Number 1 (Google Cloab)

*Transformers Version : 4.34.1*
*Pytorch Version : 2.0.1 (Latest)*
*Python Version : 3.10*
*Accelerate Version : 0.23.0*
*GPU Count : 1*
*GPU Name : Tesla T4*


### Environment Number 2 (Kaggle)

*Transformers Version : 4.34.1*
*Pytorch Version : 2.0.1 (Latest)*
*Python Version : 3.10*
*Accelerate Version : 0.23.0*
*GPU Count : 2*
*GPU Name : Tesla T4 (FSDP) and (DDP)*

### Environment Number 3 (AWS)

*Transformers Version : 4.34.1*
*Pytorch Version : 2.0.1 (Latest)*
*Python Version : 3.10*
*Accelerate Version : 0.23.0*
*GPU Count : 1*
*GPU Name : Tesla T4*




### 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, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

## The first code down below is the code that is causing the error. I cannot enable gradient_checkpointing for this model when calling `.from_pretrained` but it's possible to enable gradient checkpointing when creating a new model from scratch using a config. and i have tested `GPTNeo` models and got the same results.

### This is the example of the way that is not working.
```py
from accelerate import Accelerator
from transformers import AutoTokenizer, OPTForCausalLM
from transformers.optimization import Adafactor
from torch.utils.data import DataLoader
from Dataset import CustomDataset

accelerator = Accelerator()
device = accelerator.device
print(f'Currently using {device}')


tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m')
tokenizer.pad_token = tokenizer.eos_token



model = OPTForCausalLM.from_pretrained('facebook/opt-350m')
model.gradient_checkpointing_enable()
Train_DataLoader = DataLoader(CustomDataset('/content/HFA/data.db', tokenizer=tokenizer, dataset_line_max_length=1024), batch_size=3, shuffle=False)

optimizer = Adafactor(model.parameters(),
                      lr=7e-5,
                      clip_threshold=1.0,
                      decay_rate=-1.0,
                      weight_decay=0.05,
                      scale_parameter=False,
                      relative_step=False,
                      warmup_init=False)

model, optimizer, Train_DataLoader = accelerator.prepare(model, optimizer, Train_DataLoader)






for i in range(20):
    for batch_idx, batch in enumerate(Train_DataLoader):
        optimizer.zero_grad()
        input_ids, attention_mask, labels = batch
        sm_output = model(input_ids, attention_mask=attention_mask, labels=labels)
        s_loss = sm_output[0]
        
        print(f'Current Loss is {s_loss}\n', end='\r')
        print(f'Total Batch {len(Train_DataLoader)}/{batch_idx}\n\n', end='\r')
        accelerator.backward(s_loss)
        optimizer.step()
    print(f'Current Epoch is {i}\n\n\n', end='\r')
        
    if (i + 1) % 2 == 0:
       model.save_pretrained('/content/HFA/')
```

### This is the working code that let me enable gradient checkpointing.

```py
from accelerate import Accelerator
from transformers import AutoModelForCausalLM, AutoTokenizer, OPTForCausalLM, OPTConfig
from transformers.optimization import Adafactor
from torch.utils.data import DataLoader
import torch.nn as nn
from Dataset import CustomDataset
import torch

accelerator = Accelerator()
device = accelerator.device
print(f'Currently using {device}')




tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m')
tokenizer.pad_token = tokenizer.eos_token

model_config = OPTConfig(vocab_size=50272, hidden_size=624, num_hidden_layers=24, ffn_dim=4096, max_position_embeddings=2048, do_layer_norm_before=False, word_embed_proj_dim=768, num_attention_heads=16, activation_function="silu", attention_dropout=0.0, layer_norm_elementwise_affine=True, layerdrop=0.0, init_std=0.02, dropout=0.1, enable_bias=True, eos_token_id=2, bos_token_id=2)

model = OPTForCausalLM(model_config)
model.gradient_checkpointing_enable()
Train_DataLoader = DataLoader(CustomDataset('/content/HFA/data.db', tokenizer=tokenizer, dataset_line_max_length=1024), batch_size=3, shuffle=False)

optimizer = Adafactor(model.parameters(),
                      lr=8e-5,
                      clip_threshold=1.0,
                      decay_rate=-1.0,
                      weight_decay=0.05,
                      scale_parameter=False,
                      relative_step=False,
                      warmup_init=False)

model, optimizer, Train_DataLoader = accelerator.prepare(model, optimizer, Train_DataLoader)






for i in range(20):
    for batch_idx, batch in enumerate(Train_DataLoader):
        optimizer.zero_grad()
        input_ids, attention_mask, labels = batch
        output = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = output[0]
        
        print(f'Current Loss is {loss}\n', end='\r')
        print(f'Total Batch {len(Train_DataLoader)}/{batch_idx}\n\n', end='\r')
        accelerator.backward(loss)
        optimizer.step()
    print(f'Current Epoch is {i}\n\n\n', end='\r')
        
    if (i + 1) % 2 == 0:
       model.save_pretrained('/content/HFA/')
```

### Expected behavior

*In both of the code's provided training process continues without problems, but when using the first code i cannot enable gradient checkpointing. while using the second code allow me to enable gradient checkpointing.*

Comments:
Hi... this seems to be a good first bug. Please assign it to me.

Sure, feel free to open a PR and link it to this issue 😉 

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.