SIGN IN SIGN UP

[c10d] Bind local rank before eager NCCL2 initialization (#193237)

NCCL2 initializes eagerly, so selecting torch.cuda.current_device() before honoring LOCAL_RANK can bind forked workers to the same GPU. Prefer LOCAL_RANK and make DistributedTestBase bind the selected accelerator before init_process_group; the helper also accepts an explicit backend for compatibility tests.

Test Plan:

```bash
CUDA_VISIBLE_DEVICES=2,3 VIRTUAL_ENV=/home/tristanr/pytorch/.venv PATH=/home/tristanr/pytorch/.venv/bin:$PATH /home/tristanr/.local/bin/uv run --no-project test/distributed/test_c10d_ops_nccl.py ProcessGroupNCCLOpTest.test_all_gather_v
CUDA_VISIBLE_DEVICES=2,3 VIRTUAL_ENV=/home/tristanr/pytorch/.venv PATH=/home/tristanr/pytorch/.venv/bin:$PATH /home/tristanr/.local/bin/uv run --no-project test/distributed/test_functional_differentials.py TestFunctionalDifferentialsWithCompileCUDA.test_all_gather_into_tensor_autograd_compile_cuda
```

Authored with an AI assistant.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/193237
Approved by: https://github.com/kapilsh
T
Tristan Rice committed
05b558d4e79ea7a049dbc351a9067f0bf2a63faf
Parent: 26b9ddd
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/25/2026, 6:43:08 AM