Skip to content
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

Jump to bottom

Distribute GPUs in round robin mode for distributed_test #46389

Closed
wants to merge 1 commit into from
Closed

Distribute GPUs in round robin mode for distributed_test #46389

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 2 additions & 4 deletions torch/testing/_internal/distributed/distributed_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,16 +362,14 @@ def _init_multigpu_helper(self):
"""
nGPUs = torch.cuda.device_count()
world_size = dist.get_world_size()
visible_devices = range(nGPUs)

if BACKEND == "nccl":
apply_hack_for_nccl()

nGPUs_per_process = nGPUs // world_size
rank_to_GPU = {
i: list(
visible_devices[i * nGPUs_per_process: (i + 1) * nGPUs_per_process]
)
# Each rank has to get the GPU with the index equal to its rank
i: [i + gpu_num * world_size for gpu_num in range(nGPUs_per_process)]
Copy link
Contributor

Choose a reason for hiding this comment

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

Hey @Flamefire

Would be correct if I assume world_size=2 and 4 GPUs in total, then we have the following? Before this change, we have:

rank 0 -> gpu [0, 1]
rank 1 -> gpu [2, 3]

After this change, we have:

rank 0 -> gpu [0, 2]
rank 1 -> gpu [1, 3]

Copy link
Contributor

Choose a reason for hiding this comment

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

Could you please help me understand why the above change fixes the problem described below? Thx.

Due to NCCL communicator reuse this then leads to rank 0 using the (kinda) temporary communicator while the other processes might use other GPUs leading to them trying to create a new communicator and waiting for rank 0 until that creates a new (potentially unrelated) one.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes your assumption is correct. For an in-depth analysis see the issue where I posted many details, here only the summary:

During the barrier that happens very early (on creation of the process group) each process creates a communicator with GPU idx equal to its rank:

pytorch/torch/lib/c10d/ProcessGroupNCCL.cpp

Line 1389 in 2e2fe8c

int16_t deviceIdx = static_cast<int16_t>(rank_ % numGPUs);

The problem with the old distribution is, that rank 1 (in your example) wants to use GPU 2 afterwards and hence needs a new communicator. But rank 0 wants to (continue to) use GPU 0 and hence does not need a new communicator. But because creation of a communicator is a collective operation this will fail due to rank 1 waiting for rank 0 that never appears.

Later Rank 0 might want to create a communicator for GPU 0 and 1 and joins the still waiting rank 1 in creating one, but now there is a mismatch: Rank 0 is already further down and expects 4 total GPU ranks (2 per process) while rank 1 only expects 2. This leads to a (correct) system error in NCCL code, but the real problem is already earlier.

for i in range(world_size)
}
return rank_to_GPU
Expand Down