国外购物网站建设,黄骅市医院,广东省建设合同备案网站,网站文章模板深度学习相较传统机器学习模型#xff0c;对算力有更高的要求。尤其是随着深度学习的飞速发展#xff0c;模型体量也不断增长。于是#xff0c;前几年#xff0c;我们看到了芯片行业的百家争鸣和性能指标的快速提升。正当大家觉得算力问题已经得到较大程度的缓解时#xf…深度学习相较传统机器学习模型对算力有更高的要求。尤其是随着深度学习的飞速发展模型体量也不断增长。于是前几年我们看到了芯片行业的百家争鸣和性能指标的快速提升。正当大家觉得算力问题已经得到较大程度的缓解时大语言模型LLM, Large language model的兴起又带来了前所未有的挑战。当网络模型达到一定量级后比如参数量达到10B级别表现出In-context learningInstruction following和Step-by-step reasoning等涌现能力Emergent abilities。这些能力是以往模型所不具备的因此LLM引起了学术界与工业界的浓厚兴趣以及广泛关注。但是得到这些能力的代价是模型的急剧膨胀。比较知名的如GPT-21.5BMegatron-LM8.3BT511BTuring-NLG17BLLaMA65BGPT-3175BJurassic-1178BGopher280BPaLM540BSwitch Transformer1.6T参数量很快从B时代到达了T时代。如今在这些“庞然大物”面前那些历史上曾经的“巨无霸”像GPT110M、BERT340M看起来都显得有些迷你了。根据论文《Machine Learning Model Sizes and the Parameter Gap》中提到的从1950到2018年语言模型的模型大小增长7个数量级但从2018年到2022年4年就暴涨5个数量级。网络模型的陡然增大随之而来的是一系列数据集、训练方法、计算框架相关的问题。关于这些年大模型发展的回顾以及对这些问题业界的探索的阐述可参考论文《A Survey of Large Language Models》。LLM所涉及的东西非常多这篇文章主要讨论的限于对计算框架带来的挑战及如何用分布式的方法解决。
从网络模型结构方面好消息是自从2017年Google在《Attention Is All You Need》中提出Transformer后近几年模型的结构开始趋同于它。它不仅在于效果好而且相比之前的如RNN结构更便于计算。Transformer网络中的主要负载是MHA和FF而它们本质上都是或大部分都是GEMM。而GEMM的计算是相对成熟的。这就将前沿的AI与成熟的计算理论与实践结合起来了。另一方面Transformer是相同模块的重复堆叠这也给并行计算如pipelining提供了便利。因此近几年主要的LLM也都是基于Transformer。
那模型变大后对计算框架会带来什么样的影响呢最直接的就是memory不够用了。由于传统CPU对神经网络这种调度并行的workload效率不高因此目前主流的都是用GPU或者各种加速器。但是这些芯片的memory比CPU的memory要小很多。
在模型的训练过程中通常需要参数量数倍的memory。其中消耗memory的大户一般有这么几个
Weight模型的权重参数Activation中间层的输入与输出Gradient Optimizer state梯度及优化器状态如Adam中的momentum和variance
其它还有些如temporary buffer或者fragmented memory等由于占得不多因此优化时放的精力也少。推理的话optimizer state就不用了但weight与activation是训练与推理时都逃不了的。前面提到模型的size在短短数年增长了几个数量级但相比之下GPU的memory增长大概只有几倍。这让本就不太够用的GPU memory显得更加捉襟见肘。论文《ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning》中分析了训练过程中的memory需求。它将weightgradient和optimizer state归为model state。混合精度和使用Adam优化器情况下单个参数应对model state需要用到20 bytes。这样对于大模型的memory需求远远超过了单卡GPUA100的话是80G的承受范围。对于100 B参数量的模型光model state就需要64 GPU如果是1T参数量的模型则需要512 GPU。论文《ZeRO: Memory Optimization Towards Training A Trillion Parameter Models》中提到参数量1.5 B的GPT-2光是model state就需要至少24 G。另外activation需要8G左右。这意味着一张V100基本就撑爆了。论文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》中指出530 B的模型需要超过10 T的memory用于model state。Activation则会占用大约16.9 T通过gradient accumulation优化可减少到8.8 G。此外为了减少计算量使用KV cache的话也会增加不少的memory开销。这使得大模型计算的memory问题变得非常严重很多时候它直接决定模型是否能跑得起来。
业界从多个角度提出方案来缓解memory的问题如
Activation checkpointing / rematerizationHeterogeneous Memory ManagementCompressionMixed precision trainingMemory efficient optimizerLive analysis Reuse
本文主要关于分布式的方法因此对上面这些点不作展开。减少内存使用一定程度上解决“跑不起来”的问题。跑起来后就需要考虑“跑得更快”的问题。毕竟大家都在折腾大模型发展非常迅速能早些得到结果就意味着领先。一个好的想法如果实验它需要几年那等结果出来黄花菜都凉了。这就要求框架在计算大型AI模型时还需要提高硬件效率最终提升训练效率。
进入正题前先过一些相关概念。
为了更快得到结果直观的想法是会投入更多的资源来扩展系统而系统的扩展可分为两种方式scale up和Scale out。通俗点说前者是为单个计算节点添加更多的资源如CPUmemory等后者是添加更多的计算节点。理想情况下我们希望吞吐是随着资源的投入线性增长的虽然现实情况中由于同步、通信开销等原因总会有些开销。要达到近线性的加速就需要算法与系统是易被并行的而且硬件资源的使用平衡。
在分布式计算下这些可并行的任务就可以放到多卡或多机上。这就涉及到通信问题。否则额外计算资源带来的性能收益可能还补不上通信开销。硬件上可以通过NVLinkPCIe等可参见论文《Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect》。多机可以通过EthernetGDRGPU Direct RDMA等机制。为了使用更方便Nvidia提供了NCCL通信库利用多种通信机制向上提供诸如AllReduce、AllGather、ReduceScatter等常见的collective primitive。其它类似的通信库还有OpenMPI、Gloo。
在讲分布式并行前先提下一些并行相关概念。首先workload得并行起来才能分布到多卡或者多机上。并行是分布的基础。关于在深度学习领域的并行可以参见2018年的综述文章《Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis》。在此不展开。
关于并行带来的加速影响涉及两个重要定律即Amdahl和Gustafson定律。Amdahl定律可能是计算机领域最广为人知的定律之一它定量地给出了提升系统的部分性能对整个系统性能的影响。它有一个假设就是问题的规模是固定的。论文《Reevaluating Amdahl’s Law》中指出现实中人们一般是为了解决更大规模问题而增加更多资源。Gustafson定律以此为假设给出了增加计算资源对系统性能的影响。这两个定律从两个不同角度阐述了加速比处理器个数与串行比例之间的关系。前者强调了要降低程度的串行化比例否则光摞计算资源性价比不高。后者强调当并行化部分足够大加速比是可以随着计算资源线性增长的。
对于大模型的计算目前流行的分布式并行方法可以大体分为几类
Data parallelism
AI模型在训练时一般都会将多个样本数据组成mini-batch一起计算。DP利用了这一点将mini-batch进一步分成多个更小的mini-batch再分给多个worker去处理。在这种模式下每个worker都有完整的模型包括weightoptimizer state等。每个worker都可以独立完成整个模型的梯度计算它们计算出来的梯度最终需要聚合到一起。
其实对于大模型的分布式计算框架的研究其实从很早就开始了。2012年Google的论文《Large Scale Distributed Deep Networks》中提到的DistBelief也就是TensorFlow的前身采用了DP和MP等并行技术以及改进的优化器可以用于在集群中训练参数超过1B的模型。只是由于当时的局限性它用的是CPU集群而且是基于parameter server的中心式架构。导致机器数一多后PS成为通信瓶颈影响scalability。因此后来越来越多使用AllReduce方法因为可以充分利用所有device的带宽。可参见Massively Scale Your Deep Learning Training with NCCL 2.4。
按同步方式还可以分为BSPBulk synchronous parallel与ASPAsynchronous parallel两种。前者会有同步开销后者没有同步开销但无法保证收敛。因此还有一种SSPStale synchronous parallel方式是这两者的折中。即每个worker允许使用stale的weight处理下一个mini-batch。这个staleness的阀值是可以定义的。
2017年论文《Integrated Model and Data Parallelism in Training Neural Networks》中将DP分为两种一种是Batch parallelism即平时常见的那种另一种是Domain parallelism。它将输入图片切分成多块交给多个worker处理。这种方式中每个worker与batch parallelism一样也包含完整的模型参数但只处理图像的一部分。它在视觉网络前几层activation比较大的情况下有比较好的通信复杂度。
DP由于实现简单在各种主流AI框架中基本都有支持。如PyTorch的DDPDistributedDataParallel可参见论文《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》。TensorFlow中的tf.distribute.Strategy可参见Distributed training with TensorFlow和Multi-GPU and distributed training。MXNet中可以参考Run MXNet on Multiple CPU/GPUs with Data Parallelism。
最基础的DP的优点除了实现简单外它的计算与通信效率也很高。但是它的缺点也很明显一是当worker数增多时mini-batch也会随之变大而过大的mini-batch会让训练不稳定甚至不收敛。此外它最大的问题还在于它无法解决单卡放不下大模型的问题。
为了解决DP无法减少memory footprint的问题2019年的论文《ZeRO: Memory Optimization Towards Training A Trillion Parameter Models》提出了ZeROZero Redundancy Optimizer。ZeRO-DP用ZeRO扩展了DP在达到DP的计算与通信效率的同时还能达到MP的内存效率。它通过将model stateparameter, gradient与optimizer state在data parallel processes间进行切分从而消除memory redundancy同时它使用dynamic communication schedule利用model state时间上的使用特性减少通信量。根据切分model state的程度不同ZeRO可分为三级或三个stage。即ZeRO-1切分optimizer stateZeRO-2切分optimizer state与gradientZeRO-3切分optimizer stategradient和parameter。第三级可以支持在1024张Nvidia GPU上训练T级别参数的模型。
而后论文《ZeRO-Offload: Democratizing Billion-Scale Model Training》中基于ZeRO提出ZeRO-Offload改进ZeRO-DP。它将optimizer state放在CPU memory中同时每个process将自己的那份梯度也放到CPU memory中。ADAM优化器的更新在CPU上完成更新完了再放回GPU。再使用GPU上的all-gather收集所有更新好的参数。论文《ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning》中提出ZeRO-Infinity利用了NVMe memory进一步改进了ZeRO-Offload。这些属于异构内存管理的范畴不是讨论重点就不展开了。
它有点类似于后面要提到的TP因为它会切tensor分到多个worker上。但它与TP的显著差别是TP是切完后在worker内部算完再聚合计算结果因此它需要对代码有修改。而ZeRO只是利用多个worker来分布式地存tensor在算之前会先聚合它们算完了再把除了自己要维护的部分丢掉从而减少memory占用。总得来说MP会传输activation而ZeRO不会。MP处理的切分的输入而ZeRO处理的完整的输入。这样的话基本不需要修改代码。有些地方也将它区别于DP和TP单独作为一类来介绍。这里为了乘法同时因为论文中称为ZeRO-DP因此放到DP这一节来一并提了。
ZeRO实现于Microsoft自家的框架DeepSpeed中。由于它的有效性其思想也被其它的框架所吸收。比如PyTorch中的FSDPFully Sharded Data Parallel就是受ZeRO与FairScale的Fully Sharded Data Parallel的影响。详细可以参见Introducing PyTorch Fully Sharded Data Parallel (FSDP) API和GETTING STARTED WITH FULLY SHARDED DATA PARALLEL(FSDP)。
在DP中由于每个data-parallel process都会有完整的模型它们需要一致地被更新。这个weight的更新操作需要在所有的模型replica上进行。这是一个冗余操作。2020年的论文《Automatic Cross-Replica Sharding of Weight Updates in Data-Parallel Training》针对这个问题提出了shareded DP。该方法在XLA框架上实现用于Cloud TPU。它利用all-reduce操作可分为reduce-scatter与all-gather两个阶段的特点使用静态分析与变换使多个replica共同分担weight update计算。
关于DP的优化也有不少相关工作如论文《Preemptive All-reduce Scheduling for Expediting Distributed DNN Training》中的PACE是一个communication scheduler用于抢占式调度all-reduce操作从而最大化计算与通信的overlapping论文《Supporting Very Large Models using Automatic Dataflow Graph Partitioning》中的ByteScheduler也是用于提升了DP中的overlapping ratio。
Model parallelism
MP是一类将model state在多个device上切分的方法。它的优点是可以让多卡分担memory占用从而解决模型放不下单卡的问题。但缺点是要使切分的工作量平衡不太容易计算间有依赖activation需要传输导致通信量增加。
对于一个模型切法有很多。有些是比较通用的有些适用于特定网络。因此MP也可细分为好几类
Tensor parallelism
又称为operator parallelismoperator-level parallelism或horizontal MP。它将单个算子在多个device间切分。通常用于单个算子的weight就占用大量GPU memory的情况如embedding layer。
最早在论文《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》中使用了TP来并行化连续的FF层Transformer中的self-attention与feed-forward。它利用矩阵的分块相乘性质将weight切分从而将相乘的计算量分到多个worker上。由于目前流行的Transformer网络其主要计算量为矩阵相乘。TP应用于GEMM就可以看作并行矩阵相乘Parallel GEMM问题。并行矩阵相乘按并行的维度可分为
1D上面提到的Megatron中将weight在单个维度上切分是1D的。1.5D如论文《Communication-Avoiding Parallel Sparse-Dense Matrix-Matrix Multiplication》2D如经典的Cannon’s algorithm和SUMMAScalable Universal Matrix Multiplication Algorithm。还有论文《An Efficient 2D Method for Training Super-Large Deep Learning Models》中的Optimus使用SUMMA来提升内存与通信效率。2.5D如论文《Tesseract: Parallelize the Tensor Parallelism Efficiently》是基于SUMMA与2.5D Matrix Multiplication。论文《Communication-Optimal Parallel 2.5D Matrix Multiplication and LU Factorization Algorithms》提升了Cannon’s algorithm的效率。3D原始3D并行矩阵乘法可参见论文《A three-dimensional approach to parallel matrix multiplication. IBM Journal of Research and Development》。论文《Maximizing Parallelism in Distributed Training for Huge Neural Networks》做了改进消除了原算法中广播数据的冗余。论文《Communication-Optimal Parallel Recursive Rectangular Matrix Multiplication》提出CARMA是一种达到asymptotic lower bound的递归算法。论文《Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication》提出COSMACommunication Optimal S-partition-based Matrix multiplication Algorithm使用red-blue pebble游戏来建模数据依赖并推出I/O最优schedule。
Colossal-AI中集成了多种并行矩阵相乘方法有兴趣可以参考。
论文《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》中提出的GShard模块包含一组轻量级的annotation API和XLA compiler的扩展。它采用automatic sharding的方法用于有600 B参数的带Sparsely-Gated Mixture-of-Experts的transformer网络。
基于GShard论文《GSPMD: General and Scalable Parallelization for ML Computation Graphs》中的GSPMD使用了tensor sharding annotations来以一种统一的方式达到不同的并行DP, in-layer MP和spatial partitioning等。切分的方式很灵活不同层可以用不同的方法切分同一层可以使用不同方法切分。GSPMD泛化了其后端GShard以partial tiling和manual subgroups扩展了sharding representation实现了像priority sharding propagation, recursive partitioning和pipepiling等特性。该方法可用于上千个TPUv3 cores上的视觉、语音和语言模型。
Pipeline parallelism
又称为vertical MP。与TP切分单个算子不同PP是将整个网络在层间进行切分。这样整个计算图就分为几个子图。对于有依赖的子图计算时就像工厂的流水线一样没有依赖的部分可以在多device上并行执行。PP与TP相比通信量更小。在TP中每个Transformer layer需要两次activation上的all-reduce操作。而对于PP只需要点到点的通信。因此PP需要更少的通信但它的问题是会有bubble和weight staleness。因此两者是正交互补的关系。
PP早期比较经典的工作大致开始于2018年包括Microsoft论文《PipeDream: Fast and Efficient Pipeline Parallel DNN Training》中提出的PipeDream和Google论文《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》中提出的GPipe。
PipeDream以一种通用和自动的方式结合PP, MP和DP。PipeDream将DNN切成多个Stage。为了将这个过程自动化它通过profiler估计每个layer的计算时间与output size。基于这些估计optimizer对之进行切分。切分算法的目标是最小化总训练时间等价于最小化最慢的stage的时间。该问题具有optimal sub-problem property即给定机器数量最大化吞吐的pipeline包含了最少机器数量时最大化吞吐的sub-pipeline因此可以用DP求解。Work scheduling使用1F1Bone-forward-one-backward来避免GPU的idle时间。将forward和backward的执行交替这样就不需要等到所有的forward完成再执行backward从而提升硬件利用率。它采用异步backward update的pipeline中每个mini-batch的forward和backward pass使用不同的参数即weight staleness问题这可能导致训练无法收敛。另外不同的stage可能面临不同程度的staleness。这也会导致收敛性受影响。PipeDream使用了两种技术一是Weight Stashing维护多份weight每个active mini-batch一份二是Vertical Sync用于估计potential inconsistency across stages。
与PipeDream不同GPipe使用的是同步SGD。因此它会产生bubble。为了减少bubble的开销GPipe采用batch splitting技术将mini-batch切分为更小的micro-batches然后对其用pipeline。每个micro-batch称为一个chunk。同步发生在mini-batch结束时。一个mini-batch中的micro-batches的梯度进行累加后在mini-batch的结尾一次应用。这个同步产生的bubble当device越多越严重。当M 4 x KM为micro-batchesK为accelerator数量时bubble overhead可以忽略。另外由于它需要将中间activation的结果存下来因此内存需求很大。
简单的pipelining会因为inconsistent weight的问题影响模型收敛性。为了解决这种问题需要在memory footprint与throughput间做trade-off。以前面的两种方法为例它维护一份weight并周期性地进行pipeline flush。而PipeDream采用的是weight stashing策略。它没有周期性的pipeline flush代价是需要维护多份weight。前者降低了throughput后者增加了memory footprint。
2020年的论文《Memory-Efficient Pipeline-Parallel DNN Training》提出的PipeDream-2BW2BW是double-buffered weight updates的缩写及使内存更加高效。它使用weight gradient coalescing strategy和double buffering of weights技术。该方法利用了每个输入产生的梯度不需要马上应用因此可以被累加从而减少需要维护的weight。它将1F1B与2BW scheme结合在每个wokrer上最多stash两份weight。避免pipeline flush同时降低memory footprint。梯度以micro-batches的粒度来计算。多个micro-batches的更新会累加然后以batch为粒度来应用梯度。PipeDream-2BW结合了PP与DP它将模型分为多个stage放到多个worker中用PP然后将每个stage复制多份用DP。PipeDream-2BW的planner基于cost model通过搜索的方法决定如何在worker间切分模型。其变体PipeDream-Flush有比2BW更低的memory footprint但吞吐也低一些。它维护一份weight引入周期性的pipeline flush来保证consistent weight。
Megatron在Scaling Language Model Training to a Trillion Parameters Using Megatron一文中提到Interleaved schedule。它不限制将连续的layer分到一个device上而是允许多个stage分到一个device上。在Interleaved 1F1B schedule中device上的stage依次执行。这样的优点是可以减少bubble但缺点是由于stage数量加倍导致通信量增加。
2019年论文《PipeMare: Asynchronous Pipeline Parallel DNN Training》提出的PipeMare是一种异步方法文中它和前面提到的同步方法PipeDream与GPipe从delay, pipeline utilization和memory usage角度做了对比。现有PP方法为了保证sync execution即forward中的wegith和backward中用于计算gradient的是同一份会引入了bubble或者存micro-batch的多份weight。这些方法会降低硬件效率。PipeMare去除了sync execution的限制。为了解决由此引入的对收敛性的影响它提出两种技术Learning rate reschedulingstep size与delay反比和Discrepancy correction调节backward中的weight。
2020年论文《DAPPLE: A Pipelined Data Parallel Approach for Training Large Models》中的DAPPLE是一个结合了DP和PP的同步的分布式训练框架。其中的Profiler以DNN模型为输入然后对每一层的执行时间activation size和parameter size进行profile。基于这些profile结果以及硬件配置等信息并行策略的planner基于动态规划算法解决partition与placementdevice mapping问题从而得到执行时间最优的DP与PP混合策略。另外它还提出一种运行时调度算法来减少device memory的占用。它的特点是可以确定和交叉的方式来执行forward与backward从而让每个pipeline任务的内存可以尽早释放。
2020年论文《Pipelined Backpropagation at Scale: Training Large Models without Batches》改进了PBPipelined Backpropagation方法。PP一般会将mini-batch切成micro-batch并组成pipeline。同步方式中需要等pipeline空了才开始下一轮mini-batch的参数更新。但这个pipeline的fill和drain的过程会严重降低硬件利用率尤其是当batch size相对pipeline stage数较小时。而PB是一种异步训练方法它不需要等pipeline清空因此避免了fill和drain带来的开销。但是它会带来inconsistent weight和stale gradient的问题导致准确率下降和训练不稳定。另一方面fine-grained pipelining将每一层作为一个pipeline stage也就是每个process处理一层被发现在一些场景下有很好的性能。但它会使得pipeline stage变多从而加重weight inconsistency问题。另外small micro-batch size能减少fine-grained PB中的gradient delay与weight inconsistency问题。针对fine-grained, small batch场景下PB的准确率退化问题本文引入SCSpike Compensation和LWPLinear Weight Prediction两个方法来解决。
2020年论文《HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism》结合PP与DP应用在异构GPU cluster上。这些cluster中有不同型号的GPU有一些可能是几代前的单卡无法进行训练。HetPipeHeterogeneous Pipeline中一组GPU称为virtual worker以pipeline方式处理minibatch。多个virtual worker采用DP。另外还提出Wave Synchronous ParallelWSP。所谓wave指在一个virtual worker中可以并行处理的一串minibatch。每个virtual worker一次性更新wave的所有minibatch而不是每个minibatch更新一次到PS。它会通过profile得到perf model。然后将partitioning问题建模成linear programming问题最后通过CPLEX求解。
2020年论文《GEMS: GPU-Enabled Memory-Aware Model-Parallelism System for Distributed DNN Training》中的GEMS是GPU上的Memory-Aware MP系统。它结合了MP和DP。为了克服现有系统资源利用率低和实现复杂的缺点它提出了两种设计MASTMemory Aware Synchronize Training和MASTERMemory Aware Synchronized Training with Enhanced Replications。通过训练额外一份模型replica来填满做完一次forward与backward后空闲的硬件资源。在1024张V100上来训练1K层的ResNet并达到相当高的scaling-efficiency。
2021年论文《Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines》提出的Chimera也是一种PP方法。它结合了bidirectional pipeline方法。其主要思想是结合不同方向的两个pipeline。与最新的同步pipeline方法相比它减少最多50%的bubble。它是一种同步方法因此不影响accuracy和convergence。与SOTA方法具有相同的peak activation memory用量且memory消耗更加平衡。
2021年论文《PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers》中的automated elastic piplining提升了Transformer分布式训练的效率。调节PP与DP的程度。Freeze training的研究表明神经网络中的参数常常是由下而上收敛的因此就可以动态分配资源使它们聚焦于某一些active layer。在训练过程中它渐进地识别并冻结一些layer将资源用于训练其它的active layer。PipeTransformer由几部分组成1) Freeze Algorithm从training loop中采样并作出freezing的决策。2) AutoPipe是一个elastic pipeline module将frozen layer剔除剩下的active layer打包塞入GPU。3) AutoDP得到AutoPipe传来的pipeline length信息它会spawn更多的pipeline replica尽可能提升DP宽度。4) AutoCache是跨pipeline的caching module。它有专门的模块在shared memory中维护cache这样所有的训练进程可以实时访问。当activation太大无法放入CPU memory它还会将之swap到disk并自动做pre-fetching。
可以看到以上方法大致可分为两步和异步两种。同步方法像GPipe, GEMS, DAPPLE, Chimera通过periodic pipeline flush保证weight的一致性从而保证了收敛性。因此更convergence-friendly但有bubble。而异步方法(如PipeDream, PipeMare)没有flush虽然吞吐会更高但会导致forward与backward的weight版本mismatch如AMPNet和PipeMare或者虽然weight版本一致但是会有weight staleness问题如PipeDream。这样虽然能减少或者消除bubble但有stale weight的问题。为了使得weight版本一致就需要保留多份weight从而使memory footprint增大。
Sequence parallelism
相比前两种SP是一种更加针对特定模型的方法。语言类的模型一般要处理或者生成token序列SP在这个维度上做并行。
2022年论文《Sequence Parallelism: Long Sequence Training from System Perspective》中讨论了sequence parallelism。Transformer模型中self-attention的memory消耗与sequence length是平方关系因此大的sequence length常常会限制大模型的运行。现有的并行方法都要求整个sequence放在一个device上。SP将activation沿着sequence维度进行切分然后分到多个device上让每个device只维护一部分sequence。为了跨device计算attention scoreSP使用了ring-style communication提出Ring Self-AttentionRSA。其实现基于PyTorch且与DP, PP和TP兼容因此称为4D parallelism。该方法和DP类似需要参数和optimizer state在所有device上拷贝因此对于大模型仍然不是很友好。
2021年论文《TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models》中的TeraPipe利用了Transformer-based model的auto-regressive性质用于single training sequence内的PP。
2022年的论文《Reducing Activation Recomputation in Large Transformer Models》提出两种技术sequence parallelism和selective activation recomputation。两者都可以将内存需求减少约一半。其中的SP结合了TP。
Expert Parallelism
在不显著增加计算量的情况下增加模型capacity一个方法是利用网络结构上的sparsity注意不要与诸如矩阵计算中的sparsity混淆。近几年比较火的一种网络范式就是Mix of Expert(MoE)。这种结构是由2017年论文《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》提出。它拥有一个很好的性质就是可以在计算代码不增加的情况下增加参数量。方法是添加sparsely activated expert block并添加一个parameterized routing function在它们前面。
2021年论文《FastMoE: A Fast Mixture-of-Expert Training System》中的FastMoE是一个开源的基于PyTorch的分布式的MoE训练系统。FastMoE支持将expoert分布到多个计算节点的多个worker上。
2023年论文《A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training》提出了DeepSpeed-TED。它结合了ZeRO的DPMegatronLM的TP和DeepSpeed-MoE的expert parallelism。其中expert Parallelism利用了expert block可以并行计算的特点将它们放到不同的GPU上。
Hybrid parallelism
2022年论文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》中从memory和compute efficiency角度比较了几种并行方法的优劣。
DP的compute efficiency高易于实现但是随着worker增加batch size也需要对应地增加但增加batch size会影响收敛。Memory efficiency由于复制了model和optimizer因此不高。梯度需要在worker间通信与聚合而梯度的通信开销随着模型增大而增大。这会让通信带宽限制compute efficiency。TP的memory footprint随着worker数量增加同比减少。每个worker的计算量减少会影响compute efficiency。PP的memory footprint随着pipeline stage同比减少但不减少activation的memory footprint。PP的通信开销最小它要求切分尽可能均衡。
实际工程中几种并行方式是可以同时使用的这样可以结合各种并行方法的优势。前面提到的很多论文中的工作其实都结合了多种并行方式比如PipeDream就结合了DP和MP。另外像论文《Integrated Model and Data Parallelism in Training Neural Networks》中也结合了DP与MP。
结合多种并行方法的一般原则是节点内使用TP节点间使用PP。再用DP复制到更多节点。比如论文《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》提到Megatron-Turing NLG 530B在280 A100集群中在node内部使用8-way TP在node间使用35-way PP。然后用DP扩展到千卡。
基础设施
为了支持灵活的分布式并行不能啥都从头手撸还需要一些软件基础设施还提高效率。业界在这方面也有一些相关探索。
2018年论文《Mesh-TensorFlow: Deep Learning for Supercomputers》中的Mesh-TensorFlow支持SPMD式的per-operator partitioning。Mesh-TF引入了一种tensor partition的表达。它将processors视作多维mesh。张量可以沿着mesh维度复制或切分。Mesh-TensorFlow引入了一种可以指定分布式TP计算的语言可以编译计算图加上collective primitive。使用的案例比如Switch Transformers。
2021年论文《CoCoNet: Co-Optimizing Computation and Communication for Distributed Machine Learning》提到目前机器学习中的计算与通信是独立的抽象由不同的库实现。而打破这个边界可以解锁更多优化。比如 Interface模型参数放在非连续buffer中而要调用collective communication如AllReduce前需要拷贝到一个单独的buffer中Fusion产生执行多个通信与计算操作的kernel从而减少memory bandwidthReorder将计算移到通信前或后这样可以分散计算或者使能fusionOverlapping多个计算和通信操作以细粒度编排充分利用计算与网络资源。
2021年论文《CoCoNet: Co-Optimizing Computation and Communication for Distributed Machine Learning》中的CoCoNETCommunication and Computation optimization for neural networks使用了嵌入于C的DSL表达包含计算与通信操作的kernel。受Halide启发CoCoNET包含scheduling语言用一组变换指定执行调度。然后由AutoTuner自动应用这些变换优化程序。最后Code generator根据程序和schedule生成高性能kernel。
2022年论文《MSCCL: Microsoft Collective Communication Library》 GC3提供面向数据的DSL用于写自定义的collective communication algorithmoptimizing compiler将之lower成可执行形式。最后基于interpreter runtime来执行。它可用于AllReduce和AllToAll操作。GC3是一种内嵌于Python的声明式语言采用chunk-oriented programming style受dataflow编译语言启发。它可以为特定的topology和input size又或者是自定义collective如AllToNext产生MPI抽象的新算法。
2022年论文《OneFlow: Redesign the Distributed Deep Learning Framework from Scratch》中的OneFlow提出了SBPsplit, broadcast, partial-value抽象包含SB和P分别对应split, broadcast, partial-value。它用于指定global tensor和对应local tensor的对应关系。这个对应关系称为SBP signature。通过对应的API用户就无需接触low-level的通信原语只需指定placement和SBP signature。基于这些compiler可以自动生成DP, MP和PP的执行计划。
自动化策略
上面讲了分布式并行的方法以及相应的基础设施那剩下的核心问题就是切分策略了。
最简单就是方式当然是手工制定切分策略。如2021年论文《GSPMD: General and Scalable Parallelization for ML Computation Graphs》中提出的GSPMD需要用户通过annotation的方式指定sharding strategy。也就是说这种方式中如何切分就需要使用者来制定。
随着网络模型和硬件越来越大越来越复杂手工配置的缺点就越来越明显了。精力花费不说模型或者运行环境一变可能就得重来。
更自动化一些的方式是针对特定的网络结构与硬件特性根据经验制定切分的规则。如2014年的论文《One weird trick for parallelizing convolutional neural networks》中讨论针对CNN的Conv与FC制定的并行策略。Conv由于计算量大但参数少适合DP而FC相比之下正好相反适合MP。2022年的论文《Efficiently Scaling Transformer Inference》讨论了TPU v43D torus topology上运行PaLM模型的一些方法。其中就针对网络中的Attention与FF结合硬件特性提出了切分的策略。
根据给定网络模型与硬件配置找最优的并行策略本质上是一个组合优化问题。对于这样一类问题更加理想一些的方式是自动化地找出最优解。常见的方式比如整数线性规划启发式搜索或者强化学习。
对于MP的策略搜索其实也可以看作device placement问题。对于该问题几年前就有不少相关工作。虽然大部分是针对单机多卡或者CPU与GPU的异构计算等小规模场景但原理是类似的。其中很多是基于强化学习方法。如2017年的论文《Device Placement Optimization with Reinforcement Learning》中的ColocRL使用强化学习让RNN策略网络来预测计算图中操作的placement从而来解决CPU与多GPU环境中的device placement问题。2018年的论文《A Hierachical Model for Device Placement》是对它的改进引入了更加灵活的方式免除了需要人工对算子分组的繁琐。还有2018年的论文《Spotlight: Optimizing Device Placement for Training Deep Neural Networks》中的Spotlight基于强化学习中的PPO算法。2018年的论文《Post: Device Placement with Cross-entropy Minimization and Proximal Policy Optimization》结合了cross-entropy minimization和PPO算法。前者用于batch间后者用于batch内从而提高学习效率。2018年的论文《Placeto: Efficient Progressive Device Placement Optimization》中的Placeto也是基于强化学习。它通过迭代式生成placement方案和使用graph embedding提高了模型的泛化能力。2019年的论文《REGAL: Transfer Learning For Fast Optimization of Computation Graphs》中提出的REGALREINFORCE-based Genetic Algorithm Learning则是结合了强化学习与遗传算法更具体地BRKGA。该算法由强化学习REINFORCE算法训练GNN表示的策略网络负责输出分布用于遗传算法中的mutation和crossover。
2018年的论文《Beyond Data and Model Parallelism for Deep Neural Networks》中的FlexFlow使用自动搜索的方法来找图中operation在考虑DP与MP情况下的最优切分。FlexFlow是一个高用编排TP计算的DP框架。它使用Markov Chain Monte Carlo算法来进行策略的搜索。其中会考虑sharding strategy和placement。
2019年的论文《Optimizing Multi-GPU Parallelization Strategies for Deep Learning Training》提出MP与DP结合可以有效解决DP的scaling与statistical效率降低的问题。为了找到混合DP与MP下的最优策略它使用了基于整数线性规划的工具-DLPlacer用于求解device placement问题最大化并行度减少数据通信开销。它以最小化单步的训练时间为目标给出operation在硬件资源上的分配方案schedule与activation等数据的通信路径。
2020年的论文《Efficient Algorithms for Device Placement of DNN Graph Operators》基于cost model自动切分模型到多个device。Device placement问题中的核心是组合优化问题。文中使用基于Integer ProgrammingIP与Dynamic ProgrammingDP来解决。并且支持non-contiguous fragment of graph。类似的DP也在PipeDream中使用区别是PipeDream只支持线性的layer-granularity graphs。
2020年的论文《Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training》是一种在DP中对weight update做自动切分的方法。在DP中所有的replica都会做weight update计算因为weight没有batch维。一种方案是weight update计算sharding。但简单的做法会导致data formatting和replica间通信开销极大增加。该文自动将weight update计算在replica间做sharding。它基于XLA基于static analysis为了正确性与性能和graph transformation。不需要修改模型代码。
2020年的论文《Optimizing Distributed Training Deployment in Heterogeneous GPU Clusters》提出的HeteroG是一个用于异构GPU集群中加速深度神经网络训练的模块。HeteroG使用operation-level的hybrid parallelismcommunication architecture selection和execution scheduling。它基于GNN-based learning和组合优化。
2021年的论文《Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training》中介绍了Amazon的SageMaker。它采用module server架构支持PPTP和DP多种并行方法并可以对模型进行自动的切分与分布。切分算法的目标是尽可能使device上的计算与访存平衡并减少通信次数。它还为PP提供专用的通信后端。它管理node内与node间的device-to-device通信而无需依赖NCLL。
2021年的论文《Varuna: Scalable, Low-cost Training of Massive Deep Learning Models》中的Varuna自动配置分布式(MP与DP)的训练它可以利用spot VMs而非dedicated VMs使成本降低5倍。它解决三个挑战慢及不稳定的网络资源可能被抢占对用户透明。对于网络问题它的micro-batch调度算法可能更加高效应对network jitter另一方面stage的切分还考虑了网络带宽。对于抢占问题它采用了一种dynamic, semantics-preserving reconfiguration方法称为Job morphing。它可以在训练过程中配置pipeline depth等参数从而适配资源。为了找到更优的配置它有提供了micro-benchmark-driven simulation机制。对于易用性问题它引入了自动找模型中cut-points的机制。
2022年的论文《Accelerating Large-Scale Distributed Neural Network Training with SPMD Parallelism》根据DNN model与cluster specification自动找sharding strategy。它通过overlap通信与计算来减少通信开销。两个device上计算两个micro-batch可以并行一个通信时另一个计算。构建training time cost model并设计基于DP的搜索算法。基本思想是将每个device上分配的计算图复制一份。每个用一半的输入数据训练。这样两个micro-batch就能形成pipeline从而将计算与通信overlap。 基于PyTorch实现系统HiDup。它的输入为模型和集群的 spec。输出是可用于运行在多个GPU上的模型。
2022年论文《Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning》中的Alpa是一个用于GPU cluster的DL编译系统。它构建于Jax之上支持自动地将JAX函数并行化分析计算图并根据计算图与cluster产生并行执行计划从而结合DPTP和PP将DL运行在分布式集群上。Alpa将并行分为两个层次inter-operator与intra-operator并行将它们分别看作子问题来求解。后者由于通信量需求大一般分配在由高带宽连接的设备中。它将device cluster切成若干device mesh。首先Inter-op pass使用DP将stage分配给mesh。然后Intra-op pass使用ILP优化给定stage与mesh上的并行执行计划再将cost反馈给inter-op pass。这些信息帮助inter-op pass使用DP最小化inter-op的执行时间从而获得更优的stage与mesh划分策略。此外runtime orchestration pass用来满足相邻stage的通信需求。
2023年的论文《Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models》是一个基于PyTorch的自动化的intra-op 不考虑inter-op并行系统。可以对分布式并行和activation checkpoint进行联合优化。它使用anayzer模块收集硬件性能和计算图开销信息采用采用two-stage的求解器找到执行计划最后将PyTorch模型编译成分布式的模型。Analyzer包含三个部分symbolic profiler收集由PyTroch FX模块生成的静态计算图的计算和访存开销。Cluster detector收集硬件特性和集群的拓扑结构信息。Tensor layout manager用于确定tensor layout的conversion策略。并行执行计划的搜索分为intra-op与activation checkpoint两层。前者采用的是Alpa中用于处理intra-op的ILP求解器后者采用Rotor算法。Generator将生成的执行计划应用于原图然后将优化过的图编译成PyTorch代码。其中的pass会插入通信节点对参数做sharding和对参数做reshape conversion。
2023年的论文《OSDP: Optimal Sharded Data Parallel for Distributed Deep Learning》中的OSDP(Optimal Sharded Data Parallel)是结合DP与MP的自动并行训练系统。它在内存使用与硬件利用率间做trade-off自动产生最大化吞吐率的分布式的计算图。它指出ZeRO系统的两个缺点与vanilla DP比有额外50%的通信开销中间的巨大张量会使peak memory usage超限。对于第一个问题它对于每个operator确定是否要做parameter sharding。对于第二个问题它使用operator splitting技术。OSDP包含三个modulesProfiler, Search Engine和Scheduler。Profiler负责根据Search Engine给的execution plan与batch size基于cost model估计memory和time cost。Search Engine基于Profiler的cost估计用DFS遍历所有operator生成memory不超过限制下的最优execution plan。Scheduler迭代地收集Search Engine给出的plan与throughput作为candidate同时增加batch size直到所需memory超过device限制。
