TheRock/rocm-systems/projects/rccl/bindings/nccl4py

nccl4py: Python Bindings for RCCL

Python bindings for the ROCm Communication Collectives Library (RCCL), AMD's drop-in replacement for NVIDIA NCCL on ROCm. This package is a fork of the upstream NVIDIA nccl4py v0.2.0 and provides both low-level Cython bindings and a high-level Pythonic API for RCCL collective operations.

Highlights

• Drop-in API: keeps the upstream nccl.bindings / nccl.core import paths, so applications written against NCCL on CUDA work unchanged on RCCL on ROCm.
• HIP-backed cuda.core shim: ROCm hosts get from cuda.core import ... resolved via nccl._hip_compat.cuda_core_shim (backed by hip-python) without pulling NVIDIA cuda-bindings / cuda-core. The shim is registered in sys.modules on import nccl, so from cuda.core import ... resolves to the HIP backend after import nccl has run; it is not shipped as a top-level cuda package on disk to avoid shadowing co-installed distributions that contribute to the cuda.* namespace (cuda-bindings pulled transitively by PyTorch / RAPIDS / Triton / nsight tooling).
• RCCL-only collectives: wrappers for ncclAllReduceWithBias and ncclAllToAllv (which exist in librccl.so but have no upstream NCCL equivalent) are exposed via nccl.bindings.rocm_extensions.
• No CUDA toolchain at build time: Cython extensions treat cudaStream_t as an opaque void*; librccl.so is resolved at runtime via dlopen. pip install . works on a ROCm host without CUDA_HOME.

Requirements

• Runtime: librccl.so resolved at first use. Discovery order:
  1. NCCL_LIBRARY=<path> (explicit override; path or SONAME).
  2. ctypes.util.find_library("rccl") (the loader's ldconfig cache).
  3. $ROCM_PATH/lib/librccl.so, then $HIP_PATH/lib/librccl.so.
  4. Bare librccl.so, resolved via LD_LIBRARY_PATH / the standard loader search paths.
• Python: 3.10 or later.
• Python deps: hip-python, numpy, packaging (resolved automatically by pip).

Installation

pip install . is the supported install path on ROCm:

pip install .

Note: The Makefile and CMakeLists.txt shipped in this directory are inherited from upstream NVIDIA nccl4py and target NVIDIA hardware (they require CUDA_HOME / nvcc and pull CUDA-only torch / cupy wheels). They are out of scope for the ROCm bring-up and tracked separately under release engineering; use pip install on ROCm hosts.

For a development environment with the test extras:

pip install -e .
pip install pytest pytest-mpi pytest-cov mpi4py ml-dtypes

Tests

The unit-style files (tests/test_loader_stubs.py, tests/test_shim_surface.py, tests/test_rocm_extensions.py) run under a single rank without any GPU and are useful as smoke tests:

pytest tests/test_loader_stubs.py tests/test_shim_surface.py tests/test_rocm_extensions.py

The collective-correctness and PyTorch / CuPy interop tests are marked @pytest.mark.mpi and require mpirun -np N plus visible GPUs:

mpirun -np 8 pytest -m mpi tests/

tests/conftest.py skips MPI-tests gracefully on hosts without mpi4py, torch, or any visible GPU; the 8-rank tests xfail when fewer than 8 GPUs are visible.

Examples

examples/01_basic/ contains minimal, self-contained scripts:

mpirun -np 4 python examples/01_basic/01_allreduce.py
mpirun -np 2 python examples/01_basic/02_send_recv.py

Layout

• nccl/bindings/ - Cython bindings to the RCCL/NCCL C ABI; vendored from upstream nccl4py with the patches needed to drop the CUDA driver-types include and to resolve librccl.so via dlopen.
• nccl/_hip_compat/ - HIP/ROCm-only compatibility layer; not a public API. Hosts the cuda_core_shim package registered in sys.modules as cuda.core by nccl/__init__.py.
• nccl/core/ - the high-level Pythonic surface (Communicator, Buffer, Group, ...).

References

• RCCL
• RCCL documentation
• Upstream NVIDIA nccl4py