【NCCL】入门（一）：介绍和简单demo

介绍

NCCL（NVIDIA Collective Communications Library）是由NVIDIA开发的一种用于高性能GPU集群的通信库。它旨在提供高效的GPU间通信和协作能力，以加速分布式深度学习和其他GPU密集型计算任务。

NCCL支持在多个GPU之间进行并行计算和通信。它可以在多个GPU之间实现高效的数据传输和同步，以利用集群中的所有GPU资源。被广泛用于分布式深度学习训练中，特别是在使用多个GPU进行模型训练时。它提供了一致的接口和通信原语，使不同GPU之间的数据交换和同步变得简单和高效。NCCL支持跨多个节点的通信，使得在GPU集群中进行分布式计算和通信变得容易。它提供了跨节点的通信原语，如跨节点的点对点通信、全局同步和归约操作。

安装NCCL

1. 在https://developer.nvidia.com/nccl，根据操作系统和CUDA版本选择适当的版本进行下载。

2. 使用root权限运行安装脚本

3. 安装完成后，设置相应的环境变量，以便系统能够正确找到NCCL库。在.bashrc（或.bash_profile）文件中添加以下行：

    ```
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/nccl/lib
    export C_INCLUDE_PATH=$C_INCLUDE_PATH:/usr/local/nccl/include
    ```
    
    其中`/usr/local/nccl`替换为你实际的NCCL库的安装路径。
   

4. 使用ldconfig -p | grep libnccl检查

代码

流程

import torch
import torch.distributed as dist
import os

# 初始化进程组
dist.init_process_group(backend='nccl')

# 获取当前进程相关信息
rank = dist.get_rank()
world_size = dist.get_world_size()

# 创建输入数据（示例：每个进程创建一个大小为2的张量）
input_data = torch.tensor([rank, rank + 1]).cuda()

# 创建输出数据（示例：每个进程创建一个大小为2的张量）
output_data = torch.tensor([0, 0]).cuda()

# 同步所有进程
dist.barrier()

# 执行全局同步操作
dist.all_reduce(input_data, output_data, op=dist.ReduceOp.SUM)

# 输出结果
print(f"Rank {rank}: Input data: {input_data}")
print(f"Rank {rank}: Output data: {output_data}")

# 释放进程组资源
dist.destroy_process_group()

初始化函数dist.init_process_group

dist.init_process_group是PyTorch的分布式训练工具包（torch.distributed）中的一个函数，用于初始化分布式训练的进程组。它是进行多机多卡训练时必须的一步。

init_process_group(backend, init_method='env://', rank=-1, world_size=-1, group_name='')

参数说明：
backend：指定分布式训练的后端，常用的后端包括gloo、nccl等。gloo是一个基于TCP的高性能后端，适用于CPU和GPU的训练；nccl是NVIDIA提供的高性能后端，适用于GPU训练。
init_method：指定进程组的初始化方法，可以是文件路径、URL或环境变量。默认值是’env://'，表示使用环境变量来初始化进程组。
rank：当前进程的排名（从0开始），默认值是-1，表示自动获取当前进程的排名。
world_size：进程组中的进程总数，默认值是-1，表示自动获取进程组的总数。
group_name：进程组的名称，默认为空字符串。
在分布式训练中，所有参与训练的进程需要调用init_process_group来初始化进程组，确保它们能够互相通信和同步。通常，这个函数会在每个训练脚本的开始部分被调用。