张家港网站网络公司,网站开发涉及技术,佛山市网能建设有限公司,申请做网站、论坛版主来源#xff1a;AI科技评论编译#xff1a;琰琰大规模深度神经网络训练仍是一项艰巨的挑战#xff0c;因为动辄百亿、千亿参数量的语言模型#xff0c;需要更多的 GPU 内存和时间周期。这篇文章从如何多GPU训练大模型的角度#xff0c;回顾了现有的并行训练范式#xff0… 来源AI科技评论编译琰琰大规模深度神经网络训练仍是一项艰巨的挑战因为动辄百亿、千亿参数量的语言模型需要更多的 GPU 内存和时间周期。这篇文章从如何多GPU训练大模型的角度回顾了现有的并行训练范式以及主流的模型架构和内存优化设计方法。本文作者Lilian Weng现为OpenAI应用人工智能研究负责人主要从事机器学习、深度学习和网络科学研究 。她本科毕业于香港大学硕士就读于北京大学信息系统与计算机科学系之后前往印度安纳大学布鲁顿分校攻读博士。Lilian Weng经常在个人博客分享学习和工作笔记感兴趣的可以戳这里https://lilianweng.github.io/lil-log/。“炼大模型”已成为人工智能领域的主流研发趋势。从GPT-3的1750亿到如今悟道2.0的1.75万亿超大语言模型在 NLP 基准任务中不断刷新SOTA。然而参数和数据集的快速增长让 GPU 算力开始捉襟见肘。单个GPU内存已经远远不能满足大模型的需求。如阿里用480块GPU训练千亿模型英伟达用3072块GPU训练万亿模型谷歌用2048块TPU训练1.6万亿模型1 TPU约等于23 GPU。如何利用上百块GPU上训练大规模语言模型并行计算是一种行之有效的方法。近日OpenAI 研究员Lilian Weng分享干货文章从并行训练数据并行、模型并行、管道并行、张量并行、混合专家、内存节省设计CPU卸载、重新激活计算、混合精度训练、高效存储优化器三个方面回顾了现阶段多GPU训练大模型的主流方法。AI科技评论这篇文章编译如下供想成为“炼丹师”的朋友参考。1并行训练大规模神经网络模型对存储空间有强烈的需求单个GPU的内存已经远远不够。究其原因一方面模型权重有数百亿个浮点数随机梯度下降和Adam优化需要极高的计算成本另一方面在预训练阶段大模型与大规模语料库的配对需要很长的时间周期。综合来看跨GPU并行计算显得尤为重要。并行计算在数据、模型架构和张量等不同维度上都可以操作接下来本文将具体介绍一些主流方法数据并行 数据并行 Data parallelism DP最简单的方法是将相同的模型权重复制到worker节点并分配一部分数据以同时进行处理。我们知道如果模型的参数量大于单个GPU节点的内存DP无法正常工作GeePS架构Cui等人2016的解决思路是使用有限的GPU内存。也就是如果模型太大无法嵌入到一台机器就将暂时未使用的参数卸载回CPU。数据交换传输通常在后端进行不干扰训练计算在每个Mini-batch计算结束后worker需要同步梯度或权重以保证学习效率。现有的同步方法有两种各自优缺点如下1、批量同步并行BSPworker在每个Mini-batch结束时同步数据这种方法保证了模型权重传递的及时性但每台机器都必须排队等待其他机器发送梯度。2、异步并行ASP每个GPU采用异步方式处理数据这种方法避免了不同机器之间的相互等待或暂停但影响了权重传递的时效降低了统计学习效率。而且即使增加计算时长也不会加快训练的收敛速度。 在中间某些地方的每一次迭代1都需要同步全局梯度。自Pytorch v1.5版Li等人2021年提出后该特征在分布式数据并行Distribution Data ParallelDDP中被称为“梯度累积gradient accumulation”。分桶梯度(bucketing gradients)避免立即执行AllReduce操作而是将多个梯度存储到一个AllReduce中以提高吞吐量并基于计算图优化计算和通信调度。图1Pytorch DDP的伪代码来源Li等人2021年模型并行 模型并行Model parallelismMP用于解决模型权重不能适应单个节点的情况在这里计算和模型参数都需要跨多台机器进行处理。在数据并行中每个worker承载着整个模型的完整副本而MP只在一个worker上分配部分模型参数因此对内存和计算的需求要小很多。 深度神经网络包含一堆垂直层如果逐层拆分将连续的小层分配到工作层分区操作起来并不难但通过大量具有顺序依赖性的Workers来运行每个数据batch会花费大量的时间计算资源的利用率也严重不足。图2一个包含4个垂直层的模型并行设置由于顺序的依赖性每个数据依次由一个worker处理这个过程会出现大量多余时间“气泡”来源Huang等人2019年管道并行管道并行Pipeline parallelismPP是将模型并行与数据并行结合起来以减少低效时间“气泡”的过程。主要思想是将Mini-batch拆分为更多个微批次microbatch并使每个阶段worker能够同时处理。需要注意的是每个微批次需要两次传递一次向前一次向后。worker分区的数量称为管道深度不同worker分区之间的通信仅传输激活向前和梯度向后。这些通道的调度方式以及梯度的聚合方式在不同的方法中有所不同。 在GPipeHuang et al.2019方法中多个微批次处理结束时会同时聚合梯度和应用。同步梯度下降保证了学习的一致性和效率与worker数量无关。如图3所示“气泡”仍然存在但比图2少了很多。给定m个均匀分割的微批次和d个分区假设每个微批次向前和向后都需要一个时间单位则气泡的分数为 GPipe论文表明如果微批次的数量超过分区数量4倍m4d则“气泡”开销几乎可以忽略不计。 图3带有4个微批次 和4个分区的GPipe的并行管道来源Huang等人2019年 GPipe在吞吐量方面与设备数量成线性关系设备数量越多吞吐量越大。不过如果模型参数在worker中分布不均匀这种线性关系不会稳定出现。 PipeDreamNarayanan等人2019年方法要求每个worker交替处理向前和向后传递的消息1F1B。它将每个模型分区命名为“stage”每个stage worker可以有多个副本来并行运行数据。这个过程使用循环负载平衡策略在多个副本之间分配工作以确保相同minibatch 向前和向后的传递发生在同一副本上。 图4PipeDream中1F1B微批次调度的图示来源Harlap等人2018年由于PipeDream没有在所有worker batch结束时同步全局梯度1F1B 很容易导致不同版本的模型权重的微批次向前和向后传递降低学习效率。对此PipeDream提供了一些解决的思路权重存储每个worker跟踪多个模型版本给定数据 batch 的向前和向后传递相同版本的权重。垂直同步不同模型权重版本与激活和梯度一起在全局worker之间传递计算采用上一个worker传播的相对应的隐藏版本。这个过程确保了worker之间的版本一致性不同于GPipe采用异步计算。 在训练开始时PipeDream会先分析模型每一层的计算内存和时间成本然后将层划分为不同的stage进行优化。图5上图为VGG16在ILSVRC12上的运行结果ASP异步并行BSP批量同步并行下图为不同并行配置的训练时间加速度来源Harlap等人2018年 后来有学者提出了PipeDream两种变体主要思路是通过减少模型版本来缓解内存占用Narayanan等人2021年。其中PipeDream-flush增加了定期刷新全局同步管道的功能就像GPipe一样这种方式虽然牺牲了一点吞吐量但显著减少了内存占用例如仅需要维护单一版本的模型权重。图6PipeDream flush中管道调度图示来源Narayanan等人2021年PipeDream-2BW维护两个版本的模型权重“2BW”代表“双缓冲权重double-buffered weights”它会在每个微批次生成一个新的模型版本KKd。由于一些剩余的向后传递仍然依赖于旧版本新的模型版本无法立即取代旧版本但因为只保存了两个版本内存占用的也被大大降低了。图7PipeDream-2BW 中的流水线调度示意图来源Narayanan et al. 2021张量并行模型并行和管道并行都会垂直拆分模型而张量并行Tensor ParallelismTP是将张量运算的计算水平划分到多个设备上。以Transformer为例。Transformer架构主要由多层MLP和自注意力块组成。Megatron-LMShoeybi et al.2020采用了一种简单的方法来并行计算层内MLP和自注意力。MLP层包含GEMM通用矩阵乘法和非线性GeLU传输。如果按列拆分权重矩阵A可以得到注意力块根据上述分区并行运行GEMM的 查询Q、键K和 权重V然后与另一个GEMM组合以生成头注意力结果。图8Megatron-LM中提出的关键Transformer组件的张量平行性说明。来源Shoeybi等人2020年今年Narayanan等人将管道、张量和数据并行与新的管道调度策略相结合提出了一种名为PTD-P的新方法。该方法不仅在设备上能够定位一组连续的层“模型块”该可以为每个wokers分配多个较小的连续层子集块例如设备1具有第1、2、9、10层设备2具有第3、4、11、12层每个具有两个模型块每个batch中微批次的数量应精确除以wokers数量mm。如果每个worker有v个模型块那么与GPipe调度相比管道的“气泡”时间可以减少 v 倍。图9上图与PipeDream flush中的默认1F1B管道明细表相同下图为交错的1F1B管线一览表来源Narayanan等人2022混合专家MoE为了突破模型大小的限制谷歌后来提出一种混合专家MoE方法其核心理念是综合学习它假设多个弱学习者组合起来就会拥有一个强学习者。在深度神经网络中混合专家MoE通过连接多个专家的门机制gating mechanism实现集成Shazeer等人2017。门机制激活不同网络的专家以产生不同的输出。作者在论文将其命名为“稀疏门控专家混合层sparsely gated MoE”。仅一个MoE层包含1前馈网络专家n2可训练的门控网络G通过学习n个专家的概率分布将流量路由到几个特定的专家。 根据门控输出并非每个专家都必须进行评估。当专家的数量太大时可以考虑使用两层MoE。 图10专家混合MoE层的图示门控网络只选择并激活了n个专家中的2个来源Shazeer等人2017年G将输入与可训练权重矩阵Gg相乘然后执行softmax由于这个过程会产生密集的门控制向量不利于节省计算资源而且时也不需要评估专家。所以MoE层仅保留了顶部k值并通过向G中添加高斯噪声改进负载平衡这种机制被称为噪声top-k门。 为了避免门控网络可能始终偏向少数强势专家的自我强化效应Shazeer等人2017提出了通过额外重要损失的软约束以鼓励所有专家拥有相同的权重。其数值相当于每个专家的分批平均值变异系数的平方 其中CV是变异系数失重的waux是可调节的超参数。由于每个专家网络只能获得小部分训练样本“收缩批次问题”所以在MoE中应该尽可能使用大batch但这又会受到GPU内存的限制。数据并行和模型并行的应用可以提高模型的吞吐量。图1110亿单词的语言建模基准左模型从左到右包含4、32、256、256、1024和4096名专家右40亿参数的MoE模型在不同计算预算下的性能来源Shazeer等人2017年GShardLepikhin等人2020年通过自动分片将MoE transformer 模型的参数扩展到了6000亿。MoE transformer 用MoE层取代其他每一个前馈网络层。需要说明的是在多台机器上MoE transformer 仅在MoE层分片其他层只是复制。 GShard中的门控功能G有几种改进设计专家容量通过一位专家的令牌数量不应超过“专家容量”的阈值。如果令牌被路由到已达到容量的专家则令牌将被标记为“溢出”并且门输出将更改为零向量。本地组调度令牌被均匀地划分为多个本地组专家能力在组水平上得到加强。辅助损失与原始MoE aux损失相似添加辅助损失可以最小化路由到每个专家的数据的均方。随机路由以与其权重成比例的概率选择第二位最佳专家否则GShard遵循随机路由增加随机性。图12GShard中带辅助损失的组水平top-2门机制的伪代码来源Lepikhin等人2020年 Switch TransformerFedus et al.2021用稀疏开关FFN层取代了密集前馈层每个输入仅路由到一个专家网络将模型规模扩展到数万亿个参数。负载平衡的辅助损失是给定n个专家fi是路由到第i个专家的令牌分数pi是门控网络预测的专家i的路由概率。 图13Switch transformer稀疏Switch FFN层位于蓝色框来源Fedus等人2021年为提高训练稳定性switch transformer采用以下设计选择精度使用FP32精度以提高模型局部的稳定性并降低FP32张量的通信成本。FP32精度仅在路由器功能主体内使用结果将还原到FP16。 较小的初始化权重矩阵的初始化从平均μ0且标准偏差的正态分布中采样同时将Transformer初始化参数从s1减小到s0.1 。使用更高的专家辍学率微调通常适用于小数据集增加每个专家的辍学率以避免过度拟合。他们发现所有层中的辍学率增加会导致性能下降。在论文中他们在非专家层中使用了0.1的辍学率但在专家FF层中使用了0.4的辍学率。switch transformer论文总结了用于训练大型模型的不同数据和模型并行策略并给出了一个很好的示例 图14第一行为如何在多个GPU内核拆分模型权重顶部每种颜色代表一个权重矩阵第二行为各种数据并行策略的说明不同颜色表示不同的标记集来源Fedus等人2021年3其他节省内存的设计 CPU卸载如果GPU内存已满可以将暂时未使用的数据卸载到CPU并在以后需要时将其读回Rhu等人2016。不过这种方法近年来并不太流行因为它会延长模型训练的时间。 激活重新计算 激活重新计算也称“激活检查点”或“梯度检查点”Chen et al2016其核心思路是牺牲计算时间来换取内存空间。它减少了训练 ℓ 层深层神经网络到的内存开销每个batch只消耗额外的前向传递计算。具体来说该方法将ℓ层网络平均划分为d个分区仅保存分区边界的激活并在workers之间进行通信。计算梯度仍然需要在分区内层进行中间激活以便在向后过程中重新计算梯度。在激活重新计算的情况下用于训练M(ℓ) 是它的最低成本是激活重新计算的方法可以得出与模型大小有关次线性内存开销如下图 图15不同节省内存算法的内存开销。sharing中间结果使用的内存在不再需要时被回收。inplace将输出直接保存到输入值的内存中来源Chen等人2016混合精度训练此前NarangNarangMicikevicius等人2018年介绍了一种使用半精度浮点FP16数训练模型而不损失模型精度的方法。图16一层混合精度训练程序来源NarangMicikevicius等人2018年 其中涉及三种关键技术全精度权重复制保持累积梯度的模型权重的全精度FP32复制。对于向前和向后的传递的信息做四舍五入至半精度处理因为每次梯度更新即梯度X学习率太小可能无法完全包含在FP16范围内。缩放损失放大损失以更好地处理小幅度的梯度见图16放大梯度以使其向可表示范围的右侧部分包含较大的值移动从而保留可能丢失的值。算术精度对于常见的网络算法如矢量点积、矢量元素求和归约将部分结果累加到FP32中然后输出保存为FP16。逐点操作可以在FP16或FP32中执行。图17:全精确的梯度直方图在这项实验中图像分类、更快的R-CNN等不需要损失缩放但其他网络如多盒SSD、大LSTM语言模型是需要损失缩放的。压缩Compression 模型权重在向前和向后传递的过程中会消耗大量内存。考虑到这两种传递方式会花费大量时间2018年Jain Jain et al2018提出了一种数据编码策略即在第一次传递后压缩中间结果然后将其解码用于反向传播。 Jain和团队研发的Gist系统包含两种编码方案一是特定于层的无损编码包括 ReLU-Pool和 ReLU-Conv模式二是有攻击性的有损编码主要使用延迟精度缩减DPR。需要注意的是第一次使用特征图时应保持高精度第二次使用时要适度降低精度。这项实验表明Gist可以在5个最佳图像分类DNN上减少2倍的内存开销平均减少1.8倍性能开销仅为4%。 内存高效优化器优化器也会消耗内存。以主流的Adam优化器为例其内部需要维护动量和方差这两者与梯度和模型参数比例基本相同。这意味着我们需要节省4倍模型权重的内存。 为了减少内存消耗学术界已经提出了几款主流优化器。与Adam相比AdafactorShazeer et al.2018优化器没有存储全部动量和变化只跟踪移动平均数的每行和每列总和然后根据这些总和估计二阶矩。SM3Anil et al.2019优化器采用了一种不同的自适应优化方法。ZeRORajbhandari et al.2019零冗余优化器节省了大型模型训练在两方面的内存消耗大多数内存由模型状态消耗包括优化器状态例如Adam动量和方差、梯度和参数。混合精度训练也需要大量内存因为除了FP16版本之外优化器还需要保存FP32参数和其他优化器状态的副本。 未被激活、临时缓冲区以及不可用的碎片内存消耗论文中称为剩余状态。ZeRO结合了ZeRO-DP和ZeRO-R两种方法。ZeRO-DP是一种增强的数据并行避免了模型状态的简单冗余。它以动态的方式跨多个并行数据划分优化器状态、梯度和参数以最小化通信量。ZeRO-R使用分区激活二次计算、恒定缓冲区大小和动态内存碎片以优化剩余状态的内存消耗。参考资料[1] Li et al. “PyTorch Distributed: Experiences on Accelerating Data Parallel Training” VLDB 2020.[2] Cui et al. “GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server” EuroSys 2016[3] Shoeybi et al. “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.” arXiv preprint arXiv:1909.08053 (2019).[4] Narayanan et al. “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.” arXiv preprint arXiv:2104.04473 (2021).[5] Huang et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.” arXiv preprint arXiv:1811.06965 (2018).[6] Narayanan et al. “PipeDream: Generalized Pipeline Parallelism for DNN Training.” SOSP 2019.[7] Narayanan et al. “Memory-Efficient Pipeline-Parallel DNN Training.” ICML 2021.[8] Shazeer et al. “The Sparsely-Gated Mixture-of-Experts Layer Noam.” arXiv preprint arXiv:1701.06538 (2017).[9] Lepikhin et al. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” arXiv preprint arXiv:2006.16668 (2020).[10] Fedus et al. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” arXiv preprint arXiv:2101.03961 (2021).[11] Narang Micikevicius, et al. “Mixed precision training.” ICLR 2018.[12] Chen et al. 2016 “Training Deep Nets with Sublinear Memory Cost.” arXiv preprint arXiv:1604.06174 (2016).[13] Jain et al. “Gist: Efficient data encoding for deep neural network training.” ISCA 2018.[14] Shazeer Stern. “Adafactor: Adaptive learning rates with sublinear memory cost.” arXiv preprint arXiv:1804.04235 (2018).[15] Anil et al. “Memory-Efficient Adaptive Optimization.” arXiv preprint arXiv:1901.11150 (2019).[16] Rajbhandari et al. “ZeRO: Memory Optimization Towards Training A Trillion Parameter Models Samyam.” arXiv preprint arXiv:1910.02054 (2019).编译链接https://lilianweng.github.io/lil-log/2021/09/24/train-large-neural-networks.html未来智能实验室的主要工作包括建立AI智能系统智商评测体系开展世界人工智能智商评测开展互联网城市大脑研究计划构建互联网城市大脑技术和企业图谱为提升企业行业与城市的智能水平服务。每日推荐范围未来科技发展趋势的学习型文章。目前线上平台已收藏上千篇精华前沿科技文章和报告。 如果您对实验室的研究感兴趣欢迎加入未来智能实验室线上平台。扫描以下二维码或点击本文左下角“阅读原文”
