Trainer/accelerate crashes when loading checkpoint using FSDP: sync_module_states ValueError
### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.31
- Python version: 3.11.4
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- 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?

@sgugger @pacman100 

### 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

1. Use `run_clm.py` (https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py) to train a large model using the HuggingFace Trainer, use FSDP and save checkpoints. For example:
```
torchrun --nproc_per_node=4 --master_port=XXXXX experiments/run_clm.py \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_name openwebtext \
    --streaming \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --do_train \
    --max_steps 1000 \
    --output_dir output_dir/ \
    --block_size 512 \
    --save_steps 10 \
    --save_total_limit 1 \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap "LlamaDecoderLayer" \
    --tf32 True \
    --bf16 True \
    --gradient_checkpointing \
```
2. Kill training after a checkpoint has been saved. Then, resume training from the checkpoint with the `resume_from_checkpoint` training argument.
3. Observed behavior: crashes when loading checkpoint model:
```
Traceback (most recent call last):                                                                                                                                            
  File ".../run_clm.py", line 638, in <module>                                                                                     
    main()                                                                                                                                                                    
  File ".../run_clm.py", line 584, in main                                                                                         
    main()                                                                                                                                                                    
  File ".../run_clm.py", line 584, in main                                                                                         
    train_result = trainer.train(resume_from_checkpoint=checkpoint)                                                                                                           
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                           
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train                                                 
    main()                                                                                                                                                                    
    train_result = trainer.train(resume_from_checkpoint=checkpoint)                                                                                                           
  File ".../run_clm.py", line 584, in main                                                                                         
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                           
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train                                                 
    train_result = trainer.train(resume_from_checkpoint=checkpoint)                                                                                                           
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                           
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train                                                 
    self._load_from_checkpoint(resume_from_checkpoint)                                                                                                                        
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint                                 
    self._load_from_checkpoint(resume_from_checkpoint)                                                                                                                        
      File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint                             
self._load_from_checkpoint(resume_from_checkpoint)                                                                                                                            
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint                                 
    load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)                                                                      
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 79, in load_fsdp_model
    load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)                                                                      
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 79, in load_fsdp_model
    load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)                                                                      
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 79, in load_fsdp_model
    raise ValueError(
ValueError: Set the `sync_module_states` flag to `True` so that model states are synced across processes when initializing FSDP object
    raise ValueError(
ValueError: Set the `sync_module_states` flag to `True` so that model states are synced across processes when initializing FSDP object
    raise ValueError(
ValueError: Set the `sync_module_states` flag to `True` so that model states are synced across processes when initializing FSDP object
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 45997 closing signal SIGTERM                                                                            
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 1 (pid: 45998) of binary: miniconda3.../bin/python
```

### Expected behavior

Expected behavior: can resume training from checkpoint using FSDP.
Hello, as the error mentions, `sync_module_states` needs to be True. For this please use the accelerate launcher with the related accelerate config as specified in the docs here: https://huggingface.co/docs/transformers/main_classes/trainer#using-accelerate-launcher-with-trainer<|||||>Thank you, I will try that. Is it no longer supported to use the Trainer without accelerate launch, such as with torchrun? Loading from checkpoints worked fine for me before I updated to the latest versions.<|||||>@pacman100 I tried using accelerate launch and setting `sync_module_states` to True, but I still cannot load checkpoints, but with a different error this time:
```
[INFO|trainer.py:2020] 2023-07-25 16:53:23,504 >> Loading model from weights/llama-2/debug/checkpoint-10/.
Traceback (most recent call last):
Traceback (most recent call last):
  File ".../run_clm.py", line 638, in <module>
  File ".../run_clm.py", line 638, in <module>
Traceback (most recent call last):
  File ".../run_clm.py", line 638, in <module>
        main()main()

    main()
  File ".../run_clm.py", line 584, in main
  File ".../run_clm.py", line 584, in main
  File ".../run_clm.py", line 584, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
                 train_result = trainer.train(resume_from_checkpoint=checkpoint) 
         ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^^^^^^^^^^^^^^^^ ^^ ^^ ^^ ^^ ^^ ^^ ^ ^^ ^^ ^^ ^^^ ^^ ^^ ^^ ^^ ^^ ^^ ^^ ^^^^^^^^^^^^^^^^^^^^
^^^^  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train
^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train
^^^^^^^^^^^^^^^^^^^^^^^
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 1528, in train
        self._load_from_checkpoint(resume_from_checkpoint)self._load_from_checkpoint(resume_from_checkpoint)

    self._load_from_checkpoint(resume_from_checkpoint)
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint
  File "miniconda3.../lib/python3.11/site-packages/transformers/trainer.py", line 2055, in _load_from_checkpoint
        load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)

    load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 73, in load_fsdp_model
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 73, in load_fsdp_model
  File "miniconda3.../lib/python3.11/site-packages/accelerate/utils/fsdp_utils.py", line 73, in load_fsdp_model
    with FSDP.state_dict_type(
  File "miniconda3.../lib/python3.11/contextlib.py", line 144, in __exit__
    with FSDP.state_dict_type(
  File "miniconda3.../lib/python3.11/contextlib.py", line 144, in __exit__
    with FSDP.state_dict_type(
  File "miniconda3.../lib/python3.11/contextlib.py", line 144, in __exit__
    next(self.gen)
    next(self.gen)
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 720, in state_dict_type
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 720, in state_dict_type
    next(self.gen)
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 720, in state_dict_type
    FullyShardedDataParallel.set_state_dict_type(
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 608, in set_state_dict_type
    FullyShardedDataParallel.set_state_dict_type(
    FullyShardedDataParallel.set_state_dict_type(
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 608, in set_state_dict_type
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 608, in set_state_dict_type
    state_dict_config_type = _state_dict_type_to_config[state_dict_type]
                           state_dict_config_type = _state_dict_type_to_config[state_dict_type] 
     ~~~~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ^ ^ ^ ^ ^ ^ ^~^~^~^~^~^~^~^~^~^~^~
~~~~~~KeyError~: ~None~
~~~~~~^^^^^^^^^^^^^^^^^
KeyError: None
    state_dict_config_type = _state_dict_type_to_config[state_dict_type]
                             ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
KeyError: None
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 109772 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 1 (pid: 109773) of binary: miniconda3.../bin/python
Traceback (most recent call last):
  File "miniconda3.../bin/accelerate", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "miniconda3.../lib/python3.11/site-packages/accelerate/commands/accelerate_cli.py", line 45, in main
    args.func(args)
  File "miniconda3.../lib/python3.11/site-packages/accelerate/commands/launch.py", line 966, in launch_command
    multi_gpu_launcher(args)
  File "miniconda3.../lib/python3.11/site-packages/accelerate/commands/launch.py", line 646, in multi_gpu_launcher
    distrib_run.run(args)
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/run.py", line 785, in run
    elastic_launch(
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "miniconda3.../lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 250, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:   
```

