RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn while model parameters requires_grad is True 
Author: zhongshsh
Pull request: no

### System Info

- `transformers` version: 4.36.2
- Platform: Linux-5.15.0-18-shopee-generic-x86_64-with-glibc2.35
- Python version: 3.11.5
- Huggingface_hub version: 0.20.2
- Safetensors version: 0.4.0
- Accelerate version: 0.27.0.dev0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 8
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - 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.1.1+cu121 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker 

### Information

- [X] The official example scripts
- [ ] 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 load `Mixtral` using device_map as `auto` (have to use `auto` otherwise OOM), then I make some of the model parameters' grad `True` (e.g., model.model.layers[0]), and use `Trainer` to fine-tune it. 

I launch the code by `python xx.py` as explaining in https://github.com/huggingface/accelerate/issues/1840#issuecomment-1683105994.

error message I got as follows: 
```
Traceback (most recent call last):
  File "xx.py", line 461, in <module>
    trainer.train()
  File "miniconda3/lib/python3.11/site-packages/transformers/trainer.py", line 1537, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "miniconda3/lib/python3.11/site-packages/transformers/trainer.py", line 1864, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "miniconda3/lib/python3.11/site-packages/transformers/trainer.py", line 2763, in training_step
    self.accelerator.backward(loss)
  File "miniconda3/lib/python3.11/site-packages/accelerate/accelerator.py", line 1964, in backward
    loss.backward(**kwargs)
  File "miniconda3/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "miniconda3/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
```


If I set other module such as `lm_head` grad as True, the code can run successfully. I guess whether the above error caused by setting `device_map` as `auto`?

### Expected behavior

make some of `Mixtral` parameters' grad `True`, and use `Trainer` to fine-tune it. 

Comments:
I solved the problem by setting gradient_checkpointing_kwargs={"use_reentrant":False}, but I don’t know why.

cc @younesbelkada 

@ArthurZucker hi, I also have another problem. When I loaded Mixtral using `device_map` as `auto` in a node of 8 GPUs (A100 80G), it would not use the 8-th GPU although I changed the `device_map` to `balance`. 

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')
```

How can I do to load the model in all GPUs balancedly? 

I tried to use `max_memory` to solve my problem, but it would cause other problems when using Trainer. For example, I couldn't save the model when using RAM (although I set the `max_memory` to be large (14G per GPU), it still used RAM).

Hi @zhongshsh 
Indeed to fix that issue you need either to
1- call `model.enable_input_require_grads()`
or
2- pass `gradient_checkpointing_kwargs={"use_reentrant":False}` 
Regarding your second point have you tried with `"balanced"` (instead of `"balance"`)? Although I think `"auto"` should evenly split the model across all available GPUs cc @SunMarc 

Hi @younesbelkada thx for your reply. 

Yes, I try `balanced`. You can reproduce this problem easily by using code 
```
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')  # or balanced
```

cc @younesbelkada 

Hi @zhongshsh 
Thanks ! 
I think accelerate is hitting a corner case with mixtral when computing automatically the device map - can you try to set a custom `max_memory` map ? Something like `max_memory={0:"20GB", ... 7: "20GB"}` and pass it to `from_pretrained`

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.