大模型分布式训练并行技术（二）-数据并行-六虎

本文为稀土技能社区首发签约文章，30天内制止转载，30天后未获授权制止转载，侵权必究！

近年来，跟着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规划参数，传统的单机单卡形式现已无法满足超大模型进行练习的要求。因而，咱们需求依据单机多卡、甚至是多机多卡进行分布式大模型的练习。

而运用AI集群，使深度学习算法更好地从大量数据中高效地练习出功能优良的大模型是分布式机器学习的首要方针。为了完结该方针，一般需求依据硬件资源与数据/模型规划的匹配状况，考虑对核算使命、练习数据和模型进行区分，从而进行分布式存储和分布式练习。因而，分布式练习相关技能值得咱们进行深入分析其背后的机理。

下面主要对大模型进行分布式练习的并行技能进行解说，本系列大体分九篇文章进行解说。

大模型分布式练习并行技能（一）-概述
大模型分布式练习并行技能（二）-数据并行
大模型分布式练习并行技能（三）-流水线并行
大模型分布式练习并行技能（四）-张量并行
大模型分布式练习并行技能（五）-序列并行
大模型分布式练习并行技能（六）-多维混合并行
大模型分布式练习并行技能（七）-主动并行
大模型分布式练习并行技能（八）-MOE并行
大模型分布式练习并行技能（九）-总结

本文为分布式练习并行技能的第二篇：数据并行。因为其原理相对比较简略，因而，在日常会应用顶用的比较多。

简述

所谓数据并行，就是因为练习数据集太大；因而，将数据集分为N份，每一份别离装载到N个GPU节点中，一起，每个GPU节点持有一个完好的模型副本，别离依据每个GPU中的数据去进行梯度求导。然后，在GPU0上对每个GPU中的梯度进行累加，最后，再将GPU0聚合后的成果播送到其他GPU节点。

注意：这里是以GPU0作为参数服务器，除此之外，还可以运用CPU作为参数服务器。可是这种场景的练习速度一般会慢于运用GPU0作为参数服务器（一般状况下，GPU与CPU之间通讯运用PCIe，而GPU与GPU之间通讯运用Nvlink）。

当然，还可以将参数服务器分布在一切GPU节点上面，每个GPU只更新其间一部分梯度。

当然，数据并行不仅仅指对练习的数据并行操作，还可以对网络模型梯度、权重参数、优化器状况等数据进行并行。

下面主要以PyTorch中数据并行的发展为主线讲述现有一些数据并行办法。

数据并行（PyTorch DP）

数据并行(torch.nn.DataParallel)，这是Pytorch最早供给的一种数据并行办法，它依据单进程多线程进行完结的，它运用一个进程来核算模型权重，在每个批处理期间将数据分发到每个GPU。

DataParallel 的核算进程如下所示：

将 inputs 从主 GPU 分发到一切 GPU 上。
将 model 从主 GPU 分发到一切 GPU 上。
每个 GPU 别离独立进行前向传达，得到 outputs。
将每个 GPU 的 outputs 发回主 GPU。
在主 GPU 上，经过 loss function 核算出 loss，对 loss function 求导，求出丢失梯度。
核算得到的梯度分发到一切 GPU 上。
反向传达核算参数梯度。
将一切梯度回传到主 GPU，经过梯度更新模型权重。
不断重复上面的进程。

它运用十分简略，仅需一行代码即可完结。

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
output = net(input_var)  # input_var can be on any device, including CPU

可是它的缺陷也很明显：

单进程多线程带来的问题：DataParallel运用单进程多线程进行完结的，方便了信息的交流，但受困于 GIL，会带来功能开支，速度很慢。而且，只能在单台服务器（单机多卡）上运用（不支撑分布式）。一起，不能运用 Apex 进行混合精度练习。
功率问题，主卡功能和通讯开支容易成为瓶颈，GPU 运用率一般很低：数据集需求先拷贝到主进程，然后再分片（split）到每个设备上；权重参数只在主卡（GPU0）上更新，需求每次迭代前向一切设备做一次同步；每次迭代的网络输出需求集合到主卡（GPU0）上。因而，通讯很快成为一个瓶颈。除此之外，这将导致主卡和其他卡之间，GPU运用率严峻不均衡（比如：主卡运用了10G显存，而其他卡只运用了2G显存，batch size略微设置大一点主卡的显存就OOM了）。
不支撑模型并行，因为其自身的局限性，没办法与模型并行组合运用。

