GPU分布式训练缩写眼保健操：PP、DP、DDP、MP、ZeRO、FSDP

学习的第一步是让眼睛先熟悉

分布式计算的几个原生问题

在谈论分布式计算或大规模服务系统时，可能会谈论这几个问题
  1. 扩展的方向：纵向或横向
    1. 纵向扩展：比如一个复杂的业务系统，将其按功能和开发形式，一分为二，这是纵向扩展；纵向扩展后，微服务系统A和系统B，可以分开地考虑自己的迭代，部署方式，资源增长形式等问题。又比如TP数据库，做MYSQL分库分表的垂直拆分，将多写的表与多读的表，或按业务的性质划分不同的实例也是一种纵向扩展。纵向扩展的思想是分而治之，不同的问题分开看待，消耗了一定的团队/硬件资源，但清晰了权责，缩小了问题的边界。
    2. 横向扩展：把同一个问题切分成identity的若干个变小的问题。比如上述微服务系统A，是计算密集型应用，计算紧张时，可以部署若干套。输入的请求划分成若干份；若干份资源 对应 若干份请求，问题和解决方案被切分成更小等价的局部。再说TP数据库，MYSQL做分表分表的水平拆分，假设用户表的规模对单个实例来说压力过大，则可以通过路由的形式，拆分成N个等价的实例。AP系统更不用说，横向扩展是基本结构；计算100w个数字的avg，拆分成100个计算实例，分别计算1w个数的avg，将100个结果做合并，mapreduce的过程即横向扩展。
  1. shard与replica
    1. Shard是横向扩展的概念，表明横向扩展成几份，称呼shard数，做sharding，sharding集群等名词
    2. Replica是在横向扩展基础上衍生出的高可用需求。在无状态系统里做sharding，不需要考虑这个问题，节点之间是完全一致的。在有状态系统里做sharding需要考虑冗余，典型地如分布式存储系统。在有状态系统里，看似横向扩展的shardA和shardB，在支持的计算和存储的类型这些元信息上一致，但存储的内容本身就是一种有状态：对用户A的统计和用户B的统计计算过程一致，但用户A的结果绝不可能替代用户B。因此系统需要冗余来提供可用性，冗余本身也是一种并行度。计算100w个用户的avg，这100w个用户shard成10份，先拆分成10个负责10w用户的子任务；每个shard有3个replica，每个shard的计算任务，理论上“有机会”纵向扩展成3个计算3w用户的子任务。做不做这件事看系统的特性去论证收益。
  1. 通信模式与一致性：P2P or Master-Slave
    1. 做了sharding，做了replica，replica之间的数据就会产生一致性要求，即N份理论上要求一致的数据，如何保持一致？
    2. 一种模式是Master-Slave，典型如Raft/Paxos等分布式协议，通过选举投票机制确定“话事人”/Master，话事人执政期间，数据一切以Master为主。代价是大多数工作需要由master承接，承担“业务”的开销与“管理”的开销大头。优点是对于一致性来说，可靠且高效。
    3. 一种是P2P模型，replica之间没有一个明确的master角色，靠互通有无来决定是否有新数据。如果“外部人”不等他们同步，则只能取到不一定可靠的数据，数据还可能被回滚。但如果“外部人”愿意等待，数据最终会一致，此为最终一致性。代价就是不可靠，低效，好处是对每个节点的性能要求比较平均。
  1. 纵向系统里的流与批
    1. spark或者starrocks或者flink的分布式引擎做的是通过分析任务，先纵向拆分，再在纵向切分好的子任务里确定如何横向扩展。区别在于spark一类的批系统是按阶段（stage）在时间轴上串行执行的。而flink一类的流系统是按数据的流动在时间轴上并行执行的。
    2. 打个比方，spark系统的数据是快照性质的，集合当前数据的一个计算；就像修100km长的路，建设局是一段一段去修，有的路段难度小，用了100人的团队，有的地方难度大，需要用1000人，但不会在修下一段时，上一段的团队，资源还占着不放。而flink更像一个维护公路运作的团队，100km的路，每10km分为一个服务小组，不知道什么时候有多少车车会来，不可能等车来了再去组建团队，更不可能去现修马路。
    3. 微批：在flink大行其道之前，spark生态利用其完善的计算生态去完成流的计算，批的拆分就是流，从宏观层面不能说错误。
  1. 存算分离 vs 存算一体
    1. TP系统存算一体较多，AP系统存算分离较多。存算一体：性能好适合实时计算；存算分离：性能较差，但存储可以单独考虑成本与规模，计算有机会弹性扩容，如spark与hdfs，redshift等。
    2. 在GPU计算里，”存“可以看做显存或内存，而不是磁盘存储。

