Enable Unwrapping for Model State Dicts (FSDP)

[alex-jw-brooks]

What does this PR do?

Fixes #2958

There are more details and a repro case in the issue above - it appears the unwrap args is unused when getting the state dict and running FSDP, which seems to cause shared pointers/tied weights to behave inconsistently.

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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[muellerzr]

Thanks! Makes sense to me :)

(Test failures are unrelated)

[alex-jw-brooks]

Awesome, thank you! 😄

[evkogs]

Dude its wrong...

This breaks FSDP

Instead model should be unwraped after FSDP figured out what was the config, like this:

        elif self.distributed_type == DistributedType.FSDP:
            from torch.distributed.fsdp import FullStateDictConfig, StateDictType
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
            with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config):

                print(f"using this config: {StateDictType.FULL_STATE_DICT}")

                if unwrap:
                    model = self.unwrap_model(model)

                state_dict = model.state_dict()

Otherwise with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config): receive regular torch module and has no idea how to handle it

I think we need to create a new pool request but with more robust testing, I'm still unsure that my fix is correct