面试官：你大模型是怎么分布式训练的？_大模型多台机器分布式训练-优快云博客

本文链接：https://blog.youkuaiyun.com/Code1994/article/details/143993101

一、背景

大模型很大，大到一张卡放不下，数据很多，多到一张卡加载会爆。前几年搞的分布式计算的思想用在多GPU上是个很不错的思路，不过两者有很大的差异，这里不西索了。分布式并行有很多中方式，有DP(数据并行)、TP(张量并行)、PP(流水线并行)等。

二、介绍

2.1 数据并行

DataParallel（DP）：采用的Parameter Server模式，只能单机多卡。在DP中，每个GPU上都拷贝一份完整的模型，每个GPU上处理batch的一部分数据，所有GPU算出来的梯度进行累加后，再传回各GPU用于更新参数（PS架构特性）。DP多采用参数服务器这一编程框架，一般由若个计算Worker和1个梯度聚合Server组成。Server与每个Worker通讯，Worker间并不通讯。因此Server承担了系统所有的通讯压力。基于此DP常用于单机多卡场景。
DistributedDataParallel（DDP）：采用的All-Reduce模式，可以单机多卡，也可多机多卡。DDP将通讯压力均衡地分到每个GPU上，使得跨机器的数据并行得以高效实现（AR架构特性）。DP和DDP的总通讯量相同，但因负载不均的原因，DP需要耗费更多的时间搬运数据。
ZeRO：零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式，旨在降低存储。

对比：

DP是单进程多线程的，只能在单机上工作；DDP是多进程的，可以在多级多卡上工作。DP通常比DDP慢，主要原因有：1）DP是单进程的，受到GIL的限制；2）DP每个step都需要拷贝模型，以及划分数据和收集输出；
DDP可以与模型并行相结合；
DP的通信成本随着卡数线性增长，DDP支持Ring-AllReduce，通信成本是固定的。

2.2 张量并行

TensorParallel(TP): 将计算图中的层内的参数（张量）切分到不同设备（即层内并行），每个设备只拥有模型的一部分，以减少内存负荷，我们称之为张量模型并行。说白了就是对权重矩阵和X进行拆分，按照行和列拆开，然后并行计算。当然，不同层如Liner层，att层的计算方式或者说并行手段不一样，需要有针对性分析。TP也称为层内并行。

2.3 流水线并行

PipeLineParallel(PP):流水线并行，顾名思义肯定是将模型拆分为前后串起来的流水，当模型太大，一块GPU放不下时，流水线并行将模型的不同层放到不同的GPU上，通过切割mini-batch实现对训练数据的流水线处理，提升GPU计算通讯比。同时通过re-materialization机制降低显存消耗。但在实际应用中，流水线并行并不特别流行，主要原因是模型能否均匀切割，影响了整体计算效率，这就需要去做手调。

三、数据并行(DP)

3.1 流程

计算流程：

若干块计算GPU，如图中GPU0~GPU2；其中某块GPU负责收集和下发梯度
在每块计算GPU上都拷贝一份完整的模型参数。
把一份数据X（例如一个batch）均匀分给不同的计算GPU。
每块计算GPU做一轮FWD和BWD后，算得一份梯度G。
每块计算GPU将自己的梯度push给梯度收集GPU，做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义。
梯度收集GPU聚合完毕后，计算GPU从它那pull下完整的梯度结果，用于更新模型参数W。更新完毕后，计算GPU上的模型参数依然保持一致。
聚合再下发梯度的操作，称为AllReduce。

3.2 优缺点

方便是方便，缺点：存储开销大，通讯开销大

3.3 代码

一个案例代码如下所示

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

# 假设我们有一个简单的数据集类
class SimpleDataset(Dataset):
def __init__(self, data, target):
        self.data = data
        self.target = target

def __len__(self):
return len(self.data)

def __getitem__(self, idx):
return self.data[idx], self.target[idx]

# 假设我们有一个简单的神经网络模型
class SimpleModel(nn.Module):
def __init__(self, input_dim):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(input_dim, 1)

def forward(self, x):
return torch.sigmoid(self.fc(x))

# 假设我们有一些数据
n_sample = 100
n_dim = 10
batch_size = 10
X = torch.randn(n_sample, n_dim)
Y = torch.randint(0, 2, (n_sample, )).float()

dataset = SimpleDataset(X, Y)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# ===== 注意：刚创建的模型是在 cpu 上的 ===== #
device_ids = [0, 1, 2]
model = SimpleModel(n_dim).to(device_ids[0])

# 重点1：对上面的model再包装
model = nn.DataParallel(model, device_ids=device_ids)


