526互联

NVIDIA Collective Communications Library (NCCL)

发布时间 2023-06-02 12:04:59作者: stardsd

一、简介

NVIDIA Collective Communications Library (NCCL) 是一个多 GPU 和多节点通信原语库，具有拓扑感知能力，可以轻松集成到应用程序中。 集体通信算法采用许多协同工作的处理器来聚合数据。 NCCL 不是成熟的并行编程框架； 相反，它是一个专注于加速集体通信原语的库。

NCCL 提供诸如 all-gather、all-reduce、broadcast、reduce、reduce-scatter 以及点对点发送和接收等例程，这些例程经过优化以在 PCIe 和 NVLink 高速互连上实现高带宽和低延迟 一个节点和跨节点 1 的 NVIDIA Mellanox 网络。 NCCL 使用简单的 C API，可以通过多种编程语言轻松访问。 NCCL 紧跟 MPI（消息传递接口，message passing interface） 定义的流行集合 API。

Caffe2、Chainer、MxNet、PyTorch 和 TensorFlow 等领先的深度学习框架都集成了 NCCL，以加速多 GPU 多节点系统上的深度学习训练。 NCCL 可以作为 NVIDIA HPC SDK 的一部分下载，也可以作为 Ubuntu 和 Red Hat1 的单独软件包下载。

二、安装

要安装 NCCL 库，您需要按照以下步骤 1：

  1. 前往：NVIDIA NCCL 主页2。 https://docs.nvidia.com/deeplearning/nccl/install-guide/index.html 
  2. 单击下载。
  3. 完成简短调查并单击提交。
  4. 接受条款和条件。 显示可用的 NCCL 下载版本列表。
  5. 选择要安装的 NCCL 版本。 显示可用资源列表。 请参阅以下部分，根据您使用的 Linux 发行版选择正确的软件包。
  6. 使用适合您的 Linux 发行版的命令将 NCCL 库安装到您的系统上。 有关详细信息，请参阅 NCCL 安装指南1 中的安装 NCCL。
  7. 修改您的应用程序以链接到该库。 在您的应用程序中包含头文件 nccl.h。 创建一个通讯器。 有关详细信息，请参阅 NCCL 开发人员指南中的创建通信器

三、Pytorch实现

  1. 在您的系统上安装 NCCL 库。 
  2. 在您的系统上安装 PyTorch。 您可以将 NVIDIA NGC 容器用于已包含 NCCL2 的 PyTorch，或者您可以从源代码构建 PyTorch 并使用本地安装的 NCCL3。https://pytorch.org/docs/stable/distributed.html
  3. 使用 NCCL 后端初始化分布式进程组。 您可以使用带有 backend=“nccl” argument4 的 init_process_group() 函数。 您还需要根据分布式设置指定 init_method、rank 和 world_size 参数。
  4. 使用带有 NCCL 后端的 torch.distributed 包中的集体通信功能。 例如，您可以使用 dist.all_reduce() 在所有进程中执行 all-reduce 操作4。 您还可以使用 DistributedDataParallel 模块来包装您的模型，并使用 NCCL 后端将其并行化到多个 GPU。

四、示例代码

如何使用 NCCL 后端跨两个进程执行 all-reduce 操作。 可使用如下代码，以在单台机器或具有不同 IP 地址的多台机器上运行此示例。

# run.py
#!/usr/bin/env python
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def run(rank, size):
    """ Distributed function to be implemented later. """
    # Create a tensor and initialize it to zero
    tensor = torch.zeros(1)
    # Add the rank of the process to the tensor
    tensor += rank
    # Print the initial value of the tensor
    print('Rank ', rank, ' has data ', tensor[0])
    # Perform an all-reduce operation using NCCL backend
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    # Print the final value of the tensor
    print('Rank ', rank, ' has data ', tensor[0])

def init_process(rank, size, fn, backend='nccl'):
    """ Initialize the distributed environment. """
    # Set the master address and port
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    # Initialize the process group with NCCL backend
    dist.init_process_group(backend, rank=rank, world_size=size)
    # Run the distributed function
    fn(rank, size)

if __name__ == "__main__":
    # Set the number of processes (world size)
    size = 2
    # Create a list of processes
    processes = []
    # Use spawn method to start multiple processes
    mp.set_start_method("spawn")
    # Loop over the number of processes
    for rank in range(size):
        # Create a new process with given rank, size and function
        p = mp.Process(target=init_process, args=(rank, size, run))
        # Start the process
        p.start()
        # Append the process to the list
        processes.append(p)
    # Wait for all processes to finish
    for p in processes:
        p.join()

　　

单机运行：

python run.py

　　

多机运行：

# On machine 192.168.0.1 (rank 0)
export MASTER_ADDR=192.168.0.1
python run.py

# On machine 192.168.0.2 (rank 1)
export MASTER_ADDR=192.168.0.2
python run.py