背景

Data Parallelism

深度学习中我们按 batch 训练神经网络,前向传播获得误差,反向传播生成梯度,根据梯度修改参数,依次循环。

在多 GPU 中,我们将 batch 细分做到 DP,计算得到 Local Gradients,接下来使用 NCCL AllReduce 来共享局部梯度,求和得到总梯度再继续。

DP 的限制是不能强扩展,因为需要更大的单个 batch 来保证 GPU 效率。而如果 batch 过大会导致训练精度和效率下降。

Pipeline parallelism

因此我们将模型拆分,让每个 GPU 只包含部分层,每个 GPU 蕴含的层通过 NCCL 的 P2P 传递数据。实现 Pipeline parallelism

在反向传播时,由于每个 GPU 也只包含当前层的梯度,所以要结合 DP,用 NCCL AllReduce 实现每个都得到总梯度。

Tensor Parallelism

Allgather 非常快

Megatron

Quote

NCCL

Quote

  • NCCL 为 GPU 到 GPU 的通信提供原语。
  • 它通过 CUDA 流语义,在 GPU 上作为启动 CUDA 内核的库来实现,并与计算内核交错执行。
  • 初始化
    • CPU 端通信: 即使是 GPU 之间的通信,也需要 CPU 进行一些协调和设置。
    • 基于 Socket 的带外引导: “带外”通常指的是通过一个独立的通道进行通信,而不是通过数据传输的主通道。Socket(套接字)是一种网络通信的编程接口。这意味着 NCCL 可能使用网络套接字在不同的节点(例如,拥有多个GPU的服务器)之间建立初始连接和协调,以启动通信过程。
    • 容错: NCCL 具有处理通信失败的能力,以确保即使在出现问题时,系统也能继续运行或优雅地处理错误。
  • 系统拓扑检测
    • 包含 GPU、CPU、网卡(NICs)、PCI 交换机、NVLinks 等的图
    • 用于集体算法的图搜索
  • 点对点通信
    • 发送/接收语义: 就像邮寄包裹一样,一个 GPU 发送数据,另一个 GPU 接收数据。这是通信的基本操作。
    • 不同的协议: 根据硬件连接和性能需求,NCCL 可能会使用不同的底层通信协议来优化点对点通信的效率。
  • 集体通信算法
    • 环(Ring)、树(Tree)、SHARP

设备要求

  • NVIDIA GPU 计算能力 >= 3.5
  • glibc >= 2.17, CUDA >= 10.0
  • 推荐使用支持 NVLink 的 GPU 以获得最佳性能
  • 对于多机通信,建议使用 InfiniBand 或 RoCE 网络

NCCL in Pytorch

PyTorch 内置 NCCL 后端支持,使用非常简单:

import torch.distributed as dist
 
# 初始化进程组
dist.init_process_group(backend='nccl')
 
# 创建分布式模型
model = torch.nn.parallel.DistributedDataParallel(model)

通信协议

Simple

  • 操作
    • 移动尾指针代表发送就绪
    • 移动头指针代表接收完成
  • 特性
    • 高延迟:~6us,由于内存(同步指针)屏障
    • 满带宽
    • 流水线粒度:128~512,Chunk size 在 512K 时最佳

LL (Low Latency)

  • 分线程传输,线程越多颗粒度越细
  • 数据与 flag 一同传输,标识数据何时可用
    • Sender 在所有线程传输完成后可以开始新一轮传输
    • Receiver 在所有线程接收完毕后可以开始新一轮接收
    • 此时可以移动 head 指针,此条消息处理完成
  • 特性
    • 低延迟:~1us
    • 带宽:50% peak 50% overhead
    • 流水线粒度:16B
    • 需要系统的 8-Byte 原子性,保证 flag 不会在数据之前写入

LL128

  • Wrap:一组 CUDA
  • 特性
    • 低延迟:~2us
    • 带宽:95% peak 5% overhead
    • 流水线粒度:128B
    • 需要系统的 128-Byte 原子性,保证 flag 不会在数据之前写入(只有 NVLink)

集体通信和算法

算法

以 AllReduce 为例:

Ring

  • 算法
    • 切片
    • 分配每个 GPU 不同色块
    • 轮换并 Reduce 相同色块 x n-1
    • 继续轮换进行 AllGather 并分到各个 GPU x n-1
  • 特性
    • 简单通用,适合小规模
    • 延迟随 GPU 数量显著增加

Multi-Ring on DGX

  • 分散到每个 SM

Tree

  • 双二叉树保证出入流量平衡
  • 算法
  • 特性
    • 对数延迟
    • 小、中数据表现好
    • 难以达到理论带宽(规约计算不平衡)

Multi-Tree on DGX

SHARP/Collnet

You may wonder: IBSharp

Chain

  • 算法
    • 不使用树状规约
  • 特性
    • 带宽高:网络内部流量减少
    • 良好扩展性

Multiple inter-node chains on DGX

Direct

  • 算法
    • 先 AlltoAll 广播
    • In-network 规约
    • 再 AlltoAll收集
  • 特性
    • 需要 NWSwitch 或者全连接
    • 在每个 GPU 有自己的 NIC 表现最佳

API