optimizer = optim.SGD(model.parameters(), lr=0.01)

for epoch in range(10):
for batch_idx, (inputs, targets) in enumerate(data_loader):
        inputs, targets = inputs.to('cuda'), targets.to('cuda')
        outputs = model(inputs)

        loss = nn.BCELoss()(outputs, targets.unsqueeze(1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}')

代码来自参考[1]，测试过了没有问题，可以正常运行。复制到一个py文件中，然后点击run就可以了。

[1] https://zhuanlan.zhihu.com/p/676142368

四、分布式数据并行(DDP)

4.1 背景

DP一般用于单机多卡场景。因此，DDP作为一种更通用的解决方案出现了，既能多机，也能单机。DDP首先要解决的就是通讯问题：将Server上的通讯压力均衡转到各个Worker上。实现这一点后，可以进一步去Server，留Worker。
前文我们说过，聚合梯度 + 下发梯度这一轮操作，称为AllReduce。接下来我们介绍目前最通用的AllReduce方法：Ring-AllReduce。它由百度最先提出，非常有效地解决了数据并行中通讯负载不均的问题，使得DDP得以实现。

4.2 流程

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

可以看到GPU0和GPU1是连接的，1和2是连接的，这样依次连接形成一个环，上面的a0是的参数梯度，看到每个GPU上有a b c d这4块的梯度，对应的是模型中不同部分参数的梯度吗，所以整体上的梯度是a+b+c+d,就是a0+b0+c0+d0+a1+b1+c1+d1+a2+b2+c2+d2+a3+b3+c3+d3, 每个GPU上会基于当前分到的部分数据更新参数的梯度，更新完了之后就立马传给后面的GPU，比如对于GPU0来说，它分到了1/4的数据，那么咱就基于这1/4的数据更新w0,…,w3的参数，更新的时候更新到w4的梯度后，立马传给GPU1，然后继续更新w3，我们GPU1在GPU0更新w4的时候也不会闲着，它可以更新w3,然后给GPU2，这样类似4个人站好一个队伍，有4块大石头要从0搬到3，那么各个GPU依次搬一点，形成流水线，而不是依次把4块大石头全部搬完过去，那开销也大，也累。

这一步得到的结果如下：

可以看到每一块GPU渡海十上都有一份完整的几次聚合,注意不是所有的梯度。

第二步:All-Gather, 就是把红色那部分的梯度广播到其他GPU上去，如名字里Gather所述的一样，这操作里依然按照“相邻GPU对应位置进行通讯”的原则，但对应位置数据不再做相加，而是直接替换。All-Gather以红色块作为起点。

4.3 代码(单机分布式)

代码小的如下：

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


class DummyModel(nn.Module):
def __init__(self):
        super(DummyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.net2 = nn.Sequential(nn.Linear(10, 10), nn.LayerNorm(10))
        self.net3 = nn.Linear(10, 10)
        self.layer_norm = nn.LayerNorm(10)

def forward(self, x):
return self.layer_norm(self.net3(self.net2(self.net1(x))))


def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
if rank == 0:
        print(f"local rank: {rank}, world size: {dist.get_world_size()}")
    torch.cuda.set_device(rank)
    model = DummyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
for i in range(1000):
        outputs = ddp_model(torch.randn(20, 10).to(rank))
        labels = torch.randn(20, 10).to(rank)
        loss_fn(outputs, labels).backward()
        optimizer.step()
if rank == 0 and i % 100 == 0:
            print(f"Iteration: {i/1000 * 100} %")
if rank == 0:
        print("Training completed.")


if __name__ == '__main__':
    main()
# torchrun --nproc_per_node=2 --nnodes=1 --standalone ddp_example.py

记得好好看上面代码中的重点标记，别漏掉了。代码见参考[1]，代码写好后，使用 torchrun --nproc_per_node=2 --nnodes=1 --standalone ddp_example2.py --local_rank 0进行启动，别用参考[1]提到的命令，那个命令已经过时了。

4.4 代码(多机分布式)

代码上基本没啥改变，假设我们有两台机器，则启动代码可能要变一下

>>> #节点1
torchrun --nproc_per_node 8  --nnodes 2 --node_rank 0 --rdzv_endpoint 10.192.2.1:62111  ddp_example.py
>>> #节点2
torchrun --nproc_per_node 8  --nnodes 2 --node_rank 1 --rdzv_endpoint 10.192.2.1:62111  ddp_example.py
上面也可以不写rdzv_endpoint而直接指定

--master_addr="192.168.1.1"
         --master_port=1234

也是可以的。

其中多机分布式分为两种，一种是每个进程占用一张卡，一种是每个进程占用多张卡。

每个进程占用一张卡

注意1. dist.init_process_group里面的rank需要根据node以及GPU的数量计算；2. world_size的大小=节点数 x GPU 数量。3. ddp 里面的device_ids需要指定对应显卡。

假设一共有两台机器（节点1和节点2），每个节点上有8张卡，节点1的IP地址为192.168.0.1 占用的端口22335（端口可以更换），启动的方式如下：

>>> #节点1
>>>python python demo.py --world_size=16 --node_rank=0 --master_addr="192.168.0.1" --master_port=22335
>>> #节点2
>>>python python demo.py --world_size=16 --node_rank=1 --master_addr="192.168.0.1" --master_port=22335

两张GPU有16张卡，因此word_size为16(进程数)，就是每张卡由1个进程来控制。

每个进程控制多张卡

看单进程多卡的代码和启动方式，代码中需要注意的位置：1. dist.init_process_group里面的rank等于节点编号；2. world_size等于节点的总数量；3. DDP不需要指定device。

启动方式如下：

>>> #节点1
>>>python demo.py --world_size=2 --rank=0 --master_addr="192.168.0.1" --master_port=22335
>>> #节点2
>>>python demo.py --world_size=2 --rank=2 --master_addr="192.168.0.1" --master_port=22335

当然上面是写成了python -m，理论上要写成torchrun才可以。

这一小部分的代码和解释可以看参考3

4.5 其它

1. 模型保存

让一个rank保存模型就好了，没必要每个rank都要保存，这样太耗时了

if torch.distributed.get_rank() == 0:  #一般用0，当然，可以选任意的rank保存。  
    torch.save(net, "net.pth")

2. 时间开销比较

两者的总通讯量虽然相同，但DP存在负载不均的情况，大部分的通讯压力集中在Server上，而Server的通讯量与GPU数量呈线性关系，导致DP一般适用于单机多卡场景。而DDP通过采用Ring-AllReduce这一NCCL操作，使得通讯量均衡分布到每块GPU上，且该通讯量为一固定常量，不受GPU个数影响，因此可实现跨机器的训练。

4.6 相关概念

这里说一下rank这些东西的概念.

我们可以在cmd里面运行 torchrun -h 查看所有参数的定义

rank：用于表示进程的编号/序号（在一些结构图中rank指的是软节点，rank可以看成一个计算单位），每一个进程对应了一个rank的进程，整个分布式由许多rank完成。
node：物理节点，可以是一台机器也可以是一个容器，节点内部可以有多个GPU。
rank与local_rank：rank是指在整个分布式任务中进程的序号；local_rank是指在一个node上进程的相对序号，local_rank在node之间相互独立。
nnodes、node_rank与nproc_per_node：nnodes是指物理节点数量，node_rank是物理节点的序号；nproc_per_node是指每个物理节点上面进程的数量。
word size ：全局（一个分布式任务）中，rank的数量。

该定义参考[2]. 图解如下所示：

图中：一共有12个rank，nproc_per_node=4，nnodes=3，每个节点都一个对应的node_rank。

通信过程主要是完成模型训练过程中参数信息的传递，主要考虑通信后端和通信模式选择，后端与模式对整个训练的收敛速度影响较大，相差可达2~10倍。 在DDP中支持了几个常见的通信库，而数据处理的模式写在PyTorch底层，供用户选择的主要是后端。在初始化时需要设置：

- backend ：通信后端，可选的包括：nccl（NVIDIA推出）、gloo（Facebook推出）、mpi（OpenMPI）。从测试的效果来看，如果显卡支持nccl，建议后端选择nccl，，其它硬件（非N卡）考虑用gloo、mpi（OpenMPI）。

参考

[2] https://zhuanlan.zhihu.com/p/358974461
[3] 找不到了…
[4] https://zhuanlan.zhihu.com/p/358974461
[5] Pytorch 分布式训练DDP(torch.distributed)详解-原理-代码_pytorch ddp-优快云博客【DDP原理】

五、ZERO数据并行

5.1 思路

首先分析了一波，觉得梯度、优化器、模型参数、数据这些都可以切分，那就干脆放在不同的卡上，用的时候进行通信就可以了，这就是你那通讯来换显存。

5.2 优化状态分割

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，各得一份梯度。
（2）对梯度做一次AllReduce，得到完整的梯度G，产生单卡通讯量 2Φ2\Phi 。为了表达简明，这里通讯量我们就不再换算成byte了，而直接根据参数量来计算。
（3）得到完整梯度G，就可以对W做更新。我们知道W的更新由optimizer states和梯度共同决定。由于每块GPU上只保管部分optimizer states，因此只能将相应的W（蓝色部分）进行更新。（2）和（3）可以用下图表示：

（4）此时，每块GPU上都有部分W没有完成更新（图中白色部分）。所以我们需要对W做一次All-Gather，从别的GPU上把更新好的部分W取回来。

5.3 优化状态与梯度分割

思路是简单的，主要流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，算得一份完整的梯度（下图中绿色+白色）。
（2）对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度。例如对GPU1，它负责维护G1，因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可。汇总完毕后，白色块对GPU无用，可以从显存中移除。单卡通讯量 Φ\Phi 。（1）和（2）见下图：

（3）每块GPU用自己对应的O和G去更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次All-Gather，将别的GPU算好的W同步到自己这来。单卡通讯量 Φ\Phi

5.4 优化状态、梯度与参数分割

数据并行的流程如下：
（1）每块GPU上只保存部分参数W。将一个batch的数据分成3份，每块GPU各吃一份。
（2）做forward时，对W做一次All-Gather，取回分布在别的GPU上的W，得到一份完整的W，单卡通讯量 Φ\Phi 。forward做完，立刻把不是自己维护的W抛弃。
（3）做backward时，对W做一次All-Gather，取回完整的W，单卡通讯量 Φ\Phi 。backward做完，立刻把不是自己维护的W抛弃。
（4）做完backward，算得一份完整的梯度G，对G做一次Reduce-Scatter，从别的GPU上聚合自己维护的那部分梯度，单卡通讯量 Φ\Phi 。聚合操作结束后，立刻把不是自己维护的G抛弃。
（5）用自己维护的O和G，更新W。由于只维护部分W，因此无需再对W做任何AllReduce操作。

5.5 ZeRO-Offload与ZeRO-Infinity

最后，简单介绍一下ZeRO-Offload。它的核心思想是：显存不够，内存来凑。如果我把要存储的大头卸载(offload)到CPU上，而把计算部分放到GPU上，这样比起跨机，是不是能既降显存，也能减少一些通讯压力呢？
ZeRO-Offload的做法是：

forward和backward计算量高，因此和它们相关的部分，例如参数W（fp16），activation，就全放入GPU。
update的部分计算量低，因此和它相关的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等。

ZeRO-Infinity 则在 ZeRO-Offload 的基础上进一步优化，主要包括三个方面。一是将和 ZeRO 的结合从 ZeRO-2 延伸到了 ZeRO-3，解决了模型参数受限于单张 GPU 内存的问题；二是解决了 ZeRO-Offload 在训练 batch size 较小的时候效率较低的问题；三是除 CPU 内存外，进一步尝试利用 NVMe 的空间。

5.6 总结

从deepspeed提供的zero的启动方式来看，

ZeRO的stage-1（优化Optimizer’s State），每个GPU只保存〖自己看管的层〗的优化器状态。每个GPU在更新某层权重时，只需要向〖保管该层〗的GPU要对应的优化器状态即可。
ZeRO的stage2（优化Optimizer’s State+Gradient），在stage-1的基础上，每个GPU只负责〖自己看管的层〗的全局梯度（需要收集各个GPU上该层的梯度）即可。
ZeRO的stage3（优化Optimizer’s State+Gradient+Parameter），在stage-1和stage-2的基础上，每个GPU负责〖自己管的层〗的FP16权重，每个GPU在Forward和Back-propagation的时候只需要向〖保管该层〗的GPU要对应的模型权重即可。

zero-0就是DDP了

从左到右，训练的速度越来越慢
Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offloads
从左到右，所需GPU显存越来越少
Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offloads

5.7 调试策略

batch_size设置为1，通过梯度累积实现任意的有效batch_size
如果OOM则，设置--`gradient_checkpointing` 1 (HF Trainer)，或者 model.gradient_checkpointing_enable()
如果OOM则，尝试ZeRO stage 2
如果OOM则，尝试ZeRO stage 2 + offload_optimizer
如果OOM则，尝试ZeRO stage 3
如果OOM则，尝试offload_param到CPU
如果OOM则，尝试offload_optimizer到CPU
如果OOM则，尝试降低一些默认参数。比如使用generate时，减小beam search的搜索范围
如果OOM则，使用混合精度训练，在Ampere的GPU上使用bf16，在旧版本GPU上使用fp16
如果仍然OOM，则使用ZeRO-Infinity ，使用offload_param和offload_optimizer到NVME
一旦使用batch_size=1时，没有导致OOM，测量此时的有效吞吐量，然后尽可能增大batch_size
开始优化参数，可以关闭offload参数，或者降低ZeRO stage，然后调整batch_size，然后继续测量吞吐量，直到性能比较满意（调参可以增加66%的性能）