Correctly set CUDAGuard in nccl collectives
#130921
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/130921
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit e7cc8cf with merge base 2300bb2 ( NEW FAILURE - The following job has failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
77cf8b0 to
c74fb7a
Compare
c74fb7a to
3f0908c
Compare
CUDAGuard in PGNCCLCUDAGuard in nccl collectives
|
I guess the |
That's true. But we used to have a (currently broken) heuristic that tries to do that when the code do not set device explicitly, I'd prefer to fix the heuristic rather than leave it not working and say "this is not the recommend way". Also the warning is not printed with the default config. I add a commit to print it without extra env vars. |
| @@ -3912,13 +3910,13 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::barrier(const BarrierOptions& opts) { | |||
| // ensure that each process is on a different GPU | |||
| auto numGPUs = at::cuda::getNumGPUs(); | |||
| int16_t deviceIdx = static_cast<int16_t>(rank_ % numGPUs); | |||
| LOG(INFO) | |||
| LOG(WARNING) | |||
There was a problem hiding this comment.
I think this change makes sense to me.
|
Any update on this? |
| @@ -2752,8 +2752,6 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collectiveCoalesced( | |||
| std::make_shared<std::vector<at::Tensor>>(inputs); | |||
| } | |||
|
|
|||
| at::cuda::OptionalCUDAGuard gpuGuard; | |||
|
|
|||
There was a problem hiding this comment.
Why is this guard being removed instead of updated to device?
There was a problem hiding this comment.
The AutoNcclGroup below does that
| @@ -2591,7 +2591,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective( | |||
| work->stashed_for_allocator_safety_->push_back(input); | |||
| } | |||
|
|
|||
| at::cuda::OptionalCUDAGuard gpuGuard; | |||
There was a problem hiding this comment.
Is it the case that gpuGuard without index is not effective? Or does it default to device 0?
|
@kwen2501 thanks for your review anyway. I'd appreciate it if you could review this earlier, and your time could also be saved. |
|
Sorry I was on leave. |
Fixes #130414
#119421 removes some
CUDAGuard.set_indexcalls, leave somegpuGuardinitialized but not updated.That adds implicit call to
cudaSetDevice(0)on other ranks, and therefore lead to #130414.This PR add them back.
cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o