模型参数scale模型参数以指数级增长注意这里横轴时间是线性的纵轴模型大小是对数坐标轴。单卡很容易装不下大模型解决方式主流的解决方式一般是多卡互联。如图PLX是PCIE的管理芯片也就是GPU-CPU通信的负责单元。HCA是Infiniband的配套网卡也就是多节点通信的负责单元。NVSwitch连接NVLink是单节点内多卡通信的负责单元这里比较少见的是单节点两个CPU之间也有直接互联CPU通信不需要通过内存。多卡互联允许我们把一个内存计算量都超大的模型拆分到多卡并行推理这几乎是当前运行超大参数模型至少几百B最多10T参数量的唯一方法。集合通信当然不是有多卡模型就自动并行了还需要我们规划并行策略多卡之间需要通信这需要复杂的策略设计。其中多卡通信的基本操作就是上节讲过的几种集合通信。这里比较重要的一个算法是我们可以把all-reduce拆成reduce-scatterall-gather这也是之前讲过的这里给个具体演示图。可以看到reduce-scatter的结果的shape恰好就是all-reduce的输入shapeTPU的通信策略这里讲点题外话对比一下TPU和GPU的互联策略。下面的GPU策略就是我们熟悉的利用NVSwitch的全连接策略任何两个GPU的通信最多经过两跳。上面的TPU通信则采用了mesh网格/环面策略每个TPU只和上下左右四个TPU直连两个TPU间的通信有的很近有的很远对于一个(n,n)的二维网格最远的两个TPU间需要O(n)O(n)O(n)跳这样设计的目的是什么看起来没有什么优点实际上并非如此TPU这样设计的优点是结构简单不需要交换机且对于特定的并行策略不一定慢比如对张量并行大部分通信就是发生在局部的只需要局部之间的TPU交换mesh的延迟也可以做到很低。并且mesh的通信模式是可预测的可以基于这个模式设计并行通信策略。而对于通信流量不可预测有很多突发点对点通信时如专家并行GPU的借助交换机的all-to-all互联才会更好。也就是每个通信拓扑设计都有他适合的通信模式没有哪个模式可以适合所有情况需要根据通信策略选择合适的拓扑。当然面对MOE这种稀疏点对点突发流量多的任务TPU也做出了一些调整比如上图每四个TPU组成一个小单元小单元内全连接一跳就能到达其他TPU。8个小单元之间也全连接。这样实际上最多两跳就能从一个TPU到达任意一个其他TPU。4和8的选择是有取舍的更大当然能让集群有更多的TPU但是需要的直连网线数是随着TPU规模平方量级增长的4和8的平方还可以接受再大互联成本就过高了对于更大规模的互联把上面的8个小单元组成的大单元做一个mesh网络(6,6)一共36个大单元一共1152张TPU组成一个集群。另外一种策略是全面学习GPU互联把上面8个小单元组成的大单元通过两层路由器互联类似经典的Spine-Leaf架构可以实现完全无阻塞。另外也可以采用光互联也就是右上角的Optical Circuit Switch光路交换机主流超节点为什么不能无限扩大集群规模以NV72和昇腾384现在的两种主流超节点为例来看。NV72的单卡性能更高昇腾为了让超节点算力反超需要堆更多卡384/72。但代价是互联的开销也随之快速增长为了实现384超节点的互联需要几千个光模块几千条光纤。再大的话成本就要不可接受了所以考虑到性价比一个节点的互联规模是有限的再大不如拆成多个节点了。DP数据并行接下来循序渐进仔细讨论一下DP的策略。首先是最朴素的情况假设我们采用的是一个SGD随机梯度下降每轮需要把全部梯度求和乘上一个学习率累加到模型参数上。数据并行的思想是训练数据量很大但不同数据之间是没有依赖的只有最后需要把全部数据产生的梯度求和一次中间不同数据的计算可以交给不同GPU并行计算。那么朴素实现就是每个GPU保存完整的模型参数然后把训练数据分发给各个GPU并行计算最后再来一次通信最后汇总一次模型梯度乘上学习率算出梯度更新量再分发给每个GPU使每个GPU保存的模型参数都得到更新。这样对于计算的拓展是有效的M卡的话整体计算延迟降低M倍。代价是每一轮需要2倍模型参数的通信量梯度求和梯度更新。对于内存拓展则是没有效果每个GPU都要保存一份完整的模型参数实际上还要保存优化器的参数。如上的策略虽然简单但是分析方法是有一般性的对于一个并行策略我们关心他对内存和计算的拓展以及产生的额外通信代价。这个做法的主要问题就是每个GPU必须保存一份完整模型参数。我们前面一如多卡并行一个主要目的就是模型参数量太大一张卡装不下跑不起来。这里完全没优化这一点这肯定是不行的实际上对于实际训练场景优化器状态也不小比如Adam优化器需要保存梯度的一阶二阶动量还有累加的梯度值。加上原始梯度模型参数。实际上需要模型参数量大概5倍的显存这非常糟糕意味着我们只能训练很小的模型ZeRO-1全称是Zero Redundancy Optimizer也就是零冗余优化器是微软在DeepSpeed里提出的一个显存优化技术仔理想情况下可以在数据并行时让全局所有数据都只存储一份没有多余的存储。分三个版本三个版本循序渐进切分了优化器状态梯度模型参数。第一个版本只切分了优化器状态这是显存的大头锁优化这一步就能有很好的效果。具体实现是每个卡计算自己的数据切片得到的完整梯度。对梯度做个reduce-scatter每张卡的到某一部分参数的梯度和也就是只维护一部分参数的优化器状态每张卡利用自己的优化器更新一部分参数最后all-gather合并得到完整参数分发到每张卡上这样所有GPU都得到了利用这一轮数据反向传播后的完整模型参数。对比朴素的数据并行和一阶段的ZeRO通信量并没有增加仍然是2倍模型参数的量级。但使用的内存中有一部分优化器部分可以再不同卡上分布存储使用的显存大幅降低。也就是我们在没有额外代价的前提下优化了显存使用是个纯正面的优化。ZeRO-2二阶段把梯度也分片了注意这里的切片都是横切也就是是模型每一层的参数都切而不是切出来某些层。实现思路是每计算出一层的梯度reduce-scatter到不同rank每个rank保存这一层的一部分梯度的和。然后各个rank就可以把这一层的梯度都释放了这和前面ZeRO-1的区别就在于每算出一层的梯度就reduce-scatter一次而不是等全部梯度都算完了再做。这样每个rank都不需要保存完整梯度只需要保存自己在维护的那部分梯度切片的和这样梯度部分显存也没有冗余了最后和ZeRO-1一样用梯度切片和优化器切片更新自己负责的一部分参数切片在all-gather合并并同步到所有rank这个版本实际上也没有增加通信量只是把版本1最后才做的reduce-scatter每一层的数据的得到了就马上做这类似于IO里流式IO和非流式IO流式IO不需要像非流式那样维护一个超大缓冲区节约了内存。对于通信量两种是一样的但ZeRO-2增加了一些同步点这会导致额外延迟也就是时间换空间ZeRO关注的就是显存优化所以这是可以的。说点题外话对于更关注时间而不关注空间的情况应该反过来做比如算法竞赛和CPP编程里可以用fread函数指定一个大缓冲区一般至少1MB等这个缓冲区满了才进行一次IO这样可以减少IO系统调用的次数优化延迟但缺点是原本流式输出几乎O(1)O(1)O(1)的额外空间现在变成O(n)O(n)O(n)了。不过这一般没问题算法竞赛里空间都很松O(n)O(n)O(n)的空间一般都可以接受。ZeRO-3最后一步模型参数也可以分片存储并不需要每个rank保存完整的模型参数这一步后显存几乎只最开始的零头了(1.9/120)核心思想是类似ZeRO-2的对模型分层的思路每一层的前向传播/反向传播时只需要这一层的模型参数那么不需要每个rank的随时存着全部的模型参数可以让每个模型都只存模型参数的一个切片计算某一层的梯度时临时做个all-gather把这一层的模型参数合并送到每个rank这一层结束后把这层的参数再释放。这样每一层始终保存的参数只有他负责的那部分切片。这个操作的代价是增加了通信量以及更多的同步点每层的前向传播反向传播都需要一次all-gather。整体大约有增加了两倍模型参数量的通信量FSDPZeRO-3是一个算法实现他的一个具体实现是torch里的FSDPFully‑Sharded Data Parallel完全分片数据并行核心思想是增量式计算/通信。这就是我们在把梯度和参数做切片时的主要思想只有在计算这一层的梯度时才临时all-gather拼出这一层的模型参数这一层算完了就释放通算重叠这在前面没提到前面默认通信计算都是串行的。只有从流水线优化的角度才会分析到这一点。现代GPU的通信和计算一般都是不同单元负责的于是可以拆成流水线如图在0层的前向传播时也就是FWD0块的同时可以同时在负责通信的另一个CUDA Stream里做AG1也就是all-gather汇聚1层的参数这两部分使用不同单元无依赖。可以完全重叠。只有FWD1需要等待AG1结束这里有数据依赖需要同步。总结ZeRO三个版本都是指增加了线性的通信量在通信带宽不太小的情况下这不会带来太多的通信延迟但把显存取得了大幅优化这是值得的。实测在8卡A100集群上能跑的最大模型参数量如上图。其中stage3对比baseline显存可以严格除以GPU数也就是/8相应的能跑的最大参数量也变大了8倍。