Add fp8 (torchao)/fsdp2/torch_compile handlers and tests

[qingquansong]

What does this PR do?

[Closed the previous PR and open an new one as examples]
Add fp8/fsdp2/torch_compile handlers and tests. Since FP8 need to bundle with compile on same dedicated layers to achieve real memory reduction and speed up so put them together. FSPD2 is a plus since FSDP1 solution needs a bit hacky grad unrolling when doing with torch compile.

Note:

1. currently the FSDP2 handler are customized for HuggingFace Style LLM model architecture and only shard the layers defined in model.model.layers
2. To make FSDP1 work with FP8 for Mixtral related models with lightning we will need some dependency update including triton nightly change and set FP8 config pad inner dim to True (patch this change Support FP8 constant triton-lang/triton#4222)
3. To make FSDP1 work with torch compile for HF models, we need manual grad unrolling in both training step and validation steps, see more details here: torch.compile + FSDP1 CPU offloading + PT lightning validation loop throws an error pytorch/pytorch#139110 and solution will be given in the integration example.
4. Some FP8 dedicated functions for precomputing precompute_float8_dynamic_scale_for_fsdp and syncing sync_float8_amax_and_scale_history on Amax or scale values may need custom callbacks or change in model hooks, will share an example in the integration example.
5. Need a bit more effort on speed up CPU offloading with FSDP2: see torch.compile + FSDP1 CPU offloading + PT lightning validation loop throws an error pytorch/pytorch#139110
6. Model loading with FSDP2 sometimes consumes NaN issue when doing in configure_model function of lighting model model. Putting in the init function seem to consume more memory.

TODO:
Add examples with lightning training (FSDP1/2 + FP8 + Torch Compile)

Related PR: #20440

Fixes #<issue_number>

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--20445.org.readthedocs.build/en/20445/

[qingquansong]

Seems always being deleted 🤔 any ideas on how to keep it @lantiga