PP: Pipeline Parallelism & GPipe

Gpipe: https://arxiv.org/pdf/1811.06965
GPipe的研究大的方向是把模型纵向切分，第二个是运用了微批，第三个小优化是运用了re-materialization的技巧。
notion image
其概念，如论文中最重要的图所示，没有太多需要解释的地方，即把模型按串行的层序切分成K个部分，每个部分放进一个理论上的“加速设备”上。层间通过通信，交换必要的信息；如激活值，如梯度等。
Naive的串行方式由此引入了一个问题，在figure(b)部分所示，”Device i” 计算时，其他K-1个部分是空闲的。因此引入了micro-batch的概念，将整体的mini-batch，分成M份，“流动地”推进计算。这样有两个好处，一是计算设备的空闲时间减少了，另一个是单位时间内，mini-batch的逻辑大小可以成倍增加，有助于训练的收敛和稳定性。
第三点的优化是re-materialization。在单个设备内，假设存在若干层网络（平均是#L/#K），计算梯度所需的input，可以缓存下来，也可以不缓存。大多数情况下GPU分布式系统是memory-bound与network-bound的，计算能力是溢出的。因此不缓存下来，节省了内存消耗同时不增加通信开销，是一种很好的策略。类比spark，就是RDD的lineage机制，但那是为了节省因不可靠设备导致的开销，出发点不同。PagedAttention中也提及了内存淘汰的re-compuation策略，一种普遍的时间换空间的策略。
 
一个遗留的问题是，如何处理“跨样本的，依赖batch内全局信息”的计算，如batch norm。文中是这么描述的。
也就是说，使用了micro-batch内的数据来做mini-batch的统计量的一个预估。从逻辑上是会影响计算结果的，从统计意义上，我认为是无损的。因为文中给的实验，batch_size少说在十万级别。
 
另一个问题：Pipeline Parallelism能否像spark计算一个，分stage（layer）计算完成，再加载下一层的模型推理。当然，实现上是完全可行的。但与spark task相比，加载模型的权重对比加载next stage这样的计算图程序，需要的内存加载大得多，中间结果（activations, optimizer states）也大得多。
 
 

TP: Tensor Parallelism via Megatron-lm

纵向扩展面临的几个问题：(1) 纵向切分网络需要做“规划” (2) 切分的均衡性有一定考验，且最坏情况时，单层也许超过单卡的内存，使得进行。
Megatron-lm提出的是横向扩展的方案，即我们应当如何横向切分权重矩阵为等价的若干个小矩阵，什么时候做聚合数据等问题。
notion image
如图所述，(a) 描述的是一个两层MLP + GeLU + Dropout，应当如何切分矩阵。第一层权重A按列切分，每台设备获取全量的X与sharding分片的矩阵做乘法。分片的乘法结果，与第二层权重B，按行切分后继续相乘，随后聚合结果。至于为什么A和B两层的权重切分方式不同，以及三层MLP应该怎么做，在另一篇文章中有说明。blog.zjfwork.cc
blog.zjfwork.cc

blog.zjfwork.cc

notion image
（b）描述了在原始transformer self-attention中是怎么横向扩展的。考虑self-attention，有head_size的概念，head_size影响softmax的范围。如果按naive的思路，将attention head横向切分，则会导致softmax的计算需要跨节点，v的计算也需要跨节点。因此合理的策略是，将同一个attention head放在一台机器上，使其有能力完成一个local的self-attention计算。这也是横向扩展，只是sharding by的维度不同，(a)中描述的mlp是sharding by column or row，(b)中描述的是将attention head视为一个整体，sharding by attention head。
 
以上方法称为tensor parallelism，是model parallelism的一种。pipeline parallelism也是model parallelism的一种。

DP， DDP: Data Parallelism and Distributed Data Parallelism

