Peft intergration incompatible with Trainer/gradient checkpointing
Author: xinghaow99
Pull request: no

### System Info

- `transformers` version: 4.33.0.dev0
- Platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.31
- 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: MULTI_GPU
        - mixed_precision: no
        - use_cpu: False
        - debug: False
        - num_processes: 4
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: 0,1,2,3
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu117 (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

### Who can help?

@pacman100 @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I'm playing a model like this:
```
class CustomModel(PreTrainedModel):
    def __init__(self, config: PretrainedConfig):
        self.config = config
        super().__init__(self.config)
        self.peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
        )
        self.model = AutoModelForCausalLM.from_pretrained('EleutherAI/pythia-70m', 
                                                          )
        self.config.update(self.model.config.to_dict())
        self.tokenizer = AutoTokenizer.from_pretrained('EleutherAI/pythia-70m')

        self.model.gradient_checkpointing_enable()

        # self.model = get_peft_model(self.model, self.peft_config)
        self.model.add_adapter(self.peft_config)
        self.model.config.use_cache = False

    def test_grad_exist(self):
        input_ids = self.tokenizer("Hello, my dog is cute", return_tensors="pt").input_ids.to(self.device)
        print('Requires Grad:', self.model(input_ids, return_dict=True, output_hidden_states=True).hidden_states[-1][:, -1, :].requires_grad)

    def get_exemplar_embeddings(self, exemplar_input_ids, exemplar_attention_mask):
        hidden_states = self.model(input_ids=exemplar_input_ids, attention_mask=exemplar_attention_mask, return_dict=True, output_hidden_states=True).hidden_states
        exemplar_embeddings = hidden_states[-1][:, -1, :]
        return exemplar_embeddings
        
    def forward(self,
                exemplar_input_ids: Optional[torch.LongTensor] = None,
                exemplar_attention_mask: Optional[torch.LongTensor] = None,
                input_ids: Optional[torch.LongTensor] = None,
                attention_mask: Optional[torch.LongTensor] = None,
                labels: Optional[torch.LongTensor] = None,
                ):
        self.test_grad_exist()
        self.model.enable_adapters()
        self.test_grad_exist()
        exemplar_embeddings = self.get_exemplar_embeddings(exemplar_input_ids, exemplar_attention_mask) # [num_exemplars, hidden_size]
        self.model.disable_adapters()
        attention_mask = F.pad(attention_mask, (exemplar_embeddings.shape[0]+1, 0), value=1) # [batch_size, 1 + num_exemplars + seq_len]
        inputs_embeds = self.model.get_input_embeddings()(input_ids) # [batch_size, seq_len, hidden_size]
        bos_embeddings = self.model.get_input_embeddings()(torch.tensor(self.tokenizer.bos_token_id, device=self.device).repeat(inputs_embeds.shape[0], 1)) # [batch_size, 1, hidden_size]
        inputs_embeds = torch.cat((bos_embeddings,exemplar_embeddings.unsqueeze(0).repeat(inputs_embeds.shape[0], 1, 1), inputs_embeds), dim=1) # [batch_size, 1 + num_exemplars + seq_len, hidden_size]
        outputs = self.model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels, return_dict=True) # CausalLMOutputWithPast
        return outputs
```

The idea is to forward the peft model twice in a single step: once with the adapters to get the embeddings, and once without the adapters but with the embeddings as input.
But somehow I met this error using the current implementation: `RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn`.

I'm not sure why this is happening. 
Some extra information:
1. For this dummy forward, things seem to be functional. `self.test_grad_exist()` outputs True.
```exemplar = ["Hello, my cat is cute"]
exemplar_input_ids = model.tokenizer(exemplar, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device)
exemplar_attention_mask = torch.ones_like(exemplar_input_ids)
input_str = ["Hello, my dog is cute"]
input_ids = model.tokenizer(input_str, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device)
attention_mask = torch.ones_like(input_ids)
labels = [[-100]*2 + input_ids.tolist()[0]]
labels = torch.tensor(labels).to(model.device)
outputs = model(exemplar_input_ids=exemplar_input_ids, exemplar_attention_mask=exemplar_attention_mask, input_ids=input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss.backward()
```
2. When training it with a Trainer, I got the mentioned error.
3. When training it with a Trainer but disable gradient checkpointing, it also works.
4. When using `get_peft_model` directly in the initialization, and use `get_base_model` for operation in forward, it also works.


### Expected behavior

`test_grad_exist()` should always output True and no errors occur.

Comments:
Hi @xinghaow99 thanks for the issue and reproducer, I managed to reproduce the issue, https://github.com/huggingface/transformers/pull/25846 should fix it , the issue is that you need to make sure the input `requires_grad` are force-set to `True`, we do it in PEFT: https://github.com/huggingface/peft/blob/85013987aa82aa1af3da1236b6902556ce3e483e/src/peft/peft_model.py#L334 that should explain why you are not facing the issue with a `PeftModel`. 
Make sure also to call `self.model.gradient_checkpointing_enable()` after adding an adapter


```python
class CustomModel(PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.model = AutoModelForCausalLM.from_pretrained('EleutherAI/pythia-70m')
        self.peft_config = LoraConfig(
            task_type="CAUSAL_LM", inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
        )
        self.config.update(self.model.config.to_dict())
        self.tokenizer = AutoTokenizer.from_pretrained('EleutherAI/pythia-70m')

        # self.model = get_peft_model(self.model, self.peft_config)
        self.model.add_adapter(self.peft_config)
        self.model.config.use_cache = False

        self.model.gradient_checkpointing_enable()

    def test_grad_exist(self):
        input_ids = self.tokenizer("Hello, my dog is cute", return_tensors="pt").input_ids.to(self.device)
        print('Requires Grad:', self.model(input_ids, return_dict=True, output_hidden_states=True).hidden_states[-1][:, -1, :].requires_grad)

    def get_exemplar_embeddings(self, exemplar_input_ids, exemplar_attention_mask):
        hidden_states = self.model(input_ids=exemplar_input_ids, attention_mask=exemplar_attention_mask, return_dict=True, output_hidden_states=True).hidden_states
        exemplar_embeddings = hidden_states[-1][:, -1, :]
        return exemplar_embeddings
        
    def forward(self,
                exemplar_input_ids: Optional[torch.LongTensor] = None,
                exemplar_attention_mask: Optional[torch.LongTensor] = None,
                input_ids: Optional[torch.LongTensor] = None,
                attention_mask: Optional[torch.LongTensor] = None,
                labels: Optional[torch.LongTensor] = None,
                ):
        self.test_grad_exist()
        self.model.enable_adapters()
        self.test_grad_exist()
        exemplar_embeddings = self.get_exemplar_embeddings(exemplar_input_ids, exemplar_attention_mask) # [num_exemplars, hidden_size]
        self.model.disable_adapters()
        attention_mask = F.pad(attention_mask, (exemplar_embeddings.shape[0]+1, 0), value=1) # [batch_size, 1 + num_exemplars + seq_len]
        inputs_embeds = self.model.get_input_embeddings()(input_ids) # [batch_size, seq_len, hidden_size]
        bos_embeddings = self.model.get_input_embeddings()(torch.tensor(self.tokenizer.bos_token_id, device=self.device).repeat(inputs_embeds.shape[0], 1)) # [batch_size, 1, hidden_size]
        inputs_embeds = torch.cat((bos_embeddings,exemplar_embeddings.unsqueeze(0).repeat(inputs_embeds.shape[0], 1, 1), inputs_embeds), dim=1) # [batch_size, 1 + num_exemplars + seq_len, hidden_size]
        outputs = self.model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels, return_dict=True) # CausalLMOutputWithPast
        return outputs
```

Hi @younesbelkada. Thanks for the quick fix. Awesome work!