NVIDIA Collective Communication Library (NCCL) Runtime for CUDA 12

Warnings

• breaking NCCL versions are tightly coupled with CUDA Toolkit versions and the CUDA version used to compile deep learning frameworks (like PyTorch or TensorFlow). Mismatches can lead to runtime errors, silent performance degradation, or unexpected behavior.
  Fix: Ensure that the `nvidia-nccl-cu12` package, your system's CUDA Toolkit, and the CUDA version used by your deep learning framework are all compatible. Consult the NVIDIA documentation or framework-specific guides for compatibility matrices. For PyTorch, `torch.cuda.is_available()` and `torch.version.cuda` can help verify. For `nccl4py`, use `pip install "nccl4py[cu12]"` to ensure correct CUDA 12 support.
• gotcha The `nvidia-nccl-cu12` package itself primarily provides the `libnccl.so` shared library. Direct Python API calls are not exposed through this package. Instead, Python users interact with NCCL through higher-level libraries like `nccl4py` (official bindings) or as a backend to distributed training modules in frameworks like PyTorch (`torch.distributed`) or TensorFlow (`tf.distribute`).
  Fix: To use NCCL directly from Python, install and import `nccl4py`. If using with a deep learning framework, configure its distributed module to use the NCCL backend. Avoid `import nccl` for direct API calls, as this package is a runtime provider.
• gotcha Conflicts can arise if multiple NCCL installations are present on the system (e.g., `nvidia-nccl-cu12` from PyPI, a system-wide `apt`/`dnf` installed NCCL, or one bundled with a deep learning framework). The linker's search path (`LD_LIBRARY_PATH`) can affect which `libnccl.so` is loaded, potentially leading to incorrect versions being used.
  Fix: Prefer using `nvidia-nccl-cu12` installed via pip for consistency within Python environments. If system-wide NCCL is necessary, carefully manage `LD_LIBRARY_PATH` to ensure the correct `libnccl.so` is prioritized. Frameworks like PyTorch often statically link NCCL, mitigating some of these issues, but custom builds might need `USE_SYSTEM_NCCL` flags.

Install

• pip install nvidia-nccl-cu12 Default Install
• pip install "nccl4py[cu12]" # Official Python bindings With NCCL4Py Bindings

Imports

• NcclCommunicator

  from nccl.core import NcclCommunicator

  Primary high-level API for official nccl4py bindings.
• lib

  from nccl.bindings import lib

  Low-level API for official nccl4py bindings, less common for direct user interaction.
• dist

  import torch.distributed as dist

  nvidia-nccl-cu12 is the runtime. Python users typically interact via higher-level frameworks like PyTorch's distributed module, or dedicated bindings like nccl4py.
• tf.distribute.NcclAllReduce

  import tensorflow as tf
  strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.NcclAllReduce())

  nvidia-nccl-cu12 is the runtime. TensorFlow utilizes NCCL as a backend for its distribution strategies.

Quickstart

This quickstart demonstrates how NCCL is typically used indirectly via PyTorch's `torch.distributed` module for multi-GPU collective communication, specifically an `all_reduce` operation. NCCL provides the underlying high-performance backend. A proper distributed launcher (e.g., `torch.distributed.launch` or `mpirun`) is required to run this code across multiple processes/GPUs. For direct Python bindings, consider `nccl4py` for explicit NCCL API calls.

import os
import torch
import torch.distributed as dist

# This quickstart assumes a multi-process setup, typically launched
# via torch.distributed.launch or mpirun, where each process
# runs this script with a unique rank and world_size.

# Example environment variables (set by launch utility):
# os.environ['MASTER_ADDR'] = os.environ.get('MASTER_ADDR', 'localhost')
# os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '29500')
# os.environ['RANK'] = os.environ.get('RANK', '0')
# os.environ['WORLD_SIZE'] = os.environ.get('WORLD_SIZE', '1')

def run_distributed_example(rank, world_size):
    # Initialize the process group with NCCL backend
    print(f"Initializing process group for rank {rank}/{world_size-1}...")
    dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
    print(f"Process group initialized on rank {rank}.")

    # Set device for the current process
    torch.cuda.set_device(rank)

    # Create a tensor on the GPU
    tensor = torch.ones(10, device=f'cuda:{rank}') * (rank + 1)
    print(f"Rank {rank}: Initial tensor value: {tensor}")

    # Perform an all_reduce operation (summing tensors across all GPUs)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

    print(f"Rank {rank}: Tensor after all_reduce: {tensor}")

    # Clean up the process group
    dist.destroy_process_group()
    print(f"Rank {rank}: Process group destroyed.")

# To run this, you would typically use:
# python -m torch.distributed.launch --nproc_per_node=2 your_script.py
# Or set environment variables and run:
# MASTER_ADDR=localhost MASTER_PORT=29500 RANK=0 WORLD_SIZE=2 python your_script.py
# MASTER_ADDR=localhost MASTER_PORT=29500 RANK=1 WORLD_SIZE=2 python your_script.py

# For simplicity, if running as a single process for structural check:
if __name__ == '__main__':
    # In a real scenario, rank and world_size would be provided by a launcher.
    # This block is for structural demonstration only and will not perform
    # actual distributed communication without a proper launcher.
    try:
        rank = int(os.environ.get('RANK', '0'))
        world_size = int(os.environ.get('WORLD_SIZE', '1'))
        if torch.cuda.is_available() and world_size > 0:
             run_distributed_example(rank, world_size)
        else:
             print("CUDA not available or world_size is 0. Cannot run distributed example.")
    except RuntimeError as e:
        print(f"Error initializing distributed environment: {e}. This often happens if not run with a proper distributed launcher like torch.distributed.launch.")