[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