LLMs

主题切换

图解分布式训练（三）—— nn.parallel.DistributedDataParallel ​

来源：AiGC面试宝典 作者：宁静致远 日期：2023年09月29日

为什么需要 nn.parallel.DistributedDataParallel？ ​

多GPU并行训练的原理是将模型参数和数据分布到多个GPU上，同时利用多个GPU的计算能力加速训练过程。具体实现需要考虑以下三个核心问题：

1. 数据如何划分？ ​

由于模型需要处理的数据量通常很大，将所有数据放入单个GPU内存中可能会导致内存不足，因此需要将数据划分到多个GPU上。一般有两种划分方式：

划分方式说明
数据并行（Data Parallelism）将数据分割成多个小批次，每个GPU处理其中一个小批次，然后将梯度汇总后更新模型参数
模型并行（Model Parallelism）将模型分解成多个部分，每个GPU处理其中一个部分，并将处理结果传递给其他GPU以获得最终结果

📝 通俗解释：数据并行就像把一本很厚的书分给多个人同时阅读，每个人读一部分，最后交流心得；模型并行就像把书的制作过程分成多个环节，不同人负责不同环节（排版、印刷、装订等）。

2. 计算如何协同？ ​

由于每个GPU都需要计算模型参数的梯度并将其发送给其他GPU，因此需要使用同步机制来保证计算正确性。一般有两种同步方式：

同步方式说明
梯度同步（同步数据并行）在每个GPU上计算模型参数的梯度，然后将梯度发送到其他GPU上进行汇总，最终更新模型参数
参数同步（同步模型并行）在每个GPU上计算模型参数的梯度，然后将模型参数广播到其他GPU上进行汇总，最终更新模型参数

3. DP 的局限性 ​

DP（DataParallel）只支持单机多卡场景，在多机多卡场景下，DP的通讯问题将被放大。

  • DDP首先要解决的就是通讯问题：将主节点（Rank 0）的通讯压力均衡转移到各个工作节点（Worker）上。实现这一点后，可以进一步去中心化，留下的Worker可以直接参与计算。

📝 通俗解释：DP就像一个班长（Rank 0）要收集所有同学（其他GPU）的作业，然后自己改完再发回给同学。当同学多了，班长就忙不过来了。DDP的做法是让每个同学自己改自己的作业，然后互相交换答案，这样大家都不累。

一、DistributedDataParallel 核心 —— Ring-AllReduce ​

上节讲到DP只支持单机多卡场景，主要原因是DP无法解决数据并行中通讯负载不均的问题，而DDP能够解决该问题的核心在于Ring-AllReduce。它由百度最先提出，非常有效地解决了数据并行中通讯负载不均的问题，使得DDP得以实现。

1.1 Ring-AllReduce 介绍 ​

假设有4块GPU，每块GPU上的数据对应被切成4份。AllReduce的最终目标是让每块GPU上的数据都变成右侧汇总的样子。

[图表描述]：

  • 左侧展示了四个GPU的数据块：
    • GPU 0: a0, b0, c0, d0
    • GPU 1: a1, b1, c1, d1
    • GPU 2: a2, b2, c2, d2
    • GPU 3: a3, b3, c3, d3
  • 箭头指向右侧汇总列，表示AllReduce后的结果，每块GPU都应拥有完整的累加和：
    • $a0 + a1 + a2 + a3$
    • $b0 + b1 + b2 + b3$
    • $c0 + c1 + c2 + c3$
    • $d0 + d1 + d2 + d3$

Ring-AllReduce分两大步骤实现该目标：Reduce-Scatter（分散求和）和All-Gather（全局收集）。

📝 通俗解释：Ring-AllReduce就像一个传递游戏。4个人围成一圈，每个人手里有一张写着数字的卡片。每个人把自己的数字传给右边的人，同时接收左边人传来的数字并累加。转几圈后，每个人手里的数字都是所有人的数字之和。

1.2 Reduce-Scatter ​

Reduce-Scatter的定义网络拓扑关系，使得每个GPU只和其相邻的两块GPU通讯。每次发送对应位置的数据进行累加。由于每一次累加更新都形成一个拓扑环，因此被称为"Ring"（环）。

如果你对该问题有疑惑，下文将用图例把详细步骤介绍清楚。

[图表描述 - 拓扑结构]： 展示了四个数据块列（a0-d0, a1-d1, a2-d2, a3-d3），它们之间通过箭头首尾相连形成一个环：

  • 第一列 → 第二列 → 第三列 → 第四列 → 第一列（底部回环）

