trunk/c3c33fd3481ab28ef0037bbd9b72588fbdfec691: [CUDA][NCCL] Fix nccl.broadcast dropping the root argument (#187216)

Fixes #179908

Summary

torch.cuda.nccl.broadcast(tensors, root=N) silently ignores root and always broadcasts from tensors[0], while nccl.reduce honors root — full root-cause chain in my issue comment:

1. torch/cuda/nccl.py::broadcast passes root to torch._C._nccl_broadcast correctly.


2. THCPModule_nccl_broadcast (torch/csrc/cuda/python_nccl.cpp) parses and validates root, then calls torch::cuda::nccl::broadcast(inputs, streams, user_comms) without it.


3. torch::cuda::nccl::broadcast (torch/csrc/cuda/nccl.cpp) has no root parameter and hardcodes 0 in the ncclBcast call.

The bug is old: v0.3.0 passed root through to ncclBcast; commit de5f7b7251f ("Base for pure C++ NCCL interface", Dec 2017) dropped it when the C++ helper was extracted, so the binding has been wrong since v0.4.0 (~8.5 years).

Fix

• torch/csrc/cuda/nccl.h / nccl.cpp: add int32_t root = 0 to torch::cuda::nccl::broadcast, mirroring the existing reduce signature (defaulted, so the internal caller in torch/csrc/cuda/comm.cpp::broadcast_coalesced is unaffected); validate it in-range exactly like reduce does; pass it to ncclBcast.


• torch/csrc/cuda/python_nccl.cpp: forward the already-parsed root.


• test/distributed/test_nccl.py: extend test_broadcast with a non-zero-root regression block (root = nGPUs - 1; the root device holds a uniform_() tensor, all other devices hold zeros; assert every device ends up with the root's values). Before this fix the broadcast originates from the zeros at index 0, so the new assertions fail; after the fix they pass. It runs under the existing TEST_MULTIGPU skip and @dtypes(*broadcast_dtypes) (incl. float8) conventions.

Test plan

• New regression assertions in test/distributed/test_nccl.py::TestNCCL::test_broadcast (multi-GPU).


• lintrunner -a on the four touched files: no lint issues (CLANGTIDY skipped locally — requires a build dir).


• Compile-verified the touched translation units (nccl.cpp, python_nccl.cpp) plus the unmodified internal caller (comm.cpp) with g++ -fsyntax-only against the modified header: all clean.


• Validation gap, stated honestly: my box has a single GPU, so I could not run the multi-GPU broadcast path locally; python test/distributed/test_nccl.py -k broadcast collects and skip-passes here (OK (skipped=2)). The behavioral coverage relies on PyTorch's multi-GPU CI exercising the extended test_broadcast.

🤖 Generated with Claude Code

Pull Request resolved: #187216

Approved by: https://github.com/d4l3k