Gradient Checkpointing Fails with frozen parameters
Author: jamesharrisivi
Pull request: no

### System Info

The PEFT methods freeze the bulk of the transformer, apart from an external module. 

When I enable gradient checkpointing and train with these models or even if I simply freeze an embedding layer of a normal model, training breaks. So this problem is not specific to PEFT: gradient_checkpointing + frozen first parameter = Error

But if I do
```
 for n, param in model.named_parameters():
            param.requires_grad = True
            break
```
it trains successfully. So it seems like there is a check if the first parameter has a gradient. Ideally, I would not have to set the first parameter (embedding) to True, as I want the whole model including embeddings frozen.

```
  warnings.warn("None of the inputs have requires_grad=True. Gradients will be None")
/admin/miniconda3/envs/peft/lib/python3.10/site-packages/torch/utils/checkpoint.py:31: UserWarning: None of the inputs have requires_grad=True. Gradients will be None
  warnings.warn("None of the inputs have requires_grad=True. Gradients will be None")

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /admin/peft/model/model_training/trainer.py:48 │
│ 0 in <module>                                                                                    │
│                                                                                                  │
│                                                                      │
│                                                                                                  │
│ /admin/peft/model/model_training/trainer.py:47 │
│ 4 in main                                                                                        │
│                                                                                                  │
│   471 │   │   compute_metrics=partial(compute_metrics, metrics=metrics, preprocess_fns=preproc   │
│   472 │   │   preprocess_logits_for_metrics=preprocess_logits_for_metrics,                       │
│   473 │   )                                                                                      │
│ ❱ 474 │   trainer.train(resume_from_checkpoint=training_conf.resume_from_checkpoint)             │
│   475 │   trainer.save_model()                                                                   │
│   476 │   tokenizer.save_pretrained(output_dir)                                                  │
│   477                                                                                            │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/transformers/ │
│ trainer.py:1639 in train                                                                         │
│                                                                                                  │
│   1636 │   │   inner_training_loop = find_executable_batch_size(                                 │
│   1637 │   │   │   self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size  │
│   1638 │   │   )                                                                                 │
│ ❱ 1639 │   │   return inner_training_loop(                                                       │
│   1640 │   │   │   args=args,                                                                    │
│   1641 │   │   │   resume_from_checkpoint=resume_from_checkpoint,                                │
│   1642 │   │   │   trial=trial,                                                                  │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/transformers/ │
│ trainer.py:1906 in _inner_training_loop                                                          │
│                                                                                                  │
│   1903 │   │   │   │   │   with model.no_sync():                                                 │
│   1904 │   │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                  │
│   1905 │   │   │   │   else:                                                                     │
│ ❱ 1906 │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                      │
│   1907 │   │   │   │                                                                             │
│   1908 │   │   │   │   if (                                                                      │
│   1909 │   │   │   │   │   args.logging_nan_inf_filter                                           │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/transformers/ │
│ trainer.py:2668 in training_step                                                                 │
│                                                                                                  │
│   2665 │   │   │   │   scaled_loss.backward()                                                    │
│   2666 │   │   elif self.deepspeed:                                                              │
│   2667 │   │   │   # loss gets scaled under gradient_accumulation_steps in deepspeed             │
│ ❱ 2668 │   │   │   loss = self.deepspeed.backward(loss)                                          │
│   2669 │   │   else:                                                                             │
│   2670 │   │   │   loss.backward()                                                               │
│   2671                                                                                           │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/deepspeed/uti │
│ ls/nvtx.py:11 in wrapped_fn                                                                      │
│                                                                                                  │
│    8 │   function call."""                                                                       │
│    9 │   def wrapped_fn(*args, **kwargs):                                                        │
│   10 │   │   get_accelerator().range_push(func.__qualname__)                                     │
│ ❱ 11 │   │   ret_val = func(*args, **kwargs)                                                     │
│   12 │   │   get_accelerator().range_pop()                                                       │
│   13 │   │   return ret_val                                                                      │
│   14                                                                                             │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/deepspeed/run │
│ time/engine.py:1974 in backward                                                                  │
│                                                                                                  │
│   1971 │   │   if self.zero_optimization():                                                      │
│   1972 │   │   │   self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumula  │
│   1973 │   │   │   )                                                                             │
│ ❱ 1974 │   │   │   self.optimizer.backward(loss, retain_graph=retain_graph)                      │
│   1975 │   │   elif self.amp_enabled():                                                          │
│   1976 │   │   │   # AMP requires delaying unscale when inside gradient accumulation boundaries  │
│   1977 │   │   │   # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-i  │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/deepspeed/run │
│ time/zero/stage_1_and_2.py:2028 in backward                                                      │
│                                                                                                  │
│   2025 │   │   │   scaled_loss = self.external_loss_scale * loss                                 │
│   2026 │   │   │   scaled_loss.backward()                                                        │
│   2027 │   │   else:                                                                             │
│ ❱ 2028 │   │   │   self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)            │
│   2029 │                                                                                         │
│   2030 │   def check_overflow(self, partition_gradients=True):                                   │
│   2031 │   │   self._check_overflow(partition_gradients)                                         │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/deepspeed/run │
│ time/fp16/loss_scaler.py:54 in backward                                                          │
│                                                                                                  │
│    51 │                                                                                          │
│    52 │   def backward(self, loss, retain_graph=False):                                          │
│    53 │   │   scaled_loss = loss * self.loss_scale                                               │
│ ❱  54 │   │   scaled_loss.backward(retain_graph=retain_graph)                                    │
│    55                                                                                            │
│    56                                                                                            │
│    57 class LossScaler(LossScalerBase):                                                          │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/torch/_tensor │
│ .py:488 in backward                                                                              │
│                                                                                                  │
│    485 │   │   │   │   create_graph=create_graph,                                                │
│    486 │   │   │   │   inputs=inputs,                                                            │
│    487 │   │   │   )                                                                             │
│ ❱  488 │   │   torch.autograd.backward(                                                          │
│    489 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs                     │
│    490 │   │   )                                                                                 │
│    491                                                                                           │
│                                                                                                  │
│ /admin//miniconda3/envs/peft/lib/python3.10/site-packages/torch/autogra │
│ d/__init__.py:197 in backward                                                                    │
│                                                                                                  │
│   194 │   # The reason we repeat same the comment below is that                                  │
│   195 │   # some Python versions print out the first line of a multi-line function               │
│   196 │   # calls in the traceback and some print out the last line                              │
│ ❱ 197 │   Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the bac   │
│   198 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,                        │
│   199 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to ru   │
│   200                                                                                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
```

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] 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