[图表描述 - 数据累加过程]： 展示了Reduce-Scatter的第一步累加结果：

  • 第一列底部：$d0 + d3$（蓝色高亮）
  • 第二列顶部：$a0 + a1$（蓝色高亮）
  • 第三列第二行：$b1 + b2$（蓝色高亮）
  • 第四列第三行：$c2 + c3$（蓝色高亮）

一次累加完毕后，蓝色位置的数据块被更新，被更新的数据块将成为下一次更新的起点，继续做累加操作。

Reduce-Scatter 阶段结束状态：

GPU 1GPU 2GPU 3GPU 4
$a_0$$a_0 + a_1$$a_0+a_1+a_2$$a_0+a_1+a_2+a_3$
$b_0+b_1+b_2+b_3$$b_1$$b_1+b_2$$b_1+b_2+b_3$
$c_0+c_2+c_3$$c_0+c_1+c_2+c_3$$c_2$$c_2+c_3$
$d_0 + d_3$$d_1+d_0+d_3$$d_0+d_1+d_2+d_3$$d_3$

经过3次更新后，每块GPU上都有一块数据拥有了对应位置完整的聚合（图中高亮部分）。此时，Reduce-Scatter阶段结束，进入All-Gather阶段。目标是把这部分完整数据广播到其余GPU对应的位置上。

1.3 All-Gather ​

如名字里"Gather"所述的一样，这个操作依然按照"相邻GPU对应位置进行通讯"的原则，但对应位置数据不再做相加，而是直接替换。All-Gather以上一阶段得到的完整数据块作为起点。

[图表描述]： 图示展示了All-Gather的过程。第一行是起始状态，第二行展示了数据广播替换的过程，蓝色框表示新替换进来的完整数据块。

All-Gather 过程示意：

GPU 1GPU 2GPU 3GPU 4
$a_0+a_1+a_2+a_3$$a_0 + a_1$$a_0+a_1+a_2$$a_0+a_1+a_2+a_3$
$b_0+b_1+b_2+b_3$$b_0+b_1+b_2+b_3$$b_1+b_2$$b_1+b_2+b_3$
$c_0+c_2+c_3$$c_0+c_1+c_2+c_3$$c_0+c_1+c_2+c_3$$c_2+c_3$
$d_0 + d_3$$d_1+d_0+d_3$$d_0+d_1+d_2+d_3$$d_0+d_1+d_2+d_3$

以此类推，同样经过3轮迭代后，使得每块GPU上都汇总到了完整的数据，变成如下形式：

最终状态（所有GPU）：

数据块完整累加和
a$a_0 + a_1 + a_2 + a_3$
b$b_0 + b_1 + b_2 + b_3$
c$c_0 + c_1 + c_2 + c_3$
d$d_0 + d_1 + d_2 + d_3$

📝 通俗解释：Reduce-Scatter是"收集并分散"——把大家的数据累加起来，然后分散到每个人手里一小部分。All-Gather是"收集并复制"——把每个人手里的完整数据复制给所有人。两个阶段配合，就能让所有GPU都拥有所有数据的完整累加结果。

二、nn.parallel.DistributedDataParallel 函数介绍 ​

函数定义 ​

python
CLASS torch.nn.parallel.DistributedDataParallel(
    module, device_ids=None, output_device=None,
    dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25,
    find_unused_parameters=False, check_reduction=False
)

函数参数 ​

参数说明
module要放到多卡训练的模型（必填）
device_ids数据类型是列表，表示可用的GPU卡号
output_device数据类型也是列表，表示模型输出结果存放的卡号。如果不指定，默认放在0卡，这也是为什么多GPU训练并不是负载均衡的，一般0卡会占用更多。这里还涉及一个小知识点：如果程序开始加os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"，那么0卡（逻辑卡号）指的是2卡（物理卡号）
dim指按哪个维度进行数据划分，默认是输入数据的第一个维度，即按batchsize划分（假设数据格式是 B, C, H, W）

📝 通俗解释：DDP就像一个"分布式训练包装器"，你把模型扔进去，它会自动帮你把模型复制到多个GPU上，协调数据分配、梯度同步等工作。你只需要像正常使用单GPU模型一样写代码就行。

三、nn.parallel.DistributedDataParallel 如何多卡加速训练？ ​

