Model compilation support

Nobody has claimed this yet.

With FSDP currently the code could not be run.

If you try to add model compilation to the [training](https://github.com/Lightning-AI/lit-llama/blob/main/train.py) like:
```
...
fabric = L.Fabric(accelerator="cuda", devices=8, precision="bf16-mixed", strategy=strategy)
fabric.launch()
...

model = fabric.setup_module(model)
# compile() goes should go wrapping as per https://github.com/huggingface/transformers/commit/fb0a38b4f275727d6228fb4a78c15c6dd8480e91
# Though it does not work either even if goes before setup_module() as you'll get the same issue (see below)
model = torch.compile(model)

optimizer = torch.optim.AdamW(model.parameters(), ...)
optimizer = fabric.setup_optimizers(optimizer)
...

train(model, ...)
```

and try it via:

```
lightning run model --accelerator=cuda --devices=8 train.py ...
```

you'll get:

```
File ".../.venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 172, in __call__
return self._wrap(value).clone(**self.options())
File ".../.venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 345, in _wrap
assert getattr(
AssertionError: Dynamo only supports FSDP with use_orig_params=True
```

If I pass `use_orig_params = True` into the `FSDPStrategy()` constructor you get:

```
ValueError: The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the optimizer after setting up the model.
```

So, then one can remove `optimizer = fabric.setup_optimizers(optimizer)` as it is anyway no-op for FSDP, but even in this case I see:

```
from user code:
File ".../.venv/lib/python3.8/site-packages/lightning_utilities/core/apply_func.py", line 75, in apply_to_collection
is_namedtuple_ = is_namedtuple(data)
```
---

Also I was not able to find any tests of compiled model with the FSDP neither [here](https://github.com/Lightning-AI/lightning/blob/master/tests/tests_fabric/strategies/test_fsdp.py) nor [here](https://github.com/Lightning-AI/lightning/blob/master/tests/tests_fabric/strategies/test_fsdp_integration.py).

I wonder if anyone was able to successfully launch compiled model in a FSDP regime? Thanks a lot for the help!

---

P.S.: if i try to run similar code using HuggingFace trainer I run into the exact same `AssertionError: Dynamo only supports FSDP with use_orig_params=True` :)

Start by reproducing the failure in train.py with torch.compile and FSDPStrategy, using the reported command and configurations. Then inspect tests/tests_fabric/strategies/test_fsdp.py and test_fsdp_integration.py, where the issue notes that compiled-model coverage is missing. Done means model compilation launches successfully with FSDP and has regression coverage for the reported failures.

Written by the indexing model from the issue text.