Training with FSDP slows down the convergence speed.
Author: yuangpeng
Pull request: no

### System Info

- `transformers` version: 4.35.2
- Platform: Linux-5.4.143-2-velinux1-amd64-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.23.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: FSDP
        - mixed_precision: no
        - use_cpu: False
        - debug: False
        - num_processes: 8
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - fsdp_config: {'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP', 'fsdp_backward_prefetch_policy': 'NO_PREFETCH', 'fsdp_forward_prefetch': False, 'fsdp_offload_params': False, 'fsdp_sharding_strategy': 2, 'fsdp_state_dict_type': 'SHARDED_STATE_DICT', 'fsdp_sync_module_states': True, 'fsdp_transformer_layer_cls_to_wrap': '', 'fsdp_use_orig_params': False}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.1.2+cu118 (False)
- 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?: A800
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@pac

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

A group of same experiments with purple and green, another group of sample experiments with yellow and blue. The other parameters are the same, one uses fsdp and the other uses deepspeed.

fsdp and deepspeed have similar losses early step, with differences appearing after about 20 steps. It’s not that fsdp is not converging, fsdp is also slowly converging.

fsdp config:
```
TrainingArguments.fsdp="shard_grad_op auto_wrap"
TrainingArguments.fsdp_config=dict(fsdp_transformer_layer_cls_to_wrap=["LlamaDecoderLayer"])
```

deepspeed config:
```
{
    "bf16": {
        "enabled": true
    },
    "train_micro_batch_size_per_gpu": "auto",
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto"
    }
}
```

<img width="1442" alt="image" src="https://github.com/huggingface/transformers/assets/57125678/1894247d-7e5f-423f-83e8-72aa14cfb4e6">


### Expected behavior

Expect similar loss curves using fsdp and deepspeed

Comments:
cc @pacman100 

Although I use my own task and use my own dataset, I only added a few modules and training with full unfreeze. If you need a detailed script, I will delete the private content before uploading it in these days. @pacman100 

Hope it can solved soon!

Can you provide a minimal reproducer for this. Using the latest releases, everything is as expected. Below is the fine-tuning of Mistral on ultrachat SFT dataset. 

![Screenshot 2023-12-27 at 8 48 12 PM (1)](https://github.com/huggingface/transformers/assets/13534540/1865c960-a78a-46bd-9ee5-2a12e0ba68d2)




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.