2x slower training speed with FSDP when switching from lightning 1.9 to 2.0

Nobody has claimed this yet.

### Bug description

Hello! Thank you for the integration of fsdp in the lightning trainer - it's a game changer.

I tried to switch from `lightning==1.9.4` to the newest `lightning==2.0.4` but observed a significant slow down in training in your `lit-GPT` [repository](https://github.com/Lightning-Universe/lit-GPT/tree/main).

Previously, with `lightning==1.9.4` and `torch==2.0.1` running the command `python train.py --implementation nanogpt --batch_size 1 --block_size 8192 --strategy fsdp-gpt --model_type gpt2-xl` the training speed would be 4.04s/iteration and GPU memory consumption 4.6GB (on a 8x A100 80GB gpus machine).

When updating to `lightning==2.0.4` (and same version of torch), I had to switch from using the default `from pytorch_lightning.strategies.fully_sharded_native import DDPFullyShardedNativeStrategy` to `pl.strategies.FSDPStrategy` here https://github.com/Lightning-Universe/lit-GPT/blob/main/lightning_gpt/models.py#L287
Running the same command I now get: 9.52s/it and 6.7B memory consumption.

That's a 2x training slowdown and 47% memory increase 😢 Would you have any idea why?

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

```python
python train.py --implementation nanogpt --batch_size 1 --block_size 8192 --strategy fsdp-gpt --model_type gpt2-xl
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0): 2.0.4
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0): 2.0.1
#- Python version (e.g., 3.9): 3.9
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version:
#- GPU models and configuration: 8x A100 80GB
#- How you installed Lightning(`conda`, `pip`, source): pip in a fresh conda environment
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

cc @borda @awaelchli @carmocca

Start at lightning_gpt/models.py around line 287, where lit-GPT switches from the Lightning 1.9 strategy to pl.strategies.FSDPStrategy. Reproduce with the provided train.py command on the stated 8x A100 environment, comparing Lightning 1.9.4 and 2.0.4. Done means identifying the source of the speed and memory regression and verifying the result against the reported iteration time and memory usage.

Written by the indexing model from the issue text.