当前位置：首页 > news >正文

NCCL AI 分布式训练集合通讯库技术基本原理

news 来源：原创 2025/9/8 1:12:31

AI 分布式训练

在一个最初的 AI 模型训练场景中，由于模型自身的程序体积、输入的参数量以及样本的数据量都比较有限，一张 GPU 的显存足以放下和处理这些数据，那么一张 GPU 既可以完成训练。但对于现如今的 AI 大模型而言，模型自身的体量就已经需要拆分放置到上万张 GPU 卡中，那么多张 GPU 甚至是多台 GPU 服务器进行分布式训练就是必然的选择。

如下图所示，区别对比了单机 GPU 训练和分布式 GPU 训练的区别。

单卡 GPU 训练
分布式 GPU 训练，DP 场景，基于 NCCL 集合通信库。

NCCL 的简介

NCCL（NVIDIA Collective Communication Library）集合通讯库，是 NVIDIA 为 AI 模型开发者提供的、专用于控制多张 GPU 之间进行高效通信的 lib 库，应用于分布式 GPU 训练场景。

如下深度学习软件堆栈图所示。NCCL 的北向是 AI 模型开发框架 PyTorch、Paddle、TensorFlow 等，这些 AI 框架通过集成、调用 NCCL lib 库来控制 GPU 之间的数据通信方式。NCCL 的南向是 CUDA 库，NCCL 通过 CUDA lib 来最终控制 GPU 设备的行为。

在这里插入图片描述

NCCL 的核心功能

集合通信 Verbs API：NCCL 北向提供了多种集合通信操作 API，供上层 AI 训练框架调用。
点到点通信 Verbs API：NCCL 除了支持集合通信之外，也支持 Send/Receive 点到点通信方式。
GPU/CUDA 通信控制：NCCL 南向调用 CUDA API 来控制 GPU 通信方式。
RNIC/RDMA 数据传输：NCCL 南向调用 RMDA API 来完成跨节点之间的 RNIC 数据传输，即 GPU Direct RDMA 技术。
GPU 通信拓扑结构感知：包括单节点的 GPU、NVLink、NVSwitch、PCIe Switch、CPU 等设备之间的拓扑感知，以及多节点的 RNIC 拓扑感知。
GPU 通信路径自动优化：根据拓扑结构自动地优化 GPU 间的通信路径，自动选择使用性能最高的通信环路。例如：NCCL 会优先使用单节点内的 NVLink 高速连接进行 GPU 间通信，而不是通过较慢的跨节点 RNIC 网络接口；（注：虽然 NCCL 支持 ring 和 tree 拓扑，但目前常用 ring 环。）

NCCL 的基本工作流程

首先，完成拓扑感知；
然后，完成路径选择；
接着，进行 GPU 建链；
最后，完成 all-reduce 等数据通信。

在这里插入图片描述

NCCL 的集合通信操作方式

NCCL 最核心的功能就是提供了多种分布式训练场景中需要使用到的集合通信操作，包括：

Broadcast（广播）：将一个 GPU 上的数据广播到所有其他 GPU。常用于在训练开始时，将输入的参数分发、同步到所有 GPU。
Scatter（分散）：将一个 GPU 上的数据分散到多个 GPU 上，每个 GPU 获得数据的一部分。
Gather（收集）：将多个 GPU 上的数据收集到一个 GPU 上，与 Scatter 操作相反。
All Gather（全局收集）：将所有 GPU 上的数据收集到所有 GPU 上，每个 GPU 最终都会拥有所有数据。
Reduce（归约）：将所有 GPU 上的数据进行归约操作，并将结果存储在一个 GPU 上。
All Reduce（全局归约）：将所有 GPU 上的数据进行全局归约操作，并将结果广播到所有 GPU。这是分布式训练中最常用的操作之一，用于同步模型参数的梯度。
Reduce Scatter（归约分散）：将所有 GPU 上的数据进行归约操作，并将结果分散到多个 GPU 上。

在这里插入图片描述

NCCL 的 API 编程示例

#include <nccl.h>
#include <cuda_runtime.h>
#include <iostream>#define CHECK(call) \
{ \const cudaError_t error = call; \if (error != cudaSuccess) { \std::cerr << "Error: " << __FILE__ << ":" << __LINE__ << ", " << cudaGetErrorString(error) << std::endl; \exit(1); \} \
}int main() {int size = 1024; // 数组大小int num_gpus = 2; // GPU 数量float *h_data[num_gpus]; // 主机端数据float *d_data[num_gpus]; // 设备端数据ncclComm_t comms[num_gpus]; // NCCL 通信器cudaStream_t streams[num_gpus]; // CUDA 流// 初始化 NCCLncclUniqueId id;ncclGetUniqueId(&id);ncclCommInitAll(comms, num_gpus, id);// 分配主机和设备内存for (int i = 0; i < num_gpus; ++i) {CHECK(cudaSetDevice(i));h_data[i] = (float*)malloc(size * sizeof(float));CHECK(cudaMalloc(&d_data[i], size * sizeof(float)));CHECK(cudaStreamCreate(&streams[i]));}// 初始化数据for (int i = 0; i < num_gpus; ++i) {for (int j = 0; j < size; ++j) {h_data[i][j] = i + 1; // 每个 GPU 的数据不同}CHECK(cudaMemcpy(d_data[i], h_data[i], size * sizeof(float), cudaMemcpyHostToDevice));}// 执行 AllReduce 操作for (int i = 0; i < num_gpus; ++i) {CHECK(cudaSetDevice(i));ncclAllReduce(d_data[i], d_data[i], size, ncclFloat, ncclSum, comms[i], streams[i]);}// 同步流for (int i = 0; i < num_gpus; ++i) {CHECK(cudaSetDevice(i));CHECK(cudaStreamSynchronize(streams[i]));}// 检查结果for (int i = 0; i < num_gpus; ++i) {CHECK(cudaMemcpy(h_data[i], d_data[i], size * sizeof(float), cudaMemcpyDeviceToHost));for (int j = 0; j < size; ++j) {if (h_data[i][j] != 3.0f) { // 1 + 2 = 3std::cerr << "Error: Incorrect result on GPU " << i << std::endl;exit(1);}}}std::cout << "AllReduce test passed!" << std::endl;// 释放资源for (int i = 0; i < num_gpus; ++i) {CHECK(cudaSetDevice(i));free(h_data[i]);CHECK(cudaFree(d_data[i]));CHECK(cudaStreamDestroy(streams[i]));ncclCommDestroy(comms[i]);}return 0;
}

初始化 NCCL：首先，我们使用 ncclGetUniqueId 生成一个唯一的 ID，然后使用 ncclCommInitAll 初始化 NCCL 通信器。
分配内存：为每个 GPU 分配主机和设备内存，并创建 CUDA 流。
初始化数据：为每个 GPU 初始化数据，并将数据从主机内存复制到设备内存。
执行 AllReduce 操作：在每个 GPU 上调用 ncclAllReduce 进行全局归约操作。这里我们使用 ncclSum 作为归约操作，表示对所有 GPU 上的数据进行求和。
同步流：确保所有 GPU 上的操作都已完成。
检查结果：将结果从设备内存复制回主机内存，并验证结果是否正确。
释放资源：释放所有分配的内存和资源。

目录

文章目录

AI 分布式训练

NCCL 的简介

NCCL 的核心功能

NCCL 的基本工作流程

NCCL 的集合通信操作方式

NCCL 的 API 编程示例

相关文章：