DDP在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程将其 broadcast 到所有进程后，各进程用该梯度来独立更新参数。

  • DDP：梯度计算 → All-Reduce汇总平均 → 各进程独立更新参数
  • DP：梯度计算 → 汇总到GPU0 → 反向传播更新参数 → 广播参数给其他GPU

由于DDP各进程中的模型初始参数一致（初始时刻进行一次broadcast），而每次用于更新参数的梯度也一致，因此各进程的模型参数始终保持一致。

而在DP中，全程维护一个optimizer，对各个GPU上梯度进行求和，而在主卡进行参数更新，之后再将模型参数broadcast到其他GPU。

📝 通俗解释：DP像是一个"中央集权"模式——所有GPU把作业交给班长（GPU 0），班长改完再发回给每个人。DDP像是"民主自治"模式——每个人自己改自己的作业（梯度），然后大家投票（All-Reduce）确认答案一致后，各自更新自己的作业。很明显，DDP的效率更高，因为大家可以同时工作，不用等班长。

四、nn.parallel.DistributedDataParallel 实现流程 ​

步骤1：初始化进程组 ​

python
# 初始化使用nccl后端，官网还提供了 'gloo' 作为backend
# 知道nccl是用来GPU之间进行通信的就可以
torch.distributed.init_process_group(backend="nccl")

如果需要进行小组内集体通信，用new_group创建子分组。

📝 通俗解释：这一步就像建立一个"通讯录"，让所有GPU知道彼此的存在和联系方式。

步骤2：使用 DistributedSampler 划分数据 ​

python
# train_dataset就是自己的数据集
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
dataloader = DataLoader(
    dataset, 
    batch_size=batch_size, 
    shuffle=(train_sampler is None),
    sampler=train_sampler, 
    pin_memory=False
)

注：pin_memory（固定内存）设为True会加速训练（但官网默认参数设置是False）。

📝 通俗解释：DistributedSampler就像一个"分书员"，确保每个GPU读到的数据都不一样，避免大家读重复了浪费资源。pin_memory=True就像给数据腾出一块"VIP座位"，访问起来更快。

步骤3：创建DDP模型 ​

python
# DDP是all-reduce的，即汇总不同 GPU 计算所得的梯度，并同步计算结果
# all-reduce 后不同 GPU 中模型的梯度均为 all-reduce 之前各GPU梯度的均值
model = torch.nn.parallel.DistributedDataParallel(
    model, 
    device_ids=[args.local_rank],
    output_device=args.local_rank, 
    find_unused_parameters=True
)

📝 通俗解释：这一步把普通模型"升级"成DDP模型，就像给模型装上了"多线程通信"能力，让它能和其他GPU上的"自己"协同工作。

步骤4：命令行启动训练 ​

bash
$ python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 --node_rank=0 --master_port=6005 train.py

参数说明：

  • --nnodes：节点数量（单机多卡设为1）
  • --nproc_per_node：当前主机创建的进程数（等于可用GPU数量）
  • --node_rank：节点序号
  • --master_port：主节点端口号

📝 通俗解释：这一步是"叫大家开始干活"。每个进程独立执行训练脚本，但通过DDP机制自动协调。

整体代码汇总 ​

python
### DDP

# 引入包
import argparse
import torch.distributed as dist

# 设置可选参数
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

# 1. 初始化进程组
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

# 2. 使用DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
dataloader = DataLoader(
    dataset, 
    batch_size=batch_size, 
    shuffle=(train_sampler is None),
    sampler=train_sampler, 
    pin_memory=False
)

# 3. 创建DDP模型进行分布式训练
model = nn.parallel.DistributedDataParallel(
    model, 
    device_ids=[args.local_rank],
    output_device=args.local_rank, 
    find_unused_parameters=True
)

启动命令：

bash
python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 --node_rank=0 --master_port=6005 train.py

📝 通俗解释：相较于DP，DDP传输的数据量更少，因此速度更快，效率更高。

