accelerator.save_state() will report error while i use accelerate and fsdp
Author: lplzyp
Pull request: no

### System Info

version:
pytorch 2.0.1
accelerate 0.21.0


other:
i have already print the `object_state.non_tensors`, and it shows a number 1, and object_state looks like that:
StateInfo(tensors={'exp_avg': _PosDimTensorInfo(shape=torch.Size([65491968]), dtype=torch.float32), 'exp_avg_sq': _PosDimTensorInfo(shape=torch.Size([65491968]), dtype=torch.float32)}, scalar_tensors={}, non_tensors=1)


### Who can help?

_No response_

### Information

- [X] 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

train a model with one machine with 4 gpu


report:
```
**File "/opt/conda/envs/niuniu/lib/python3.9/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1637, in _all_gather_optim_state
    for name, non_tensor_value in object_state.non_tensors.items():
    for name, non_tensor_value in object_state.non_tensors.items():AttributeError
: 'int' object has no attribute 'items'AttributeError
: 'int' object has no attribute 'items'**
```

### Expected behavior

solve it

Comments:
cc @pacman100 

same env, accelerate + deepspeed is fine to save model or checkpoint, and i try old version like 0.20.3, it performs well.

version:
pytorch 2.0.1
accelerate 0.22.0.dev0 
transformers 4.31.0

I also met the same problem. It saves rightly when on one GPU, but fails on many GPUs.
```
aving model checkpoint to ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300
Configuration saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/config.json
Configuration saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/generation_config.json
Model weights saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/pytorch_model.bin
tokenizer config file saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/tokenizer_config.json
Special tokens file saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/special_tokens_map.json
Copy vocab file to ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/spiece.model
Traceback (most recent call last):
  File "examples/trainer_accelerate.py", line 565, in <module>
Traceback (most recent call last):
  File "examples/trainer_accelerate.py", line 565, in <module>
    main()
  File "examples/trainer_accelerate.py", line 559, in main
    trainer.train(resume_from_checkpoint= args.resume_from_checkpoint)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1539, in train
    main()
  File "examples/trainer_accelerate.py", line 559, in main
    return inner_training_loop(
  File "/userhome/dsj/zip/accelerate-smangrul-fsdp-state-dict-fix/src/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1901, in _inner_training_loop
    trainer.train(resume_from_checkpoint= args.resume_from_checkpoint)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1539, in train
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2237, in _maybe_log_save_evaluate
    return inner_training_loop(
  File "/userhome/dsj/zip/accelerate-smangrul-fsdp-state-dict-fix/src/accelerate/utils/memory.py", line 136, in decorator
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2306, in _save_checkpoint
    return function(batch_size, *args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1901, in _inner_training_loop
    save_fsdp_optimizer(
  File "/userhome/dsj/zip/accelerate-smangrul-fsdp-state-dict-fix/src/accelerate/utils/fsdp_utils.py", line 137, in save_fsdp_optimizer
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2237, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2306, in _save_checkpoint
    save_fsdp_optimizer(
  File "/userhome/dsj/zip/accelerate-smangrul-fsdp-state-dict-fix/src/accelerate/utils/fsdp_utils.py", line 137, in save_fsdp_optimizer
    optim_state = FSDP.optim_state_dict(model, optimizer)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1753, in optim_state_dict
    optim_state = FSDP.optim_state_dict(model, optimizer)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1753, in optim_state_dict
    return FullyShardedDataParallel._optim_state_dict_impl(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1154, in _optim_state_dict_impl
    return FullyShardedDataParallel._optim_state_dict_impl(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1154, in _optim_state_dict_impl
    return _optim_state_dict(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1455, in _optim_state_dict
    return _optim_state_dict(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1455, in _optim_state_dict
    _gather_orig_param_state(
      File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1690, in _gather_orig_param_state
_gather_orig_param_state(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1690, in _gather_orig_param_state
    gathered_state = _all_gather_optim_state(fsdp_state, optim_state)    gathered_state = _all_gather_optim_state(fsdp_state, optim_state)

  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1637, in _all_gather_optim_state
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1637, in _all_gather_optim_state
    for name, non_tensor_value in object_state.non_tensors.items():
AttributeError: 'int' object has no attribute 'items'
    for name, non_tensor_value in object_state.non_tensors.items():
AttributeError: 'int' object has no attribute 'items'
  5%|▌         | 300/6000 [04:02<1:16:38,  1.24it/s]
```
And when I use accelerate==0.20.3. I met an another error.
```
  5%|████████                                                                                                                                                        | 300/6000 [04:20<1:07:07,  1.42it/sS
aving model checkpoint to ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300
Configuration saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/config.json
Configuration saved in ./trained_checkpoint/model-mt0-small-raw_dataset_name-opus9_2k_select_train_labse_dev_tst_trigger_3-num_epochs-3-src_max_length-256-max_new_tokens-256-num_beams-5-validation_ratio-0.1/checkpoint-300/generation_config.json
Traceback (most recent call last):
  File "examples/trainer_accelerate.py", line 565, in <module>
    main()
  File "examples/trainer_accelerate.py", line 559, in main
    trainer.train(resume_from_checkpoint= args.resume_from_checkpoint)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/opt/conda/lib/python3.8/site-packages/accelerate/utils/memory.py", line 132, in decorator
    return function(batch_size, *args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1901, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2237, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2294, in _save_checkpoint
    self.save_model(output_dir, _internal_call=True)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2751, in save_model
    save_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir)
NameError: name 'save_fsdp_model' is not defined
```

Hello, please let us know if the issue persists on the latest release of Transformers and Accelerate.

> same env, accelerate + deepspeed is fine to save model or checkpoint, and i try old version like 0.20.3, it performs well.

In my case, I still have the above problem.  Later, I only use `deepspeed  examples/xxx.py --args xx ...` to avoid this error, not the `accelerate launch --config_file=$config_file  examples/xx.py`. 


Transformers 4.31.0 requires accelerate > 0.20.3 for FSDP to work, you can see it in their import statements : https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/trainer.py#L202C1-L208C10

You can try downgrading transformers

Fixed the issue. By default HF Trainer uses adamw_hf, use the torch implementation of adamw where everything is implemented in tensors.

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.

这是来自QQ邮箱的假期自动回复邮件。您好，我最近正在休假中，无法亲自回复您的邮件。我将在假期结束后，尽快给您回复。

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.