[Model Parallel] Add configure sharded model hook

[kaushikb11]

What does this PR do?

Adds a configure sharded model hook. This is required for both DeepSpeed and Fully Sharded. Both teams have defined a context in which you can wrap model layers to be sharded instantly.

Why does it need to be a hook?

It needs to be a hook because model layers can only be sharded instantly after torch distributed is setup. This ensures that the distributed environment is setup, then calls the hook.

We also wrap the hook with the model parallel context for DeepSpeed/Fully Sharded so the user doesn't need to do anything extra.

What about fit/test/predict?

It will be treated like the setup hook for now, where we assume we're either using fit/test/predict and can call the hook for the user. We can iterate on this in further PRs.

Before submitting

• Was this discussed/approved 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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

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

• 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

Did you have fun?

Make sure you had fun coding 🙃

[pep8speaks]

Hello @kaushikb11! Thanks for updating this PR.

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2021-03-29 19:15:47 UTC

[codecov]

Codecov Report

Merging #6679 (aa35583) into master (646cf2f) will decrease coverage by 4%.
The diff coverage is 98%.

@@           Coverage Diff           @@
##           master   #6679    +/-   ##
=======================================
- Coverage      91%     87%    -4%     
=======================================
  Files         192     192            
  Lines       12189   12231    +42     
=======================================
- Hits        11143   10659   -484     
- Misses       1046    1572   +526     

[tchaton]

Add a check that the context manager is actually being yield.

[kaushikb11]

Thanks, it made me see an Issue in implementation. Fixinggg

[tchaton]

Might want to check the hook is not called within the callbacks.

[kaushikb11]

I think we can assume Users would be aware of the implications.

[carmocca]

But those two snippets you shared don't use on_model_parallel_setup, they use model_parallel_context.

I think the question is more along the lines of "when would you override on_model_parallel_setup?"

[tchaton]

LGTM !

[ananthsub]

@SeanNaren @kaushikb11 why's this needed in the trainer? could this be part of the accelerator setup? in L440 instead?

[SeanNaren]

which part exactly? the trainer has control over the hook, we do not want the hook to be called within the accelerator i think

[ananthsub]

I see, on first read I thought that we could do this:

self.call_setup_hook(model)
self.accelerator.setup(self, model)  -> inside here apply the model sharded context and call model.configure_sharded_model()

# now that accelerator has finished doing the model sharding setup call the callbacks
self.configure_sharded_model(model)

the upside is that more of the training type + accelerator logic lives inside of the accelerator internals
but the downside is that this splits the callback hook from the sharded model setup which can be less convenient

[SeanNaren]

Agree as long as we're ok with calling the model hook within the accelerator (cc @justusschock @awaelchli)