DP，如《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》一文中所说：
DP的思想很直观，一个模型的版本训练速度是确定的，那么理论上多部署一些模型就可以加速。但加速不是免费的，是边际效用递减的，因为最终若干个模型replica之间需要通信来保持参数的一致。
思考几个场景：（1）单机单卡（2）单机多卡（3）多机多卡。前两者在本机内通信，后者通过网络通信，需要考虑的网络开销不是一个数量级。DDP主要的工作是在确保分布式正确性的前提下，加速梯度通信。
 
pytorch DDP主要通过NVDIA的NCCL提供的”collective communication primitives”实现多模型实例之间的梯度同步。文中主要提及有以下几个特性
（1）通过NCCL的ring all-reduce tree-all reduce， 从通信模型上减少了通信压力。相比master-slave结构的，更类似p2p的通信协议。
（2）bucketing gradients
通过实验证明，对于tensor数据的通信，nccl等通信库，少量多次 vs 大量少次，批量的总体性能表现更好。因此通信时机策略上会偏好梯度计算的过程中，累积足够的数据后再通信。
notion image
（3）overlapping communication with computation
文中论述了在GPU/CPU上完整计算一次反向传播的消耗时间，与all reduce一次梯度的时间在同一个数量级，意味着可以设计一种并行策略，将梯度ready的单位划分成更小的粒度，在计算梯度时，已经全部ready的梯度单位则可以先行开启all-reduce。这一部分的时间收益几乎是免费的。
（4）skipping synchronization
是一种防止死锁的策略，与性能无关

ZeRO:

ZeRO paper中声明了一个关于Megatron-lm的性能实验，以表述了一个问题，Tensor Parallelism的方法在单机多卡中可行，而多机多卡需要网络通信时，则硬件利用率很低。

ZeRO的insights

简单类比CPU分布式系统里的存算一体或存算分离，DP和MP都属于存算一体的特征；参数所在的机器就是负责其前向计算，后向计算和参数更新的机器。
（1）单纯DP中，参数被复制N份，每一完整复制品所在的机器负责一部分的数据，进行前向计算，后向计算，all-reduce后更新参数。单纯MP中，参数被切割成M份，每一份独特的分片所在的机器，负责所有数据对当前分片矩阵的前向计算，后向计算，结果同步与更新参数。
（2）DP的存算一体，带来了成倍的固定显存消耗，同时也限制了模型能塞进单卡的显存尺寸 ；MP的存算一体，虽然解除了显卡尺寸的限制，但MP alone的训练策略，无法对数据并行训练；DP + MP的模式，模型尺寸的限制因素从单卡变成了集群的显存总和，但要除以DP的并行度，因为集群的整体显存随着DP并行度线性增长。
（3）从存算分离的角度，模型/优化器状态/梯度这三者主要的显存占用（存），从全局而言，只需要一份。而计算则多多益善。这样，存算分离后，只要整个集群的整体显存塞得下单份模型的模型/优化器状态/梯度这三者，理论上就能达到了当前总显存大小下，能训练的最大模型（不考虑offload策略）。
（4）那么代价是什么，代价就是通信。ZeRO需要论证，在存算分离之后带来的计算/通信过程的改变，会让“算”与“存”的交互的额外时间开销不仅不会增加总体物理训练时间，反而会成倍加速。

ZeRO的方法与结果

1. ZeRO-DP

DP是很好的在充分利用Flops的情况下加速训练的手段，DP的问题就是memory的利用效率与单卡的限制。通过显存占用的分析，将其分为三个对象与三个ZeRO的stage:
（1） Optimizer States: 如Adam中的Momentum与Variance；而文中讨论Mixed-Precision设置下，优化器需要存储三个值，fp32版本的参数，M与V。因此在文中K = 4 * 3 = 12。Stage1是将这一部分分布在集群的所有GPU上，即每个分片存储 。 是数据并行度，在单纯ZeRO-DP里可以认为是GPU的数量。 （2）stage2/stage3继续扩展，将fp16版本的“负责存储”梯度与参数的状态分片。
（3）在ZeRO-DP里，分片的是“状态”，是“存储”，计算则是每块GPU都会计算完整的模型过程，因此 就是GPU的数量，就是一个mini-batch分多少个micro-batch的数量。
notion image
实验表明，在单机内，ZeRO与Megatron-lm的方法差距并不十分大（实现ZeRO的复杂度比Megatron-lm大很多）。当TP方法需要跨节点时，ZeRO在计算利用率上有数倍的提升。
notion image
notion image