Here is the accelerate config I am using:
```yaml
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```
I'm really not sure what is going on, any help would be greatly appreciated!<|||||>Update: it seems that using `fsdp_state_dict_type=SHARDED_STATE_DICT` fixes this issue. Not sure if it is expected behavior to not work with `FULL_STATE_DICT`.<|||||>Hello, can you share the outputs of the saved checkpoint? This seems to be an issue with the PyTorch. <|||||>I experienced a similar problem, where using `SHARDED_STATE_DICT` instead of `FULL_STATE_DICT` allowed me to load the model when resuming from checkpoint. <|||||>Hi, @chenchenygu @pacman100  i am do the same things like you, could share your training srcipt of Acceleate;
Here is my training script:
```
accelerate launch --config_file training_scripts/accelerate_config.json train.py \
    --model_name_or_path ./pre-trained-model/huggyllama/llama-7b  \
    --train_file datasets/yesno_task/datatsets/train_pruning.json \
    --validation_file datasets/yesno_task/datatsets/valid.json \
    --bf16 True \
    --output_dir model/test \
    --num_train_epochs 3 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "epoch" \
    --save_total_limit 5 \
    --learning_rate 1e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 True \
    --max_length 512 \
    --gradient_checkpointing True 
```

I meet the follow warning:
```
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1531: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
  warnings.warn(
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1531: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
  warnings.warn(
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1531: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
  warnings.warn(
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1531: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
{'loss': 2.2814, 'learning_rate': 1e-05, 'epoch': 0.62}
2023-08-02 11:29:12,413 _dedup_tensors.py:44 INFO p:MainProcess t:MainThread: Duplicate keys to remove: {}
2023-08-02 11:29:52,569 _dedup_tensors.py:44 INFO p:MainProcess t:MainThread: Duplicate keys to remove: {}
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/checkpoint/filesystem.py:157: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  if tensor.storage().size() != tensor.numel():
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/checkpoint/filesystem.py:157: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  if tensor.storage().size() != tensor.numel():
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/checkpoint/filesystem.py:157: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  if tensor.storage().size() != tensor.numel():
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/checkpoint/filesystem.py:157: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
```

However, my training script seems work, but i want solve the above wanring.
Thank you for your help !!<|||||>Updata hit OOM in training process !
```
/root/zhuxuekai/anaconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/fsdp/_state_dict_utils.py:312: UserWarning: Failed to clone() tensor with name _fsdp_wrapped_module.model.layers.30.mlp.down_proj.weight on rank 0. This may mean that this state_dict entry could point to invalid memory regions after returning from state_dict() call if this parameter is managed by FSDP. Please check clone implementation of _fsdp_wrapped_module.model.layers.30.mlp.down_proj.weight. Error: CUDA out of memory. Tried to allocate 172.00 MiB (GPU 0; 47.54 GiB total capacity; 45.40 GiB already allocated; 25.62 MiB free; 46.42 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

The major error seem to be **Please check clone implementation of _fsdp_wrapped_module.model.layers.30.mlp.down_proj.weight. Error: CUDA out of memory. **<|||||>@Xuekai-Zhu i also ran into that issue, the workaround that worked for me is https://discuss.pytorch.org/t/fsdp-failed-to-save-model-checkpoints/178232/3
but this is a hacky fix, i think the ultimate fix that is needed is https://github.com/pytorch/pytorch/issues/98823#issuecomment-1504812144
but this requires changes to the accelerate/trainer code<|||||>> @Xuekai-Zhu i also ran into that issue, the workaround that worked for me is https://discuss.pytorch.org/t/fsdp-failed-to-save-model-checkpoints/178232/3 but this is a hacky fix, i think the ultimate fix that is needed is [pytorch/pytorch#98823 (comment)](https://github.com/pytorch/pytorch/issues/98823#issuecomment-1504812144) but this requires changes to the accelerate/trainer code

@chenchenygu  Huge thank to you !!! https://discuss.pytorch.org/t/fsdp-failed-to-save-model-checkpoints/178232/3 this is also wrok for me<|||||>Hello, with PR https://github.com/huggingface/transformers/pull/24926, this should be resolved, i.e., by default it will use `FULL_STATE_DICT` with cpu_offload on rank_0 only.