Cannot save Adafactor optimizer when using Trainer with accelerate and fsdp
Author: thangphan0168
Pull request: no

### System Info

version:
pytorch==2.0.1
transformers==4.34.0
accelerate==0.23.0 


### Who can help?

_No response_

### Information

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

 ```python
  # Import stuff
  ... 
  # Arguments
  args = TrainingArguments(
          output_dir="/checkpoints/test_train_llm", # set to "/checkpoints/train_llm" when training
          per_device_train_batch_size=4,
          logging_steps=1, # Set to 50 when training
          max_steps=4, # Set to 10_000 when training
          gradient_accumulation_steps=8,
          weight_decay=0.1,
          warmup_ratio = 0.01,
          lr_scheduler_type="cosine",
          learning_rate=1e-5,
          save_steps=2, # Set to 500 when training
          fp16=True,
          push_to_hub=False,
          gradient_checkpointing=True,
          optim="adafactor",
          save_total_limit=5
  )
  model_name = "gpt2"
  tokenizer = AutoTokenizer.from_pretrained(model_name,
                                  local_files_only=False,
                                  cache_dir='/checkpoints')
  model = AutoModelForCausalLM.from_pretrained(model_name,
                                  local_files_only=False,
                                  cache_dir='/checkpoints')
  # Load and prepare dataset
  lm_datasets = ...
  # Train
  trainer = Trainer(
      model=model,
      args=args,
      train_dataset=lm_datasets,
      data_collator=default_data_collator
  )
result = trainer.train()
```

It seems that the Adafactor optimizer is not able to save
```python
Traceback (most recent call last):
  File "/app/main3.py", line 214, in <module>
    main()
  File "/app/main3.py", line 209, in main
    result = trainer.train()
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1591, in train
    return inner_training_loop(
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1984, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2339, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2408, in _save_checkpoint
    save_fsdp_optimizer(
  File "/opt/conda/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py", line 138, in save_fsdp_optimizer
    optim_state = FSDP.optim_state_dict(model, optimizer)
  File "/opt/conda/lib/python3.10/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.10/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.10/site-packages/torch/distributed/fsdp/_optim_utils.py", line 1455, in _optim_state_dict
    _gather_orig_param_state(
  File "/opt/conda/lib/python3.10/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)
  File "/opt/conda/lib/python3.10/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'
``` 

My config file to run with `accelerate launch`:
```
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```


### Expected behavior

The optimizer is saved successfully and the model is continued to be trained

Comments:
It seems that this will only be an issue when `fsdp_use_orig_params` is set to `true`, I was able to run the script when it is set to `false`