Developer Tools

PyTorch #187385: Dynamic NCCL EP linking for system NCCL wheel builds

⚡System NCCL users on wheels now get dynamically linked NCCL EP support.

Deep Dive

This commit (830e4a7) in the PyTorch repository removes the gate preventing USE_NCCL_EP from being used with USE_SYSTEM_NCCL=ON, specifically for wheel (system NCCL) builds. Previously, the NCCL Execution Provider (EP) was only built as a static library (libnccl_ep.a) when using system NCCL, which could cause conflicts. Now, the build system branches: when USE_SYSTEM_NCCL=OFF (source/dev builds), it produces a static libnccl_ep.a; when ON (wheel/CI), it produces a dynamic libnccl_ep.so that NEEDED-links against the system NCCL library (libnccl.so.2) resolved from the nccl4py wheel at runtime.

The implementation modifies `cmake/External/nccl_ep.cmake` to branch shared vs static libraries, bakes `NCCL_EP_JIT_HOME` only on the static path, and adjusts `torch/CMakeLists.txt` to NEEDED-link libnccl_ep and use the nccl4py RPATH. The `_token_switch.py` script gains a `_prepare_nccl4py` fallback that locates the nccl4py wheel and sets environment variables (`NCCL_EP_HOME`, `NCCL_HOME`) to load the extension. Tested on 4xH100 GPUs using the nvidia.nccl wheel, the suite passed all 7 tests, demonstrating that the dynamic linking works correctly and resolves dependencies at runtime.

Key Points
  • Enables USE_NCCL_EP with USE_SYSTEM_NCCL=ON for wheel builds, dropping the previous NOT gate.
  • Dynamic libnccl_ep.so is built for system NCCL, static .a for source/dev; leverages nccl4py for JIT headers.
  • Tested on 4xH100 with nvidia.nccl wheel; all 7 tests passed with correct NEEDED-link resolution.

Why It Matters

Improves NCCL EP compatibility for PyTorch wheel users, reducing static linking friction in multi-GPU setups.

📬 Get the top 10 AI stories daily