8000 Correctly set `CUDAGuard` in nccl collectives by oraluben · Pull Request #130921 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Correctly set CUDAGuard in nccl collectives #130921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

Correctly set CUDAGuard in nccl collectives #130921

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3f0908c
Set guard for nccl collectives
oraluben Jul 17, 2024
b716b5c
lint
oraluben Jul 17, 2024
976c74d
Warn when barrier uses `guessDeviceForRank`
oraluben Jul 18, 2024
e7cc8cf
remove unused guard
oraluben Jul 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 8 additions & 9 deletions torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2591,7 +2591,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective(
work->stashed_for_allocator_safety_->push_back(input);
}

at::cuda::OptionalCUDAGuard gpuGuard;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it the case that gpuGuard without index is not effective? Or does it default to device 0?

at::cuda::OptionalCUDAGuard gpuGuard(device.index());

// Start event should only be recorded before the ncclGroupStart()
if (work->timingEnabled_) {
Expand Down Expand Up @@ -2752,8 +2752,6 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collectiveCoalesced(
std::make_shared<std::vector<at::Tensor>>(inputs);
}

at::cuda::OptionalCUDAGuard gpuGuard;

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this guard being removed instead of updated to device?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The AutoNcclGroup below does that

// Start event should only be recorded before the ncclGroupStart() (which
// happens inside AutoNcclGroup guard below)
if (work->timingEnabled_) {
Expand Down Expand Up @@ -3021,7 +3019,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::pointToPoint(
}

// is gpuGuard needed for the if block below, or can i swap them
at::cuda::OptionalCUDAGuard gpuGuard;
at::cuda::OptionalCUDAGuard gpuGuard(device.index());

if (!coalescing_state_) {
// Start event should only be recorded before the ncclGroupStart()
Expand Down Expand Up @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this change makes sense to me.

<< logPrefix()
<< c10::str(
" using GPU ",
"using GPU ",
deviceIdx,
" to perform barrier as devices used by this process are currently unknown. ",
"This can potentially cause a hang if this rank to GPU mapping is incorrect.",
"This can potentially cause a hang if this rank to GPU mapping is incorrect. ",
"Specify device_ids in barrier() to force use of a particular device.");
devices.emplace_back(guessDeviceForRank());
} else {
Expand All @@ -3929,8 +3927,9 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::barrier(const BarrierOptions& opts) {

// Use one device only
auto device = devices.back();
at::Tensor barrierTensor =
at::empty({1}, at::TensorOptions().device(device).dtype(at::kByte));
at::cuda::OptionalCUDAGuard gpuGuard(device.index());
at::Tensor barrierTensor = at::empty(
{1}, at::TensorOptions().device(at::DeviceType::CUDA).dtype(at::kByte));
// All reduce to achieve the barrier
auto work = allreduce_impl(barrierTensor);

Expand Down
Loading
0