trainer fails when using torchrun for distributed run  of  transformer model wrapped with PEFT without using device_map
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?

@pacman100



### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Background about my target: 
I want to do distributed training of new architecture of LLaMA model wrapped with peft. 
Basically,  I have compressed  llama2  and modified its architecture. For saving the new model I am using torch.save and torch.load for loading the model. Hence, I am not able to use .pretrained_method for loading the model and setting device_map ='auto' for further distributed fine-tuning. I am not sure if there is  a still a way for  automatic using  the device map. 
What i did, I tried to use torchrn for the distribuited run of  the original LLaMA2 model wrapped with LoRA to see if i can use the same approach to my modified model. 

Steps:
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/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._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

I was expecting successful run of the distribuited training of LLaMA2 with PEFT similar to LLaMA2 without parameter efficient finetuning methods 

Comments:
@muellerz @pacman100, could you help me with above issue?. I would appreciate your help

Gentle ping @pacman100 @muellerzr 

Hello, when using FSDP+PEFT, you need to have a different Auto wrap policy. See these lines on how to do that: https://github.com/pacman100/DHS-LLM-Workshop/blob/3c34a1215d7b5f73bb4a1ec122011b1de1470559/chat_assistant/sft/training/train.py#L168-L172

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.