Support TP + FSDPv2 / HSDP or just FSDPv2 / HSDP

[kmehant]

What does this PR do?

1. This PR proposes having FSDP2 as a separate distributed type for accelerate residing along with the FSDPv1 implementation.
2. Furthermore, the PR also proposes use of prepare_nd_device_mesh util function to extend creation of device meshes for any combination of parallelisms. Currently it supports any combination of TP and FSDP/HSDP
3. This PR should potentially supersede [RFC] Support FSDP2 #3231

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?

@SunMarc @muellerzr
@kwen2501 from PyTorch

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[kmehant]

pulse

[SunMarc]

cc @S1ro1

[S1ro1]

FSDP2 is close to done in #3394. Then I'll take a look at supporting HSDP and including TP. Though this PR probably also gets into the issue of increased memory usage for FSDP2 because of creating optimizer on a full (non-sharded) model, discussed here

@kmehant let me know if this works for your PR as expected, but it shouldn't.

model = ...
optimizer = ...(model.parameters(),...)

model, optimizer = accelerate.prepare(model, optimizer)

This should result in a higher memory usage as the optimizer holds the original model parameters.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[kmehant]

Being tracked at #3498 :)