fsdp with custom process groups

[vchiley]

What does this PR do?

This PR enables FSDP _auto_wrap to accept and propagate custom arguments, most notably, it enables the user to propagate dist process_groups.
Custom process_groups enable users to instantiate fsdp modules which shard / synchronize parameters over a subset of accelerators. This is useful for MoE Expert layers and TensorParallel layers.

What issue(s) does this change relate to?

Mixture of Experts (MoE)
https://mosaicml.atlassian.net/browse/RESEARCH-351
https://mosaicml.atlassian.net/browse/CO-1716
https://mosaicml.atlassian.net/browse/CO-1715
https://mosaicml.atlassian.net/browse/CO-1714
https://mosaicml.atlassian.net/browse/CO-1712

TensorParallel (TP)
https://mosaicml.atlassian.net/browse/CO-1635
https://mosaicml.atlassian.net/browse/RESEARCH-442

In general this doesn't fix an issue, but enables a feature which enable TP and MoE models
mosaicml/examples#180
also see: https://github.com/mosaicml/tutel/pull/1

btw I ran this "test" on 16 gpus with different process_group configurations to validate different configurations run. For the MoE setup with pg 'self' or 'set1', I do test if the different configurations result in the expert weights not being sync'd. I do not check if params / grads are sync'd in appropriately in the TP setting; since TP params are sharded, the params shouldn't be equal, so idk what to check, just the fact that they params are gathered, have the appropriate shape, and run might be test enough.

Before submitting

• Have you read the contributor guidelines?
• Did you update any related docs and document your change?
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

Tests: this PR enables "super-users" to use advanced features and defaults to previous behavior. All former tests passing is test enough.

[dskhudia]

This also allows user to specify per module sharding strategy as well?

[vchiley]

#2006 (comment)
Yes this enables the user to pass custom args to fsdp for every fsdp wrapped module; we include some of the mosaicml tooling for easily setting the configs
cc @dskhudia

[bcui19]

LGTM, will approve when we add docs

[mvpatel2000]

I'm going to hold this PR until we get multiple approval here just because it looks scary and might break release, so we should be extra careful. Will lift once we have sign-offs from all parties. This is basically in lieu of codeowners for this part of the repo, which isn't configured.

Requiring approvals from @bcui19 @dakinggg and myself.

[dakinggg]

I think this looks good to me, left a few minor comments. Could you also test that a mosaicgpt-125m runs and loss curve looks reasonable? (i.e. make sure this PR doesn't break stuff unrelated to this PR)

[vchiley]

@bcui19 added docs

@dakinggg made suggested updts, suggested runs from diff branches here (using tutel_moe branch in examples repo since it points to this branch).
Updt: those lines are soooo similar that I had to go verify that they were running from the correct branches 😅 . We're good

Are there two loss curves in the room with us now? 👻

[dakinggg]

LGTM

[bcui19]

LGTM!

[mvpatel2000]

LGTM.