[Fix] Move init dist connection into the setup function

[SeanNaren]

What does this PR do?

Fixes #6318

This fix moves the init ddp connection for DDP into the setup function, and reorders the hook such that setup can now have access to the initialized distributed environment. This is also important for FSDP.

This fix however diverges DDP Spawn from DDP, and should be noted in the docs. As @ananthsub and @awaelchli have brought up, we may need to discuss the responsibility of hooks, primarily because we're seeing some inflexibility in the API when changing call orders of hooks/custom loops.

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 🙃

[awaelchli]

Some tests in tests\accererator\test_acceleratror_connector hang on my server. I'll investigate

[awaelchli]

DeepSpeed GPU tests had to be made into special tests, I was getting local errors since torch dist is being initialized and making them independent tests seems to fix them. Not high priority now, but we could wrap them potentially in a fixture such that torch distributed is deleted

We can't terminate the ddp connection after a trainer.fit call because we don't know if the user will call .fit or .test again right after it, and we want to enable that. The connection staying open could be a problem in pytest the way these tests are launched and stopped. A fixture could help us in the tests ensuring that no connections are open before running a new test, while still enabling the use case described above outside of the test environment.
One thing I would test is this: Take one of these deepspeed tests and copy the contents to a regular script and execute it. Does leave any connections open after it ends?

[awaelchli]

I have to call model_to_device twice in the single device plugin; once before and once within pre_dispatch. This is to handle the case if user's define layers within setup which need to be moved to the correct devices. This does mean if the user defines layer in setup, things can break using APEX since weights may not be transferred to half at initialization time. I'm going to investigate further to see if I can remove the model_to_device duplicate in the single device plugin, and fix the Apex plugin to handle this case independently

Maybe it helps to go back to before the accelerator refactor and look again what the order of calls were. The commit before acc refactor is 309ce7a
In the DDPAccelerator for example we can see the model is moved to the device after the setup hook is called.

[SeanNaren]

If anyone extended and made their own accelerator, this will be a breaking change so might need to handle a deprecation path here

[tchaton]

Do you need to rename them ?

[SeanNaren]

I don't i guess? It's just a bit concerning because if I don't rename them, connect_training_type_plugin will be calling plugin.setup, and connect will be calling training_type_plugin.connect. Just confusing function names. I think if it becomes an issue we can make this BW compatible however in most cases it seems users should be defining plugins, not accelerators.

[awaelchli]

yeah my silly todo....
"need to double check that it is the correct place"

Thanks for double checking @SeanNaren 😄

[tchaton]

LGTM ! Great job !

[ananthsub]

🚀

[Borda]

btw, shall we have this as a single message intend for 4 separate?

[Borda]

this is so cool :D

[carmocca]

Setting milestone as 1.3 as it requires separating the fix from the API change to get into 1.2.x.