当然，现在PyTorch官方建议运用DistributedDataParallel，而不是DataParallel类来进行多 GPU 练习，即便在单机多卡的状况下。那么下面咱们来看看PyTorch DDP。

分布式数据并行（PyTorch DDP）

分布式数据并行(torch.nn.DistributedDataParallel)，依据多进程进行完结的，每个进程都有独立的优化器，履行自己的更新进程。每个进程都履行相同的使命，而且每个进程都与一切其他进程通讯。进程（GPU）之间只传递梯度，这样网络通讯就不再是瓶颈。

详细流程如下：

首先将 rank=0 进程中的模型参数播送到进程组中的其他进程；
然后，每个 DDP 进程都会创立一个local Reducer来担任梯度同步。
在练习进程中，每个进程从磁盘加载 batch 数据，并将它们传递到其 GPU。每个 GPU 都有自己的前向进程，完结前向传达后，梯度在各个 GPUs 间进行 All-Reduce，每个 GPU 都收到其他 GPU 的梯度，从而可以独自进行反向传达和参数更新。
一起，每一层的梯度不依赖于前一层，所以梯度的 All-Reduce 和后向进程一起核算，以进一步缓解网络瓶颈。
在后向进程的最后，每个节点都得到了均匀梯度，这样各个 GPU 中的模型参数坚持同步。

而DataParallel是将梯度 reduce 到主卡，在主卡上更新参数，再将参数 broadcast 给其他 GPU，这样无论是主卡的负载仍是通讯开支都比 DDP 大很多)，比较于DataParallel，DistributedDataParallel办法可以更好地进行多机多卡运算，更好的进行负载均衡，运转功率也更高，虽然运用起来较为麻烦，但关于追求功能来讲是一个更好的挑选。

以下为DistributedDataParallel的简略示例，运用 torch.nn.Linear 作为本地模型，用 DDP 对其进行包装，然后在 DDP 模型上运转一次前向传达、一次反向传达和更新优化器参数进程。之后，本地模型上的参数将被更新，而且不同进程上的一切模型完全相同。

import torch
import t dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()
def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)
if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()

DP与DDP 的区别

DP 和 DDP 的主要差异有以下几点：

DP 是依据单进程多线程的完结，只用于单机状况，而DDP 是多进程完结的，每个 GPU 对应一个进程，适用于单机和多机状况，真实完结分布式练习，而且因为每个进程都是独立的 Python 解说器，DDP避免了 GIL 带来的功能开支。
参数更新的办法不同。DDP在各进程梯度核算完结之后，各进程需求将梯度进行汇总均匀，然后再由 rank=0 的进程，将其播送到一切进程后，各进程用该梯度来独立的更新参数（而 DP是梯度汇总到 GPU0，反向传达更新参数，再播送参数给其他剩余的 GPU）。因为DDP各进程中的模型，初始参数共同 (初始时间进行一次播送)，而每次用于更新参数的梯度也共同；因而，各进程的模型参数一直坚持共同（而在DP中，全程维护一个 optimizer，对各个GPU上梯度进行求均匀，而在主卡进行参数更新，之后再将模型参数播送到其他GPU）。相较于DP，DDP传输的数据量更少，练习更高效，不存在 DP 中负载不均衡的问题。现在，根本上 DP 现已被弃用。
DDP 支撑模型并行，而 DP 并不支撑，这意味如果模型太大单卡显存缺乏时，只能运用DDP。

补充阐明：DP与DDP数据传输进程

DP数据传输进程：

前向传达得到的输出成果gather到主cuda核算loss
scatter上述loss到各个cuda
各个cuda反向传达核算得到梯度后gather到主cuda后，主cuda的模型参数被更新。
主cuda将模型参数broadcast到其它cuda设备上，至此，完结权重参数值的同步。

综上，DP大概是有4次输出传输。