2. ZeRO-R

除优化器外，剩余的内存占用问题主要是(1) 激活值的存储（activations） (2) 算子的中间值等（temporary buffers）以及由此引发的external memory fragmentation（如Partitioned Attention中讨论的问题）。
 
（1）激活值
激活值的理论集群显存占用在分布式训练中跟三个要素有关，①sequence length ②模型大小，层数与hidden size ③batch_size。前两者取决于模型训练的需求，后者主要影响训练速度。在单纯ZeRO-DP中，一份计算对应一份GPU，因此一份GPU对应一份micro-batch的激活值显存占用，并没有什么sharding的空间。因此ZeRO-R中需要引入MP的概念，在基于ZeRO-DP的一个计算组内（一整份的计算和一个分片的“状态”的逻辑单位），再次横向拆分。接着一整份计算对应的激活值，就可以被MP degree降低占用。举个例子，同样16卡的双机集群，可以设置DP=16，也可以设置DP=8, MP=2或DP=2, MP=8，取决于内存与通信占用分布与训练速度之间的取舍。
同时，在一个DP组内，activation也不会layer by layer存储，而是使用activation checkpoint策略，采样sqrt(L) 层存储，其他层在需要时recompute，以此减少模型配置带来的显存增长。
（2）temporary buffers与external memory fragmentation
文中没有细说，大意是通过连续内存管理等工程方式降低了浪费。
 

ZeRO的通信量分析

设模型参数量为
  1. ZeRO-DP
    1. 以单纯DP为baseline， 单纯DP在计算出local的梯度后，对分布式的梯度做all-reduce。ring all-reduce为其通信量最佳实现，分为ring reduce-scatter + all gather两阶段。两阶段分别的单机通信量为 ，和为 。具体过程说明可见https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/
    2. 与 ，其过程均为先对partition的元素做reduce-scatter，后进行all gather，因此通信量是一致的。
    3. 对于 ，额外的过程是需要拉取一个 的参数量，因此通信量是
  1. ZeRO-R
    1. 以单纯MP为baseline，对比activation partition带来的额外通信量
    2. 结论是10%左右，计算需要考虑多种因素较为复杂。
 
 
 
 

FSDP： Fully Sharded Data Parallel

https://arxiv.org/abs/2304.11277 fsdp是zero stage3在pytorch中的一个实现 The FSDP algorithm is motivated by the ZeroRedundancyOptimizer technique from DeepSpeed but with a revised design and implementation that is aligned with the other components of PyTorch.
 

1. FSDP Unit

ZeRO paper中并没有提到，应该如何partition states，以什么方式partition。 Naive地，我想有几种可能性（1）对参数全局编号，分段partition；显然这不是很好的方式，会造成forward时layer by layer加载时的分片网络热点。（2）全局编号，hash partition；分散了layer by layer的热点通讯。但在元数据管理上显得比较单薄，比如在加载策略上，我们在无法做更多的工作，按矩阵加载 - 计算 - 释放的方式显然会容易存在computation和communication的bubble。其他还要考虑self-attention/MLP等计算内容上的不同。
如果我们能在使用时表达一种元数据，表明A和B两个参数的单位是一组，我希望这一组同时拉取/同时释放，帮助训练引擎“掌握”更多的对训练中模型的理解。 FSDP提供了一个新的概念，unit，同一unit中的参数拥有相同的生命周期。这样我们可以将较小的layer的通信fusion在一起，从而摊销通信的overhead。
notion image
 

2. Hybrid Sharding

类比分布式数据库，ZeRO只提供了一种sharding=N, replica=1的策略；而FSDP提出了Hybrid Sharding的概念，sharding states允许存在replica，以适配hardware的locality。提供了一种在模型size扩展和通信量增加之间中间参数调节方案。即假设模型的大小没有大到超过一台物理机的显存size，但又希望有多机并行训练，显然通过全局的ZeRO satge3 shard会增加网络延迟。因此，在计算上依然是以单卡做一个DP，在存储分片上，一台物理机保留一个完整副本，物理机之间为replica；可以将参数通信保持在单个硬件内完成。
notion image
 
Prev
较难无痛的各种PO眼保健操：VPG TRPO PPO DPO GRPO
Next
Megatron-LM中GEMM的推导
Loading...