from transformers import AutoModel

from peft import LoraConfig,  get_peft_model
model = AutoModel.from_pretrained('decapoda-research/llama-7b-hf')
config = LoraConfig(
            r=16,
            lora_alpha=32,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
        )
model = get_peft_model(model, config)

and then proceed to train with gradient_checkpointing enabled.

### Expected behavior

Gradient checkpointing shouldn't affect whether a subset of parameters are frozen. As PEFT models are increasingly popular as well as gradient_checkpointing  it makes sense to get to the bottom of this bug.

Comments:
cc @younesbelkada and @pacman100 

hi @jamesharrisivi 
For that you need to make sure that the input being passed to the peft model has `requires_grad` set to `True`
This is a duplicate of https://discuss.huggingface.co/t/peft-lora-gpt-neox-backward-pass-failing/35641 
Can you try to add:
```python
if hasattr(model, "enable_input_require_grads"):
    model.enable_input_require_grads()
else:
    def make_inputs_require_grad(module, input, output):
         output.requires_grad_(True)

    model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
```
somewhere in your training script, before the call to `get_peft_model` ?


I agree though we can do it directly when creating the peft model, I propose to fix it in https://github.com/huggingface/peft/pull/404

@younesbelkada Is this fix also intended to be compatible with Deepspeed?

I think it would work with DS out of the box but not sure, did you tried already on your end @bradfox2 ?

> I think it would work with DS out of the box but not sure, did you tried already on your end @bradfox2 ?

From what I saw, get_input_embeddings() wasn't defined on the DS model object.

> hi @jamesharrisivi For that you need to make sure that the input being passed to the peft model has `requires_grad` set to `True` This is a duplicate of [discuss.huggingface.co/t/peft-lora-gpt-neox-backward-pass-failing/35641](https://discuss.huggingface.co/t/peft-lora-gpt-neox-backward-pass-failing/35641) Can you try to add:
> 
> ```python
> if hasattr(model, "enable_input_require_grads"):
>     model.enable_input_require_grads()
> else:
>     def make_inputs_require_grad(module, input, output):
>          output.requires_grad_(True)
> 
>     model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
> ```
> 
> somewhere in your training script, before the call to `get_peft_model` ?

Hi:
may I ask if this is the choice for re-entrant=True by default now? BTW, the input for embedding is INT tensors, so why input.requires_grad_(True) worked here?

> hi @jamesharrisivi For that you need to make sure that the input being passed to the peft model has `requires_grad` set to `True` This is a duplicate of https://discuss.huggingface.co/t/peft-lora-gpt-neox-backward-pass-failing/35641 Can you try to add:
> 
> ```python
> if hasattr(model, "enable_input_require_grads"):
>     model.enable_input_require_grads()
> else:
>     def make_inputs_require_grad(module, input, output):
>          output.requires_grad_(True)
> 
>     model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
> ```
> 
> somewhere in your training script, before the call to `get_peft_model` ?

How can we back-prop the gradients when we use `model.generate()` efficiently?



Hey @SuperBruceJia please refrain from asking the same question on each and every single issue that is unrelated 😅 The forum is the best place to discuss this 🤗 

> How can we back-prop the gradients when we use model.generate() efficiently?

Currently you can't back-prop when using generate as that method uses `torch.no_grad()` context manager. You can overwrite the method and make sure it is not using that context manager

> > How can we back-prop the gradients when we use model.generate() efficiently?
> 
> Currently you can't back-prop when using generate as that method uses `torch.no_grad()` context manager. You can overwrite the method and make sure it is not using that context manager

@younesbelkada @ArthurZucker After commenting the `torch.no_grad()`, a huge amount of GPU memory is needed to do the `model.generate()` and a `CUDA out of memory` error would be easily triggered. Do you have any suggestions on this? 

Thank you very much, and have a nice day!

Best regards,

Shuyue
Dec. 4th, 2023