DDP数据传输进程：

前向传达的输出和loss的核算都是在每个cuda独立核算的，梯度all-reduce到一切的CUDA(传输梯度)，这样初始参数相同，para.grad也相同，反向传达后参数就仍是坚持共同的，其他没有数据传输了。

完全分片数据并行(PyTorch FSDP)

因为 PyTorch FSDP 受 DeepSpeed ZeRO 启示而获得创意，因而，下面先扼要介绍下 ZeRO。

补充阐明：ZeRO

一般来说，在模型练习的进程中，GPU上需求进行存储的参数包含了模型自身的参数、优化器状况、激活函数的输出值、梯度以及一些零时的Buffer。各种数据的占比如下图所示：

可以看到模型参数仅占模型练习进程中一切数据的一部分，当进行混合精度运算时，其间模型状况参数(优化器状况+ 梯度+ 模型参数）占到了一大半以上。因而，咱们需求想办法去除模型练习进程中的冗余数据。

针对模型状况的存储优化（去除冗余），DeepSpeed 提出了 ZeRO，ZeRO 运用的办法是分片，即每张卡只存1/N的模型状况量，这样体系内只维护一份模型状况参数。

ZeRO对模型状况（Model States）参数进行不同程度的分割，主要有三个不同等级：

ZeRO-1 : 优化器状况分片（ Optimizer States Sharding）
ZeRO-2 : 优化器状况与梯度分片（Optimizer States & Gradients Sharding）
ZeRO-3 : 优化器状况、梯度和模型权重参数分片（Optimizer States & Gradients & Parameters Sharding）

ZeRO-1：

ZeRO-1没有将模型自身进行分片，也没有将Gradient进行分片，而是只将优化器进行分片。练习进程与DDP相似。

forward进程由每个rank的GPU独自完好的完结，然后进行backward进程。在backward进程中，梯度经过allReduce进行同步。
Optimizer state 运用贪心战略依据参数量进行分片，以此保证每个rank几乎具有相同巨细的优化器内存。
每个rank只担任更新当前优化器分片的部分，因为每个rank只要分片的优化器state，所以当前rank疏忽其他的state。
在更新往后，经过播送或者allGather的办法保证一切的rank都收到最新更新往后的模型参数。

ZeRO-1 十分合适运用相似Adam进行优化的模型练习，因为Adam具有额定的参数m（momentum）与v（variance），特别是FP16混合精度练习。ZeRO-1 不合适运用SGD相似的优化器进行模型练习，因为SGD只要较少的参数内存，而且因为需求更新模型参数，导致额定的通讯成本。ZeRO-1只是解决了Optimizer state的冗余。

ZeRO-2：

比较于ZeRO-1，ZeRO-2除了对optimizer state进行切分，还对Gradient进行了切分。

像ZeRO-1一样将optimizer的参数进行分片，并安排在不同的rank上。在backward进程中，gradients被reduce操作到对应的rank上，取代了all-reduce，以此减少了通讯开支。每个rank独自更新各自担任的参数。在更新操作之后，播送或allGather保证一切的ranks接收到更新后的参数。

ZeRO-3：

为了进一步节约更多的内存，ZeRO-3提出进行模型参数的分片。相似以上两种分片办法，ranks担任模型参数的切片。可以进行参数切片的原因主要有以下两点：

All-Reduce操作可以被拆分为Reduce与allgather操作的结合。
模型的每一层具有该层的完好参数，而且整个层可以直接被一个GPU装下。所以核算前向的时分，除了当前rank需求的层之外，其他的层的参数可以扔掉。从这个层面上来说，Zero相当于数据并行+模型并行。

FSDP

完全分片数据并行(torch.distributed.fsdp.FullyShardedDataParallel)，是Pytorch最新的数据并行方案，在1.11版本引进的新特性，目的主要是用于练习大模型。咱们都知道Pytorch DDP用起来简略方便，可是要求整个模型加载到一个GPU上，这使得大模型的练习需求运用额定杂乱的设置进行模型分片。因而，为了打破模型分片的妨碍（包含模型参数，梯度，优化器状况）；一起，仍然坚持了数据并行的简略性，该新特性应运而生。

FSDP 是一种新式数据并行练习办法，但与传统的数据并行不同，传统的数据并行维护模型参数、梯度和优化器状况的每个 GPU 副本，而 FSDP 将一切这些状况跨数据并行工作线程进行分片，而且可以挑选将模型参数分片卸载到 CPU。

下图显示了 FSDP 怎么在 2 个数据并行进程中工作流程：

一般，模型层以嵌套办法用 FSDP 包装，因而，只要单个 FSDP 实例中的层需求在前向或后向核算期间将完好参数搜集到单个设备。核算完结后，搜集到的完好参数将立即开释，开释的内存可用于下一层的核算。经过这种办法，可以节约峰值 GPU 内存，从而可以扩展练习以运用更大的模型巨细或更大的批量巨细。为了进一步最大化内存功率，当实例在核算中不活动时，FSDP 可以将参数、梯度和优化器状况卸载到 CPU。

解锁ZeRO/FSDP的关键是咱们可以把DDP之中的All-Reduce操作分解为独立的 Reduce-Scatter 和 All-Gather 操作。

All-Reduce 是 Reduce-Scatter 和 All-Gather 的组合。聚合梯度的规范 All-Reduce 操作可以分解为两个独自的阶段。

Reduce-Scatter 阶段，在每个GPU上，会依据 rank 索引对 rank 之间持平的块进行求和。
All-Gather 阶段，每个GPU上的聚合梯度分片可供一切GPU运用。

经过重新整理 Reduce-Scatter 和 All-Gather，每个 DDP worker只需求存储一个参数分片和优化器状况。

在 PyTorch 中运用 FSDP 包装模型有两种办法。

主动包装（Auto Wrapping）是 DDP 的直接替代品；
手动包装（Manual Wrapping）需求对模型界说代码进行少数的更改，而且可以探究杂乱的分片战略。

主动包装（Auto Wrapping）

模型层应以嵌套办法包装在 FSDP 中，以节约峰值内存并完结通讯和核算重叠。最简略的办法是主动包装，它可以作为 DDP 的直接替代品，而无需更改其他代码。

fsdp_auto_wrap_policy参数答应指定可调用函数以运用 FSDP 递归地包裹层。 PyTorch FSDP供给的default_auto_wrap_policy函数递归地包裹参数数量大于100M的层。当然，您也可以依据需求供给自己的包装战略。

此外，可以挑选配置 cpu_offload，以便在核算中不运用包装参数时将这些参数卸载到 CPU。这可以进一步进步内存功率，但价值是主机和设备之间的数据传输开支。

下面的示例展现了怎么运用主动包装（Auto Wrapping）来包装 FSDP。

from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)
from torch.distributed.fsdp.wrap import (
   default_auto_wrap_policy,
)
import torch.nn as nn
class model(nn.Module):
   def __init__(self):
       super().__init__()
       self.layer1 = nn.Linear(8, 4)
       self.layer2 = nn.Linear(4, 16)
       self.layer3 = nn.Linear(16, 4)
