fix NCCL error with non-consecutive trainer gpus

[awaelchli]

What does this PR do?

Fixes #8139

Can be reproduced on master with the following command:

NCCL_DEBUG=INFO python pl_examples/basic_examples/simple_image_classifier.py --trainer.gpus "1,3"  --trainer.accelerator  ddp

The issue appeared after #7202, where a barrier was introduced right after init_process group but before the ddp model gets configured with device_ids. When calling torch.distributed.barrier on a number of devices, it is assumed that it is affecting devices 0, 1, 2, etc. However, when we want to trainer on gpus=[1,3] for example that's not the case.
The fix is to use the device_ids argument in the barrier call. I do not know why and how this works, don't ask me. But if you know why, let me know.

Partially related, which need to be informed about this fix:

• CUDA OOM when initializing DDP #4705
• CUDA OOM when using "ddp" mode in training  #7817

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)
  Note: this cannot really be tested, needs a > 2 GPUs machine!
• 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 🙃

[codecov]

Codecov Report

Merging #8165 (8677cba) into master (c4492ad) will decrease coverage by 5%.
The diff coverage is 80%.

@@           Coverage Diff           @@
##           master   #8165    +/-   ##
=======================================
- Coverage      93%     88%    -5%     
=======================================
  Files         211     211            
  Lines       13450   13456     +6     
=======================================
- Hits        12486   11841   -645     
- Misses        964    1615   +651     

[kaushikb11]

@awaelchli Should we add a test for this?

[awaelchli]

@kaushikb11 I can add a test but it will require >3 gpus and will never run on our CI (which only has 2). We could then ask Adrian nicely to manually run this test on every ddp-related PR, what do you think? xD

[justusschock]

@awaelchli could we maybe mock the barrier call just to see if it is actually called with the devices instead of just the number?

[tchaton]

Wow, great fix ! Awesome work @awaelchli :)

[pep8speaks]

Hello @awaelchli! Thanks for updating this PR.

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

Comment last updated at 2021-06-28 16:29:25 UTC

[awaelchli]

@kaushikb11 @justusschock thanks for your suggetions. I added a test, it's not particularly smart but it fails on master with NCCL error and passes here. Again, it will not run in our ci due to the gpu requirements, but I tested on another server by running the special test directly (master and this PR).

[ananthsub]

great catch @awaelchli !
I think there's a logline from torch.distributed that says when device ids is not explicitly passed, its inferred from the local rank, but I need to double check this.

[awaelchli]

@ananthsub This was also my suspicion, and I first tried to set torch.cuda.set_device(local_rank) before the first barrier call without success.

In the torch.distributed docs it is clearly written that it is the local rank by default for the DistributedDataParallel wrapper.

Please ensure that device_ids argument is set to be the only GPU device id that your code will be operating on. This is generally the local rank of the process. In other words, the device_ids needs to be [args.local_rank], and output_device needs to be args.local_rank in order to use this utility.

And it seems the barrier(device_ids=) need to match this. However, it is a bit a puzzle for me why the barrier needs this when other collective calls don't.