Trainer fails when using torchrun for distributed run of transformer model wrapped with PEFT 
Author: Ahmed-Roushdy
Pull request: no

### System Info

- `transformers` version: 4.33.3                                                                      [1/1854]
- Platform: Linux-5.8.0-59-generic-x86_64-with-glibc2.31                                                      
- Python version: 3.9.18
- Huggingface_hub version: 0.19.2
- Safetensors version: 0.4.0
- Accelerate version: 0.24.1
- Accelerate config:    not found
- 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?

 @muellerzr 


### Information

- [X] The official example scripts
- [ ] 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

The following code train LLaMA2 transformer model with PEFT using  torchrun
The code snippet 
```ruby
   model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    print(training_args)

    print('Start Loading Model')
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,            )    
  
    config = LoraConfig(
        r=8,
        lora_alpha=16,
        lora_dropout=0.1,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, config)
    model.print_trainable_parameters()
    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)

```

scripet to run 
```
torchrun \
    --standalone \
    --nnodes=1 \
    --nproc-per-node=7 \
    train.py \
        --model_name_or_path "meta-llama/Llama-2-7b-hf" \
        --bf16 True \
        --output_dir checkpoints/dist-LLaMa-7B \
        --num_train_epochs 3 \
        --per_device_train_batch_size 1 \
        --per_device_eval_batch_size 1 \
        --gradient_accumulation_steps 8 \
        --fsdp "full_shard auto_wrap" \
        --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
        --tf32 True
```


I got the following error 

```

Start building the trainer module
Traceback (most recent call last):
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 285, in <module>
    train()
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 278, in train
    trainer.train()
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
  File "/home//.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1675, in _inner_training_loop
    self.model = self.accelerator.prepare(self.model)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1288, in prepare
    result = tuple(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1289, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1094, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1464, in prepare_model
    model = FSDP(model, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 391, in __init__
    _auto_wrap(auto_wrap_kwargs, fsdp_kwargs, FullyShardedDataParallel)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 73, in _auto_wrap
    _recursive_wrap(**auto_wrap_kwargs, **fsdp_kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  [Previous line repeated 2 more times]
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 388, in _recursive_wrap
    return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 317, in _wrap
    return wrapper_cls(module, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 408, in __init__
    _init_param_handle_from_module(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 429, in _init_param_handle_from_module
    _init_param_handle_from_params(state, managed_params, fully_sharded_module)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 525, in _init_param_handle_from_params
    handle = FlatParamHandle(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 366, in __init__
    self._init_flat_param(params, fully_sharded_module, use_orig_params)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 440, in _init_flat_param
    raise ValueError(
ValueError: `FlatParameter` requires uniform `requires_grad`
Traceback (most recent call last):
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 285, in <module>
Traceback (most recent call last):
    train()
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 285, in <module>
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 278, in train
    trainer.train()
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1556, in train
    train()
  File "/vault/aelkordy/NLP_projects/pruning/MAmmoTH/train.py", line 278, in train
    trainer.train()
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1675, in _inner_training_loop
    return inner_training_loop(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/transformers/trainer.py", line 1675, in _inner_training_loop
    self.model = self.accelerator.prepare(self.model)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1288, in prepare
    result = tuple(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1289, in <genexpr>
    self.model = self.accelerator.prepare(self.model)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1288, in prepare
    result = tuple(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1289, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1094, in _prepare_one
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1094, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1464, in prepare_model
    return self.prepare_model(obj, device_placement=device_placement)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/accelerate/accelerator.py", line 1464, in prepare_model
    model = FSDP(model, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 391, in __init__
    _auto_wrap(auto_wrap_kwargs, fsdp_kwargs, FullyShardedDataParallel)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 73, in _auto_wrap
        model = FSDP(model, **kwargs)_recursive_wrap(**auto_wrap_kwargs, **fsdp_kwargs)

  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 391, in __init__
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    _auto_wrap(auto_wrap_kwargs, fsdp_kwargs, FullyShardedDataParallel)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 73, in _auto_wrap
    _recursive_wrap(**auto_wrap_kwargs, **fsdp_kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  [Previous line repeated 2 more times]
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 388, in _recursive_wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 370, in _recursive_wrap
    return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 317, in _wrap
    wrapped_child, num_wrapped_params = _recursive_wrap(
  [Previous line repeated 2 more times]
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 388, in _recursive_wrap
    return wrapper_cls(module, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 408, in __init__
    return _wrap(module, wrapper_cls, **kwargs), nonwrapped_numel
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/wrap.py", line 317, in _wrap
    _init_param_handle_from_module(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 429, in _init_param_handle_from_module
    return wrapper_cls(module, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 408, in __init__
    _init_param_handle_from_params(state, managed_params, fully_sharded_module)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 525, in _init_param_handle_from_params
    _init_param_handle_from_module(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 429, in _init_param_handle_from_module
    handle = FlatParamHandle(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 366, in __init__
    _init_param_handle_from_params(state, managed_params, fully_sharded_module)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/_init_utils.py", line 525, in _init_param_handle_from_params
    self._init_flat_param(params, fully_sharded_module, use_orig_params)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 440, in _init_flat_param
    handle = FlatParamHandle(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 366, in __init__
    raise ValueError(
ValueError: `FlatParameter` requires uniform `requires_grad`    
self._init_flat_param(params, fully_sharded_module, use_orig_params)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 440, in _init_flat_param
    raise ValueError(
ValueError: `FlatParameter` requires uniform `requires_grad`
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 1854106) of binary: /home/aelkordy/.conda/envs/mamoth/bin/python
Traceback (most recent call last):
  File "/home/aelkordy/.conda/envs/mamoth/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/run.py", line 794, in main
    run(args)
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/run.py", line 785, in run
    elastic_launch(
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/aelkordy/.conda/envs/mamoth/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 250, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
train.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2023-11-16_10:52:13
  host      : g1lmd1
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 1854107)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[2]:
  time      : 2023-11-16_10:52:13
  host      : g1lmd1
  rank      : 2 (local_rank: 2)
  exitcode  : 1 (pid: 1854108)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-11-16_10:52:13
  host      : g1lmd1
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 1854106)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html

```

### Expected behavior

Successful run of the distribuited training of LLaMA2 with PEFT similar to LLaMA2 without parameter efficient fine-tuning methods 

Comments:
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.

cc @pacman100 or @BenjaminBossan 

@Ahmed-Roushdy did you see these threads?

- https://github.com/pytorch/pytorch/issues/100945
- https://github.com/pytorch/pytorch/issues/104690

Perhaps the info there could help. You should not have to use a nightly PyTorch build if you use the latest PyTorch version.

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.