model = DistributedDataParallel(model())
fsdp_model = FullyShardedDataParallel(
   model(),
   fsdp_auto_wrap_policy=default_auto_wrap_policy,
   cpu_offload=CPUOffload(offload_params=True),
)

手动包装（Manual Wrapping）

经过有挑选地对模型的某些部分应用包装，手动包装关于探究杂乱的分片战略十分有用。总体设置可以传递给enable_wrap()上下文管理器。

from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)
from torch.distributed.fsdp.wrap import (
   enable_wrap,
   wrap,
)
import torch.nn as nn
from typing import Dict
class model(nn.Module):
   def __init__(self):
       super().__init__()
       self.layer1 = wrap(nn.Linear(8, 4))
       self.layer2 = nn.Linear(4, 16)
       self.layer3 = wrap(nn.Linear(16, 4))
wrapper_kwargs = Dict(cpu_offload=CPUOffload(offload_params=True))
with enable_wrap(wrapper_cls=FullyShardedDataParallel, **wrapper_kwargs):
   fsdp_model = wrap(model())

运用上述两种办法之一，用 FSDP 包装模型后，可以选用与本地练习相似的办法练习模型，详细如下所示：

optim = torch.optim.Adam(fsdp_model.parameters(), lr=0.0001)
for sample, label in next_batch():
  out = fsdp_model(input)
  loss = criterion(out, label)
  loss.backward()
  optim.step()

