Fix FSDPv2 checkpoint saving on TPU by using recursive unwrap

[Nikhil172913832]

What does this PR do?

This PR fixes checkpoint saving for FSDPv2 (SPMD) on TPU by properly unwrapping nested FSDP wrappers before extracting the model state dict.

When using FSDPv2 on TPU, models have nested FSDP wrappers around each transformer layer. The previous implementation only unwrapped the top-level wrapper, causing the saved checkpoint to contain wrapped state dict keys instead of the actual model parameters. This resulted in:

• PEFT adapters not being saved in the correct format
• Model weights appearing unchanged after training
• Missing adapter keys when loading checkpoints

The fix uses unwrap_model with recursive=True specifically for FSDPv2 to unwrap all nested wrappers, then extracts the state dict from the fully unwrapped model. This ensures clean parameter keys in saved checkpoints while maintaining backward compatibility with FSDPv1 and other training configurations.
#36004
Fixes #36004

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc @muellerzr

[Nikhil172913832]

Hi @SunMarc @muellerzr just wanted to check if you could take a look at this PR when you get a chance.