五、nn.parallel.DistributedDataParallel 参数更新流程 ​

  1. 初始化同步：process group中的训练进程都起来后，rank为0的进程会将网络初始化参数broadcast到其它每个进程中，确保每个进程中的网络都是一样的初始化的值（默认行为，可通过参数禁止）

  2. 数据分配：每个进程各自读取各自的训练数据，DistributedSampler确保了进程两两之间读到的是不一样的数据

  3. 前向计算：前向和loss的计算如今都是在每个进程上（也就是每个CUDA设备上）独立计算完成的；网络的输出不再需要gather到master进程上了，这和DP显著不一样

  4. 梯度同步：反向阶段，梯度信息通过all-reduce的MPI原语，将每个进程中计算到的梯度reduce到每个进程；也就是backward调用结束后，每个进程中的param.grad都是一样的值；注意，为了提高all-reduce的效率，梯度信息被划分成了多个buckets

  5. 参数更新：因为刚开始模型的参数是一样的，而梯度又是all-reduced的，这样更新完模型参数后，每个进程/设备上的权重参数也是一样的。因此，就无需DP那样每次迭代后需要同步一次网络参数，这个阶段的broadcast操作就不存在了

  6. Buffer同步：注意，Network中的Buffers（比如BatchNorm数据）需要在每次迭代中从rank为0的进程broadcast到进程组的其它进程上

📝 通俗解释：整个过程就像一个"自转小宇宙"：每个人独立学习（独立前向/反向计算），然后互相确认学习成果（All-Reduce梯度），最后各自更新自己的知识（更新参数）。不需要等老师（GPU 0）来统一分发作业，非常高效。

六、DP vs DDP 对比 ​

特性DataParallel (DP)DistributedDataParallel (DDP)
实现方式单进程多线程多进程
GIL问题存在Python解释器GIL性能开销避免GIL问题
负载均衡存在负载不均衡问题（0卡压力大）无负载不均衡问题
支持场景仅单机多卡支持多机多卡
参数更新梯度汇总到GPU0 → 更新 → 广播参数各进程独立更新（梯度已同步）
通信量较大较小
效率较低较高

详细说明 ​

  1. DDP通过多进程实现的。操作系统会为每个GPU创建一个进程，从而避免了Python解释器GIL带来的性能开销。而DataParallel()是通过单进程控制多线程来实现的。还有一点，DDP也不存在前面DP提到的负载不均衡问题。

  2. 参数更新的方式不同：

    • DDP：各进程梯度计算完成 → 汇总平均 → rank 0广播到所有进程 → 各进程独立更新参数
    • DP：梯度汇总到GPU0 → 反向传播更新参数 → 广播参数给其他GPU

  3. DDP支持 all-reduce、broadcast、send 和 receive 等等。通过MPI实现CPU通信，通过NCCL实现GPU通信，缓解了进程间通信开销大的问题。

📝 通俗解释：DP是"一个老师教多个学生"（单进程多线程），DDP是"多个老师同时教学生"（多进程）。显然，多老师的效率更高，而且不会让其中一个老师累死。

七、DDP 优点 ​

DistributedDataParallel是PyTorch提供的一种更加高级的多GPU并行训练方式，适用于多机多GPU的情况。DistributedDataParallel使用了数据并行和模型并行两种方式，通过将模型参数和梯度分布到不同的GPU上来充分利用多个GPU进行训练。

主要优点：

  1. 内存效率更高：每个进程只保存一份模型参数和梯度，而不是像DP那样需要复制多份
  2. 通信效率更高：使用Ring-AllReduce算法，通信负载均衡
  3. 支持多机多卡：可以跨多台机器进行分布式训练
  4. 训练速度更快：各进程独立计算，可并行处理
  5. 容错性更好：某个进程失败不影响其他进程

📝 通俗解释：DDP就像一个"专业团队"——每个人各司其职，独立工作，效率高，配合默契。而DP就像"一个老板带着几个实习生"，老板忙得不可开交，实习生却在等待。

八、DDP 缺点 ​

使用DistributedDataParallel需要一定的分布式编程经验，使用也相对比较复杂。

主要缺点：

  1. 配置复杂：需要设置进程组、初始化通信等
  2. 调试困难：多进程调试比单进程复杂
  3. 学习成本高：需要理解分布式训练的概念和机制
  4. 资源要求高：每个进程都需要独立的CPU资源

📝 通俗解释：DDP虽然好，但就像开车——手动挡（DP）简单易学，自动挡（DDP）虽然更好开，但你需要先考驾照（学习分布式知识）。

总结 ​

对比项DPDDP
使用难度简单较复杂
训练速度较慢快
内存占用高低
支持规模单机多卡多机多卡
通信效率低（单点瓶颈）高（分布式）
推荐场景快速原型验证生产级训练

📝 通俗解释：如果你是学生做作业（实验/原型），用DP就够了，简单省事；如果你是要建工厂（生产环境大规模训练），一定要用DDP，效率高、扩展性强。

基于 MIT 许可发布

版权所有 © 2024 LLMs