DDP 与 FSDP 的区别

在规范的数据并行（DistributedDataParallel）练习办法中，每个GPU上都有一个模型副本，向前和向后传递的序列只在自己的数据分片进步行运转。在这些部分核算之后，每个部分进程的参数和优化器与其他GPU共享，以便核算大局权重更新。

而在FullyShardedDataParallel练习办法中：

Model shard：每个GPU上仅存在模型的分片。
All-gather：每个GPU经过all-gather从其他GPU搜集一切权重，以在本地核算前向传达。
Forward（local）：在本地进行前向操作。前向核算和后向核算都是运用完好模型。
All-gather：然后在后向传达之前再次履行此权重搜集。
Backward（local）：本地进行后向操作。前向核算和后向核算都是运用完好模型，此刻每个GPU上也都是悉数梯度。
Reduce-Scatter：在向后传达之后，部分梯度被聚合而且经过 Reduce-Scatter 在各个GPU上分片，每个分片上的梯度是聚合之后本分片对应的那部分。
Update Weight（local）：每个GPU更新其部分权重分片。

一起，为了最大限度地进步内存功率，咱们可以在每层前向传达后丢弃悉数权重，为后续层节约内存。这可以经过将 FSDP 包装应用于网络中的每一层来完结（经过设置reshard_after_forward=True）。

总结

本文主要解说了大模型分布式练习并行技能的数据并行，并以Pytorch为主线解说了DP、DDP、FSDP三种不同的数据并行方案。

DP 主要存在如下问题：

单进程多线程形式，因为锁的机制导致线程间同步存在瓶颈。
运用普通的All-Reduce机制，一切的卡需求将梯度同步给0号节点，并由0号节点均匀梯度后反向传达，再分发给一切其他节点，意味着0号节点负载很重。
因为第二点的原因，导致0号GPU通讯成本是跟着GPU数量的上升而线性上升的。
不支撑多机多卡。

现在，因为功能问题，DP根本不用了。

而 DDP 是多进程完结的，每个 GPU 对应一个进程，适用于单机和多机状况，真实完结分布式练习，而且因为每个进程都是独立的 Python 解说器，DDP避免了 GIL 带来的功能开支。

DDP在各进程梯度核算完结之后，各进程需求将梯度进行汇总均匀，然后再由 rank=0 的进程，将其播送到一切进程后，各进程用该梯度来独立的更新参数。因为DDP各进程中的模型，初始参数共同 (初始时间进行一次播送)，而每次用于更新参数的梯度也共同；因而，各进程的模型参数一直坚持共同。相较于DP，DDP传输的数据量更少，练习更高效，不存在 DP 中负载不均衡的问题。

虽然Pytorch DDP完结了真实的分布式练习，一起，避免了DP 中负载不均衡的问题，可是，要求整个模型加载到一个GPU上，这使得大模型的练习需求运用额定杂乱的设置进行模型分片。因而，为了打破模型分片的妨碍（包含模型参数，梯度，优化器状况），一起仍然坚持了数据并行的简略性，FSDP应运而生。

如果觉得我的文章可以可以给您带来协助，期待您的点赞收藏加关注~~

大模型分布式训练并行技术（二）-数据并行

简述

数据并行（PyTorch DP）

分布式数据并行（PyTorch DDP）

DP与DDP 的区别

补充阐明：DP与DDP数据传输进程

完全分片数据并行(PyTorch FSDP)

补充阐明：ZeRO

FSDP

DDP 与 FSDP 的区别

总结

相关文章

优秀设计背后的思考：票务系统的架构分析

货拉拉大数据对BitMap的探索与实践上

2003. 每棵子树内缺失的最小基因值 : 图解三大结论秒懂本题

AI 制作 3D 素材｜基于 AI 5 天创建一个农场游戏，第 3 天

作者信息