Loading large models with fabric, FSDP and empty_init=True does not work

Nobody has claimed this yet.

### Bug description

I'm following [litgpt for how to load a large model with FSDP](https://github.com/Lightning-AI/litgpt/blob/main/litgpt/finetune/full.py#L129) and I'm getting an error related to shapes. See example script and error message

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

v2.2

### How to reproduce the bug

```python
#!/usr/bin/env python
import os.path as osp
import torch
from lightning.fabric.strategies import DDPStrategy, FSDPStrategy
from litgpt.lora import Block as LoraBlock
from litgpt.model import Block
import lightning as L
from litgpt.model import GPT
from litgpt.lora import Config

from torch.distributed.fsdp import MixedPrecision

def main():

strategy = FSDPStrategy(auto_wrap_policy={Block, LoraBlock},
sharding_strategy='HYBRID_SHARD',
mixed_precision=torch.distributed.fsdp.MixedPrecision(param_dtype=torch.bfloat16),
use_orig_params=True)

fabric = L.Fabric(
accelerator="gpu",
num_nodes=1,
devices=8,
strategy=strategy,
precision='bf16-mixed')

fabric.launch()

llm_config_path = '.../checkpoints/mistralai/Mistral-7B-v0.1/model_config.yaml'
config = Config.from_file(llm_config_path)
with fabric.init_module(empty_init=True):
model = GPT(config)
print('setting up')
fabric.setup_module(model)

llm_path = osp.join(osp.dirname(llm_config_path), "lit_model.pth")
fabric.load_raw(llm_path, model)
print('done')

main()
```

### Error messages and logs

```
File "test_init.py", line 42, in
main()
File "test_init.py", line 38, in main
fabric.load_raw(llm_path, model)
File "lightning/fabric/fabric.py", line 816, in load_raw
self._strategy.load_checkpoint(path=path, state=obj, strict=strict)
File "lightning/fabric/strategies/fsdp.py", line 548, in load_checkpoint
_load_raw_module_state_from_path(path, module=state, world_size=self.world_size, strict=strict)
File "/lightning/fabric/strategies/fsdp.py", line 888, in _load_raw_module_state_from_path
_load_raw_module_state(state_dict=_lazy_load(path), module=module, world_size=world_size, strict=strict)
File "lightning/fabric/strategies/fsdp.py", line 896, in _load_raw_module_state
module.load_state_dict(state_dict, strict=strict)
File "torch/nn/modules/module.py", line 2153, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for GPT:
size mismatch for lm_head.weight: copying a param with shape torch.Size([32000, 4096]) from checkpoint, the shape in current model is torch.Size([32768512]).
size mismatch for transformer.wte.weight: copying a param with shape torch.Size([32000, 4096]) from checkpoint, the shape in current model is torch.Size([0]).
size mismatch for transformer.ln_f.weight: copying a param with shape torch.Size([4096]) from checkpoint, the shape in current model is torch.Size([0]).
```

### Environment

```* Lightning:
- lightning: 2.2.3
- lightning-cloud: 0.5.64
- lightning-utilities: 0.9.0
- lion-pytorch: 0.1.4
- pytorch-lightning: 2.2.1
- pytorch-wpe: 0.0.1
- torch: 2.2.2+cu121
```

### More info

Looking into the implementation of `fabric.load_raw`, it seems to f[irst unwrap compiled objects and then call ](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/fabric.py#L802)`‎FSDPStrategy.load_checkpoint`, looking at the latter there is the following sentence in the doc string:

```
The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a directory of multiple files rather than a single file.
```
which seems to contradict the doc of `load_raw()`:

```
Use this for loading a raw PyTorch model checkpoint created without Fabric.
This is conceptually equivalent to ``obj.load_state_dict(torch.load(path))``, but is agnostic to the strategy
being used.
```
Very confused what the right approach is supposed to be.

cc @justusschock @lantiga

Start with the provided reproduction script and trace `fabric.load_raw` into `FSDPStrategy.load_checkpoint`, `_load_raw_module_state_from_path`, and `_load_raw_module_state`. Compare the documented raw-checkpoint behavior with the observed shape mismatches; done means the supported loading path is corrected or the documentation and behavior clearly agree, with a regression test for the reported case.

Written by the indexing model from the issue text.