Transformer 高效训练综述

Bohan Zhuang1  Jing Liu1  Zizheng Pan1  Haoyu He1  Yuetian Weng1  Chunhua Shen2
1ZIP Lab, Monash University 2Zhejiang University
Correspondence should be addressed to BZ. Email: 𝚋𝚘𝚑𝚊𝚗.𝚣𝚑𝚞𝚊𝚗𝚐@𝚐𝚖𝚊𝚒𝚕.𝚌𝚘𝚖
摘要

Transformer 的最新进展对计算资源提出了巨大的要求,这凸显了开发高效训练技术的重要性,通过有效利用计算和内存资源,使 Transformer 训练更快、成本更低、精度更高。 本次调查首次系统概述了 Transformers 的高效训练,涵盖了加速算法和硬件方面的最新进展,重点关注前者。 我们分析和比较了在训练期间节省中间张量的计算和内存成本的方法,以及硬件/算法协同设计的技术。 最后我们讨论了未来研究的挑战和有希望的领域。

1简介

深度学习是最近最深刻的方法,它彻底改变了机器学习(ML)和人工智能,并正在引领第四次工业革命。 从本质上讲,深度学习的成功取决于大量可用的计算资源和大量的标记数据。 尽管最近的发展令人兴奋,但深度学习模型,尤其是 Transformer [87],已经变得非常庞大且计算密集,导致在基础层面上面临两个紧迫的挑战。

第一个问题涉及训练基于 Transformer 的大型模型的密集计算。 一项广泛讨论的深度学习模型能源研究[82]估计,使用神经架构搜索(NAS)[79]训练 Transformer 基础模型会产生大约 626,155 磅的行星-二氧化碳变暖,相当于五辆汽车一生的排放量;随着模型变得越来越大,它们对计算的需求超过了硬件效率的提高。 例如,GPT-3 [10](ChatGPT 的前身)接受了 5000 亿个单词的训练,并配备了 1750 亿个参数。 值得注意的是,根据 GPT-3111https://lambdalabs.com/blog/demystifying-gpt-3/,单次训练运行需要 355 个 GPU 年,成本至少为 460 万美元,估计 V100 和 V100 的理论计算速度为 28 TFLOPS最低的 3 年预留云定价。 因此,深度学习的真正突破性成功,例如ChatGPT,完全由谷歌或微软等大型且富有的企业主导。 让深度学习在计算和能源效率方面适用于绿色人工智能[76],并在资源有限的情况下将人工智能民主化到更广泛的社区变得极其重要。

第二个问题是训练内存与基于注意力的模型大小成正比呈指数增长。 例如,文献中最大的语言模型从 2018 年 BERT-large [48] 的 345M,发展到现在 MT-NLG [78] 等模型的数千亿t1>配备530B参数。 因此,这些 SOTA 大规模模型需要内存高效的训练技术,以减少存储中间张量和跨加速器的数据交换(通信)的内存占用,同时确保高处理元件(PE)利用率。

在本次调查中,我们回顾了提高基于训练注意力的模型的计算和内存效率的通用技术,i.e.,Transformers,如图1 我们通过技术创新和主要用例来描述它们的特征,总结它们并建立它们之间的联系。 我们主要对提高变形金刚训练效率的算术创新感兴趣,并简要讨论硬件/算法协同设计的进步。 我们将硬件加速器设计的审查(一个广泛的类别)作为未来的工作。

Refer to caption
图1: Transformers 高效训练概述,包括计算效率、内存效率和硬件/算法协同设计视角。

2计算效率

2.1优化

优化器 为了实现梯度下降更快的收敛速度,经典的解决方案是融合动量技术,其中每一步都是最速下降方向和最近迭代位移的组合,有助于加速相关方向的梯度下降并抑制振荡。 开创性的工作包括Nesterov针对凸优化的加速梯度[67]和针对非凸问题的动量近端梯度[55]等。机器学习模型的规模优化,主导优化器是以随机方式设计的。 特别是,动量随机梯度下降(SGD)和自适应学习率估计方法 Adam [50] 被广泛用于训练深度神经网络。 根据经验,使用 Adam 训练 Transformers 的性能优于对应的 SGD,并且 [96] 揭开了随机梯度中噪声的重尾分布是 SGD 性能不佳的主要原因的神秘面纱,并通过自适应噪声削波。 默认情况下,AdamW [62] 是 Adam 的一个变体,它解耦了 L2 正则化和权重衰减,是 Transformers 最广泛使用的优化器。 最近,Google 搜索优化算法并发现了一种简单有效的优化器,称为 Lion [18] Lion仅跟踪一阶梯度的动量,其更新仅考虑sign方向,并且每个参数具有相同的幅度,这与AdamW等自适应优化器有很大不同。 在实践中,在各种基准上训练 Transformer 时,Lion 通常收敛速度更快,并且比 AdamW 更节省内存且更准确。 我们建议读者参阅[60, 9],了解有关机器学习中加速优化方法的更多详细信息。

为了提高 Transformers 的泛化能力,锐度感知最小化 (SAM) [28] 寻求基于损失景观的几何形状与泛化之间的联系,同时最小化损失值和损失锐度,i.e.,更平坦的最小值往往会提高泛化能力。 以下工作[17]将 SAM 应用到 Transformer,通过平滑损失表面观察到显着的精度增益。 然而,SAM 需要解决双层最小-最大优化问题,这几乎使训练时间增加了一倍。 为了加速优化,[23]提出了随机权重扰动来保持泛化能力和锐度敏感的子集选择策略。 最近,[24]通过将锐度估计替换为两个连续更新步骤之间的 KL 散度,设计了锐度损失的近零成本代理。

初始化 良好的初始化对于稳定训练、提高学习率、加速收敛和提高泛化能力至关重要。 因此,为了更好地初始化 Transformer,人们提出了许多工作。 具体来说,Fixup [95] 建议适当地重新调整标准初始化,以确保适当的梯度范数,以避免梯度爆炸或消失,这可以训练具有超过 10,000 层的非常深的网络,而无需添加归一化层。 基于归一化残差块计算的函数接近恒等函数(i.e.,单位方差)的见解,以下工作 ReZero [ 5]和SkipInit[22]只是简单地初始化每一层来执行恒等操作。 具体来说,他们在每个残差块的输出上添加了一个可学习的缩放乘数:

𝐱l+1=𝐱l+αlFl(𝐱l), (1)

其中𝐱lFl()是层l的输入和函数,其中函数可以是多头自注意力层(MSA)和feed-前向网络(FFN),并且 αl 简单地初始化为 0。 针对 Transformers 定制的 T-Fixup [40] 分析认为,部分优化难度来自于 Adam 优化器的早期更新不稳定,因为二阶动量的方差是无界的。 因此,跟随Fixup采用重缩放方案来初始化残差块。 所有上述方法都从所有块和训练中删除了批/层归一化,而无需学习率预热。 在训练深度视觉 Transformers (ViT) 上,[85] 提出了通道可学习的缩放因子,并根据经验观察到重新引入预热和层归一化技术可以使训练更加稳定。

除了重缩放范式之外,一些文献提出通过利用自注意力和卷积层之间的关系从新的角度改进初始化。 [19] 证明了具有 N 个头和维度 D3 的相对位置编码的多头自注意力层可以重新参数化以表达任何卷积过滤器尺寸N×N的层。 注意力可以分解为内容项和(相对)位置项,其中后者决定每个头的注意力中心和宽度。 基于这一性质,ConViT [25]通过添加软门控参数来平衡两项来学习控制局部性,从而具有将软卷积归纳偏差纳入全局自注意力的效果。

稀疏训练 稀疏训练的关键思想是直接训练稀疏子网络而不是从头开始训练完整网络而不牺牲准确性。 可靠性首先由彩票假设 (LTH) [29] 证明,即密集、随机初始化的网络包含子网络(中奖彩票),可以单独训练这些子网络以匹配原始网络的准确性。 然而,LTH 需要以交替训练-修剪-再训练的方式识别中奖彩票,这使得大型模型和数据集的训练成本极高,限制了实际收益。 有鉴于此,后续训练效率较高的工作可以大致分为三类:(i)在初始化时通过测量连接对损失的重要性来找到稀疏网络,从而消除了复杂的迭代优化计划的需要[52, 88]; (ii) 通过低成本方案在非常早期的训练阶段确定变形金刚训练中的中奖彩票,然后仅识别这些早期彩票,直到收敛[90, 16]; (iii) 在整个训练过程中使用交替修剪和生长计划来动态更新模型稀疏模式,适用于一般架构[26, 14]

过度参数化 实际的 DNN 训练过度参数化,可学习参数的数量远大于样本的数量。 据观察,过参数化在经验上提高了收敛性和泛化性,虽然有理论保证,但还不够。 早期工作[4]在数学上证明了线性神经网络中过度参数化增加深度可以加速 SGD 收敛。 [57]进一步探索了两层非线性神经网络,[3]证明SGD可以在多项式中的DNN训练目标上收敛到全局最小值时间,假设训练样本不重复,并且参数数量是训练样本数量和网络深度的多项式。 在泛化方面,[2]从理论上证明了充分超参数化(三层)的神经网络可以泛化到群体风险,并且一个有趣的特性是 在 SGD 训练轨迹上的任何点的近邻中都存在一个精确的网络,并且在随机初始化上具有很高的概率。 请注意,它与 LTH 有着深厚的联系,因为它部分解释了为什么 LTH 训练处于稀疏状态,因为由于过度参数化,低风险的良好小型子网络非常丰富。 应用于变形金刚,[58]利用超参数化理论的更快收敛和更好泛化来设计高效的训练管道:训练一个非常大的模型,然后执行早期停止并对其进行严重压缩,类似于LTH 。

大批量训练 加速训练的另一种流行方法是使用大批量,减少每个时期的迭代次数并提高计算资源利用率。 从统计角度来看,大批量训练减少了随机梯度估计的方差,因此需要调整可靠的步长以获得更好的收敛[9] 在卷积神经网络时代,[31]采用学习率的线性缩放在ImageNet上训练ResNet-50,批量大小为1小时8,192。 随后提出了更先进的步长估计方法。 广泛使用的方法是 SGD 的 LARS [92] 和 Adam 的 LAMB [93],它们分别建议对 ResNet 和 Transformers 使用分层自适应学习率。 分层适应策略可以表示为

𝐰t+1i=𝐰tiηtϕ(𝐰ti)𝜸ti𝜸ti, (2)

其中ηt,𝐰ti𝜸ti是第i层在某个时刻的学习率、参数和基于动量的梯度步骤tϕ是缩放函数。 它配备了一个归一化项,可以为爆炸梯度和平台提供鲁棒性,并且缩放项确保更新的范数与参数的范数具有相同的阶数,从而促进更快的收敛。 最近,经验表明,针对大批量训练定制的更强大的优化方法表现良好。 例如,[46]表明对一定数量的最新检查点的权重进行平均可以促进更快的训练。 [37] 中的 DeepMind 训练了 400 多个 Transformer 语言模型,模型大小和训练标记数量的比例各不相同,达到了一个实际假设,即模型大小和训练标记数量应同等缩放计算最优大语言模型训练。

增量学习 增量学习的高级概念是将原始具有挑战性的优化问题放松为一系列易于优化的子问题,其中一个子问题的解决方案可以作为后续子问题的良好初始化,以规避训练难度,与退火类似。 一些作品[30, 32]提出通过逐步堆叠层、正确地从较小的模型初始化较大的模型来加速 BERT 预训练。 [97] 与通过层下降随机深度训练变形金刚相反,它沿着时间维度和深度维度逐渐增加下降率。 AutoProg [54] 针对 ViT 进行了定制,建议使用神经架构搜索在渐进式学习过程中自动决定模型是否增长、在何处增长以及增长多少。 一个关键的观察结果是,逐步提高输入图像的分辨率(减小补丁大小)可以显着加速 ViT ,这与众所周知的训练动态一致,该训练动态侧重于早期阶段的低频结构和后期阶段的高频语义。后期。

2.2数据选择

除了模型效率之外,数据效率也是高效训练的关键因素。

词符掩蔽 词符掩蔽是自监督预训练任务中的主要方法,例如掩蔽语言建模(MLM)[48, 10]和掩蔽图像建模(MIM)[6, 35 ] 词符掩码的本质是随机掩码一些输入标记,并训练模型使用可见标记中的上下文信息来预测丢失的内容,例如词汇 ID 或像素。 由于压缩序列长度可以平方降低计算和内存复杂性,因此跳过处理屏蔽标记可以为 MLM 和 MIM 带来可观的训练效率增益。 对于MLM,[80]提出联合预训练编码器和解码器以执行语言生成任务,同时删除解码器中的屏蔽标记以节省内存和计算成本。 对于 MIM,代表性工作 [35] 表明,在视觉中,在编码器之前去除屏蔽图像块表现出更强的性能,并且总体预训练时间和内存降低 3× 或更多消耗比保留屏蔽 Token 更重要。 [56]中也发现了类似的现象,即对于语言图像预训练,随机屏蔽和移除屏蔽图像片段的预训练总体时间比原始 CLIP [72]快 3.7×

重要性采样 对数据的重要性采样,也称为数据剪枝,理论上可以保证通过优先考虑信息丰富的训练示例来加速监督学习的随机梯度算法,主要受益于方差减少。 对于 DNN,估计每个样本重要性的主要方法是使用梯度范数,并且 [47, 44] 使用不同的近似值来使这些范数的计算变得易于处理。 [70] 进一步加速了类似于早期 LTH 的采样过程,但在数据域中,简单的平均梯度范数或误差 2-范数超过了几个权重初始化可用于在训练的早期阶段识别重要的示例。 最近,[81] 展示了一个令人兴奋的分析理论,即测试误差随数据集大小的缩放可以突破幂缩放定律,并且如果配备优质数据,则至少可以减少到指数缩放修剪指标,并且它采用使用k-均值聚类的自监督指标。 它展示了基于数据重要性采样的更有效的神经缩放法则的有前途的方向。

表格1: 记忆有效训练方法总结。 缩写包括:AMP= 自动混合精度、DP = 数据并行、MP = 模型并行、TP = 张量并行、PP = 管道并行、ACT = 激活压缩训练和 PET = 参数高效调优。
Method Class
Micikevicius et al. [64] AMP
Chen et al. [15] Rematerialization
Herrmann et al. [36] Rematerialization
ZeRO-Offload [75] Offloading
Beaumont et al. [7] Offloading + Rematerization
ZeRO [73] DP+MP+AMP
Megatron-LM [77] DP+TP
GPipe [41] DP+PP
torchgpipe [49] PP+Rematerization
Megatron-LM[66] DP+TP+PP+AMP
Wang et al. [89] FP8 Training
Cambier et al. [11] FP8 Training
Mesa [69] 8-bit ACT
ACTNN [12], GACT [61] 2-bit ACT
[53, 43, 38] Addition-based PET
Bitfit [94], LoRA [39] Reparameterization-based PET

3 内存效率

除了计算负担之外,大型 Transformer 模型 e.g. 的模型规模不断增长,从 BERT [48] 345M 参数模型到GPT-3 拥有 1.75 万亿个参数,是训练的关键瓶颈,因为它们不适合单个设备的内存。 我们首先分析现有模型训练框架的内存消耗,其占用情况为1)模型状态,包括优化器状态(e.g .、Adam 中的动量和方差)、梯度和参数; 2) 激活(我们忽略临时缓冲区和空闲碎片内存,因为它们相对较小)。 我们在表1中总结了内存有效的训练方法。 下面,我们讨论优化内存使用的主要解决方案。

并行性 跨设备并行训练大型 DNN 是满足内存需求的常见做法。 基本上有两种范例:数据并行性(DP),它在不同的设备上分配小批量数据;模型并行性(MP),它在多个工作人员之间分配模型的子图。 对于DP来说,随着可用worker的增加,batch size接近线性缩放。 第 2 节中讨论的大批量训练。 2就是针对这种情况开发的。 然而,很明显DP具有较高的通信/计算效率,但内存效率较差——当模型变大时,单个设备无法存储模型副本,并且梯度的同步通信会阻碍DP的可扩展性。 因此,DP本身只适合训练中小型模型。 为了提高DP的可扩展性,Transformer的一种解决方案是参数共享[51],称为Albert,但它限制了表示能力。 最近,ZeRO [73] 将统一分区策略与 DP 相结合,其中每个数据并行过程仅处理模型状态的一个分区,在混合精度状态下工作。 为了处理非常大的 DNN,人们总是需要利用模型并行性以“垂直”方式在多个加速器上分配不同的层。 尽管MP具有良好的存储效率,但由于跨设备的大量数据传输和PE利用率较差,其通信和计算效率较低。 幸运的是,有两种策略可以进一步提高正交“水平”维度的 MP 效率,包括张量并行(TP)和管道并行(PP)。 TP 将张量操作划分为跨工作层的层,以实现更快的计算和更多的内存节省。 Megatron-LM [77] 针对基于 Transformer 的模型进行了定制,可跨 GPU 切片 MSA 和 FFN,并且在前向和后向传递中只需要一些额外的 All-Reduce 操作,从而允许它们训练模型使用 512 个 GPU 达到 83 亿个参数。 PP方面,最初是在GPipe[41]中提出的,它将输入的mini-batch分割成多个更小的微批次,使得不同的加速器(顺序层跨加速器划分)能够工作在对整个小批量应用单个同步梯度更新之前,同时处理不同的微批量。 然而,它仍然受到管道气泡(加速器空闲时间)的影响,从而降低了效率。 特别是,PyTorch 实现了 torchgpipe [49],它通过检查点执行微批次 PP,允许扩展到大量微批次以最小化气泡开销。

请注意,DP 和 MP 是正交的,因此可以同时使用两者来训练具有更高计算和内存容量的更大模型。 例如,Megatron-LM [66] 和 DeepSpeed [74] 组合张量、管道和数据并行性,将训练扩展到数千个GPU。

量化训练 训练神经网络的标准例程采用全精度(i.e., FP32)。 相比之下,量化训练通过将激活/权重/梯度压缩为低位值(e.g.、FP16 或 INT8)从头开始训练神经网络,但精度较低。 。 之前的工作已经表明,降低训练精度[98, 42]可以加速神经网络训练,并具有良好的性能。 对于 Transformers,最广泛采用的方法是自动混合精度(AMP)训练[64] 具体来说,AMP 以全精度存储权重的主副本以进行更新,而激活、梯度和权重则存储在 FP16 中以进行算术。 与全精度训练相比,AMP能够实现更快的训练/推理速度并减少网络训练期间的内存消耗。 例如,基于批量大小为 64 和图像分辨率 224×224,在 AMP 下的 RTX3090 上训练 DeiT-B [84] 为 2×比全精度训练更快(305 vs. 124 个图像/秒),并且消耗峰值 GPU 内存减少 22%(7.9GB vs. 10.2GB)。 虽然人们普遍认为在不影响模型精度的情况下训练网络至少需要 16 位[64, 21],但 NVIDIA H100 上对 FP8 训练的最新支持在 Transformer 训练上显示出了有希望的结果,其中 FP8 下的训练 DeiT-S 和 GPT [10] 可以与 16 位训练相匹配。 除了同时量化激活/权重/梯度的降低精度训练之外,激活压缩训练 (ACT) [12] 在精确计算前向传递时存储激活的低精度近似副本,这有助于减少训练期间的总体内存消耗。 然后,保存的激活在向后传递中被反量化到原始精度以计算梯度。 最近的工作[69, 61]进一步提出定制ACT以支持内存高效的Transformer训练。

再物化和卸载 Rematerialization,也称为检查点 [15],是一种广泛使用的时空权衡技术,仅存储一部分激活/权重在前向传播期间重新计算其余部分。 [15] 提供了一个在 PyTorch 中实现的简单定期计划222https://pytorch.org/,但它仅适用于同质顺序网络。 [36] 等更高级的方法实现异构网络的最佳检查点333https://gitlab.inria.fr/hiepacs/rotor 卸载而言,它是一种通过GPU和CPU之间的通信,使用CPU内存等外部内存作为GPU内存的扩展,以增加训练时内存容量的技术。 模型状态和激活可以卸载到 CPU,但最佳选择需要最小化与 GPU 之间的通信成本(i.e.,数据移动) ,减少CPU计算并最大限度地节省GPU内存。 一个代表性的工作是 ZeRO-Offload [75],它提供了使用 Adam 优化器针对混合精度训练定制的最佳卸载策略。 它将所有 fp32 模型状态和 fp16 梯度卸载到 CPU 内存上,并在 CPU 上计算 fp32 参数更新。 fp16参数保存在GPU上,前向和后向计算都在GPU上。 为了两全其美,[7]建议联合优化激活卸载和重新实现。

参数高效调整。 以HuggingFace为代表的公共模型动物园包含丰富的预训练模型,可供下载和执行,为降低训练成本做出了巨大贡献。 有效调整这些现成的模型正在成为大幅降低成本的流行训练方法。 作为普通完全微调的强大替代方案,参数高效调整(PET)仅更新少量额外参数,同时冻结预训练模型以显着减轻存储负担,可随动态部署场景扩展而无需存储每种情况都有一个单独的模型实例。 一般的 PET 方法可以分为基于加法的方法和基于重新参数化的方法。 前者将额外的可训练参数附加到预训练模型上,并且仅调整这些参数。 例如,[53, 43]将可训练参数添加到输入空间,[38]将适配器模块添加到MSA和FFN之后的每个Transformer块中两次。 然而,额外的参数在推理过程中引入了额外的计算和内存开销。 为了应对这一挑战,后者建议调整模型[94]中固有的参数或可以重新参数化到模型[39]中的新参数,从而产生不牺牲推理效率。 受到大型语言预训练模型具有低内在维度[1]这一观察的启发,代表作品LoRA[39]将自注意力权重的更新近似为两个低维权重:排名矩阵,可以在推理过程中合并到预训练的权重中。 值得注意的是,最受认可的大语言模型民主化努力之一是斯坦福羊驼[83],它是根据开源LLaMA模型[86]使用微调而成的ChatGPT 生成的 52K 指令跟踪数据。 为了廉价高效地实现它,其变体 Alpaca-LoRA 444https://github.com/tloen/alpaca-lora 进一步采用低阶 LoRA 在客户硬件上启用指令调优 LLaMA,显示训练可以在单个 RTX 4090 上在数小时内完成。

开源框架。 有几种广泛采用的原型用于大规模训练大型 Transformer 模型,其中 Microsoft DeepSpeed555https://github.com/microsoft/DeepSpeed,HPC-AI Tech Colossal-AI666https://github.com/hpcaitech/ColossalAI 和 Nvidia Megatron-LM7 5>77https://github.com/NVIDIA/Megatron-LM 具体来说,DeepSpeed主要基于[74]和ZeRO系列作品[73, 75]实现,Colossal-AI基于[8]构建>,Megatron-LM 实现了 [66] 所有这些都支持混合精度的数据和模型并行性,以及卸载和重新实现等其他通用实践。 更多用于高效分布式训练的库包括但不限于 HuggingFace Transformers888https://github.com/huggingface/transformers、MosaicML Composer999https://github.com/mosaicml/composer、百度 PaddlePaddle101010https://github.com/PaddlePaddle/Paddle, Bytedance Lightseq111111https://github.com/bytedance/lightseq、EleutherAI GPT-NeoX121212https://github.com/EleutherAI/gpt-neox 等。

4硬件/算法协同设计

除了计算和内存负担之外,设计高效的硬件加速器还可以加快 DNN 的训练和推理速度。 具体来说,与中央处理单元(CPU)相比,图形处理单元(GPU)由于并行度高,执行矩阵乘法的能力更强大。 对于专注于特定计算任务的应用程序,专用集成电路(ASIC)具有低功耗和高训练/推理速度的优势。 例如,Google 设计的张量处理单元 (TPU) 的每瓦性能比当代 CPU 和 GPU 高出 3080× [45] 然而,ASIC 不容易重新编程或适应新任务。 相比之下,现场可编程门阵列 (FPGA) 旨在根据需要重新编程以执行不同的功能,并且还可以在最终设计之前用作 ASIC 的原型。 为了进一步优化 DNN,特别是 Transformer 的训练效率,硬件算法协同设计在设计算法时考虑了硬件的约束和能力,这将在下面的小节中介绍。

稀疏矩阵乘法 为了减少 Transformers 的计算开销,稀疏通用矩阵乘法(SpGEMM)涉及将稀疏矩阵与稠密矩阵相乘,利用注意力矩阵的稀疏性来减少计算量。 有几种流行的稀疏矩阵计算库,例如 Intel Math Kernel Library131313https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html CPU 和 cuSPARSE141414https://docs.nvidia.com/cuda/cusparse/ , 尖峰155>1515https://cusplibrary.github.io/ 161616https://developer.nvidia.com/blog/accelerating-inference-with-sparsity-using-ampere-and-tensorrt/[65] 然而,由于不规则的稀疏性,SpGEMM 通常对通用处理器(例如 CPU 和 GPU)的硬件不友好。 为了解决这个问题,需要专门的硬件加速器(例如 FPGA 和 ASIC)来处理数据局部性较差的问题。 例如,OuterSPACE [68] 将矩阵乘法转换为外积过程,并通过将乘法与累加解耦来消除冗余内存访问。 为了充分利用这种减少的优势而不引入显着的开销,OuterSPACE 构建了一个具有可重新配置内存层次结构的自定义加速器,并实现了 7.9 的平均加速× 针对运行英特尔数学核心函数库的 CPU,14.0× 针对运行 CUSP 的 GPU。 此外,为了缓解高稀疏性引起的数据移动瓶颈,ViTCoD [91] 使用可学习的自动编码器将稀疏注意力压缩为更紧凑的表示,并设计编码器和解码器引擎来提高硬件利用率。

硬件感知低精度 降低计算精度会减少内存和计算量,这可以用硬件友好的定点或整数表示形式而不是浮点表示形式来实现。 因此,我们可以使用较低精度的乘法器、加法器和内存块,从而显着改善功耗和加速。 此外,低精度算法可以与剪枝和低秩近似等其他技术相结合,以实现进一步的加速。 例如,Sanger [63] 使用 4 位查询和密钥来计算稀疏注意力矩阵的量化预测。 然后,稀疏注意力掩模被重新排列成结构化块并由可重新配置的硬件处理。 以下工作 DOTA [71] 使用低秩变换和低精度计算来识别注意力中不重要的连接。 通过结合 Token 级并行性和乱序执行,DOTA 比 GPU 实现了 152.6× 加速。

高效的注意力。 除了稀疏矩阵乘法和低精度计算之外,一些开创性的工作还专注于硬件中高效、轻量级的注意力实现[33,34,20] 具体来说,A3 [33] 仅选择那些可能与给定查询具有较高相似度的键,以减少注意力的计算量。 ELSA [34] 根据哈希相似性过滤掉特定查询的不相关键以节省计算量。 凭借高效的硬件加速器,与配备 16GB 内存的 Nvidia V100 GPU 相比,ELSA 的加速速度提高了 58.1×,能源效率提高了三个数量级。 值得注意的是,FlashAttention [20] 提出利用平铺来减少 GPU 高带宽内存 (HBM) 和片上 SRAM 之间的 I/O 通信,这正在成为默认的快速且内存高效的注意力模块为了加速。

5 结论和未来研究

In this survey, we have reviewed several important factors that improve the training of Transformers: 1) appropriate initialization and optimization paradigms that can accelerate the convergence rate with fewer training iterations, resulting in lower computational costs; 2) higher data efficiency by sampling informative training samples towards more efficient neural scaling laws of test error with respect to dataset size; 3) memory-efficient techniques to meet the memory requirements for training large Transformers, which requires jointly optimizing PE utilization, memory and communication footprints across accelerators, using parallelism, low-precision arithmetic, checkpointing and offloading strategies, etc; 4) hardware and algorithm co-design to maximize the training scalability on hardware platforms.

我们最后根据审查的进展强调了几个有希望的方向:(i)联合训练和部署效率优化。 在现实中,我们通常需要将模型部署到具有不同资源限制的不同任务和平台。 因此,非常需要开发新方法来有效训练弹性超级网络,该超级网络支持单次 NAS [13] 或专家混合 [27] 后的多种不同架构配置>,针对多项任务和预算要求; (ii) 在资源有限的边缘设备上训练 [59](e.g.、电池电量和内存容量低),避免频繁传输数据而导致隐私和延迟问题; (iii) 结合词符/模型剪枝、低秩分解、轻量级架构设计、动态神经网络等高效逼近技术,以互补的方式减少模型大小和计算成本; (iv) 评估和比较有效训练方法的标准基准。 拥有这样的基准将加快这些方法的采用,并导致实际成本的降低。

参考

  • [1] A. Aghajanyan, S. Gupta, and L. Zettlemoyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. In ACL, 2021.
  • [2] Z. Allen-Zhu, Y. Li, and Y. Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In NeurIPS, 2019.
  • [3] Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over-parameterization. In ICML, 2019.
  • [4] S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In ICML, 2018.
  • [5] T. Bachlechner, B. P. Majumder, H. Mao, G. Cottrell, and J. McAuley. Rezero is all you need: Fast convergence at large depth. In UAI, 2021.
  • [6] H. Bao, L. Dong, and F. Wei. Beit: Bert pre-training of image transformers. In ICLR, 2022.
  • [7] O. Beaumont, L. Eyraud-Dubois, and A. Shilova. Efficient combination of rematerialization and offloading for training dnns. In NeurIPS, 2021.
  • [8] Z. Bian, H. Liu, B. Wang, H. Huang, Y. Li, C. Wang, F. Cui, and Y. You. Colossal-ai: A unified deep learning system for large-scale parallel training. arXiv preprint arXiv:2110.14883, 2021.
  • [9] L. Bottou, F. E. Curtis, and J. Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 2018.
  • [10] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
  • [11] L. Cambier, A. Bhiwandiwalla, T. Gong, M. Nekuii, O. H. Elibol, and H. Tang. Shifted and squeezed 8-bit floating point format for low-precision training of deep neural networks. In ICLR, 2020.
  • [12] J. Chen, L. Zheng, Z. Yao, et al. Actnn: Reducing training memory footprint via 2-bit activation compressed training. In ICML, 2021.
  • [13] M. Chen, H. Peng, J. Fu, and H. Ling. Autoformer: Searching transformers for visual recognition. In ICCV, 2021.
  • [14] T. Chen, Y. Cheng, Z. Gan, L. Yuan, L. Zhang, and Z. Wang. Chasing sparsity in vision transformers: An end-to-end exploration. In NeurIPS, 2021.
  • [15] T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.
  • [16] X. Chen, Y. Cheng, S. Wang, Z. Gan, Z. Wang, and J. Liu. Earlybert: Efficient bert training via early-bird lottery tickets. In ACL, 2021.
  • [17] X. Chen, C.-J. Hsieh, and B. Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In ICLR, 2022.
  • [18] X. Chen, C. Liang, D. Huang, E. Real, K. Wang, Y. Liu, H. Pham, X. Dong, T. Luong, C.-J. Hsieh, et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675, 2023.
  • [19] J.-B. Cordonnier, A. Loukas, and M. Jaggi. On the relationship between self-attention and convolutional layers. In ICLR, 2020.
  • [20] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In NeurIPS, 2022.
  • [21] D. Das, N. Mellempudi, D. Mudigere, D. D. Kalamkar, et al. Mixed precision training of convolutional neural networks using integer operations. In ICLR, 2018.
  • [22] S. De and S. Smith. Batch normalization biases residual blocks towards the identity function in deep networks. In NeurIPS, 2020.
  • [23] J. Du, H. Yan, J. Feng, J. T. Zhou, L. Zhen, et al. Efficient sharpness-aware minimization for improved training of neural networks. In ICLR, 2022.
  • [24] J. Du, D. Zhou, J. Feng, V. Y. Tan, and J. T. Zhou. Sharpness-aware training for free. In NeurIPS, 2022.
  • [25] S. d’Ascoli, H. Touvron, M. L. Leavitt, A. S. Morcos, G. Biroli, and L. Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. In ICML, 2021.
  • [26] U. Evci, T. Gale, J. Menick, P. S. Castro, and E. Elsen. Rigging the lottery: Making all tickets winners. In ICML, 2020.
  • [27] W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. JMLR, 2022.
  • [28] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In ICLR, 2021.
  • [29] J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In ICLR, 2019.
  • [30] L. Gong, D. He, Z. Li, T. Qin, L. Wang, and T. Liu. Efficient training of bert by progressively stacking. In ICML, 2019.
  • [31] P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, et al. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677, 2017.
  • [32] X. Gu, L. Liu, H. Yu, J. Li, C. Chen, and J. Han. On the transformer growth for progressive bert training. In NAACL, 2021.
  • [33] T. J. Ham, S. J. Jung, S. Kim, Y. H. Oh, Y. Park, et al. A^ 3: Accelerating attention mechanisms in neural networks with approximation. In HPCA, 2020.
  • [34] T. J. Ham, Y. Lee, S. H. Seo, S. Kim, H. Choi, et al. Elsa: Hardware-software co-design for efficient, lightweight self-attention mechanism in neural networks. In ISCA, 2021.
  • [35] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In CVPR, 2022.
  • [36] J. Herrmann, O. Beaumont, L. Eyraud-Dubois, J. Hermann, A. Joly, and A. Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory. arXiv preprint arXiv:1911.13214, 2019.
  • [37] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
  • [38] N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, et al. Parameter-efficient transfer learning for nlp. In ICML, 2019.
  • [39] E. J. Hu, yelong shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. In ICLR, 2022.
  • [40] X. S. Huang, F. Perez, J. Ba, and M. Volkovs. Improving transformer optimization through better initialization. In ICML, 2020.
  • [41] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In NeurIPS, 2019.
  • [42] I. Hubara, M. Courbariaux, D. Soudry, R. El-Yaniv, and Y. Bengio. Quantized neural networks: Training neural networks with low precision weights and activations. JMLR, 2017.
  • [43] M. Jia, L. Tang, B.-C. Chen, C. Cardie, S. Belongie, B. Hariharan, and S.-N. Lim. Visual prompt tuning. In ECCV, 2022.
  • [44] T. B. Johnson and C. Guestrin. Training deep models faster with robust, approximate importance sampling. In NeurIPS, 2018.
  • [45] N. P. Jouppi, C. Young, N. Patil, D. Patterson, et al. In-datacenter performance analysis of a tensor processing unit. In ISCA, 2017.
  • [46] J. Kaddour. Stop wasting my time! saving days of imagenet and bert training with latest weight averaging. In NeurIPS workshop, 2022.
  • [47] A. Katharopoulos and F. Fleuret. Not all samples are created equal: Deep learning with importance sampling. In ICML, 2018.
  • [48] J. D. M.-W. C. Kenton and L. K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2019.
  • [49] C. Kim, H. Lee, M. Jeong, W. Baek, B. Yoon, et al. torchgpipe: On-the-fly pipeline parallelism for training giant models. arXiv preprint arXiv:2004.09910, 2020.
  • [50] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
  • [51] Z. Lan, M. Chen, S. Goodman, K. Gimpel, P. Sharma, and R. Soricut. Albert: A lite bert for self-supervised learning of language representations. In ICLR, 2020.
  • [52] N. Lee, T. Ajanthan, and P. Torr. Snip: Single-shot network pruning based on connection sensitivity. In ICLR, 2019.
  • [53] B. Lester, R. Al-Rfou, and N. Constant. The power of scale for parameter-efficient prompt tuning. In EMNLP, 2021.
  • [54] C. Li, B. Zhuang, G. Wang, X. Liang, X. Chang, and Y. Yang. Automated progressive learning for efficient training of vision transformers. In CVPR, 2022.
  • [55] Q. Li, Y. Zhou, Y. Liang, and P. K. Varshney. Convergence analysis of proximal gradient with momentum for nonconvex optimization. In ICLR, 2017.
  • [56] Y. Li, H. Fan, R. Hu, C. Feichtenhofer, and K. He. Scaling language-image pre-training via masking. arXiv preprint arXiv:2212.00794, 2022.
  • [57] Y. Li and Y. Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In NeurIPS, 2018.
  • [58] Z. Li, E. Wallace, S. Shen, K. Lin, et al. Train big, then compress: Rethinking model size for efficient training and inference of transformers. In ICML, 2020.
  • [59] J. Lin, W.-M. Chen, Y. Lin, C. Gan, S. Han, et al. Mcunet: Tiny deep learning on iot devices. In NeurIPS, 2020.
  • [60] Z. Lin, H. Li, and C. Fang. Accelerated optimization for machine learning. Nature Singapore: Springer, 2020.
  • [61] X. Liu, L. Zheng, D. Wang, Y. Cen, W. Chen, X. Han, J. Chen, et al. GACT: activation compressed training for generic network architectures. In ICML, 2022.
  • [62] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In ICLR, 2019.
  • [63] L. Lu, Y. Jin, H. Bi, Z. Luo, P. Li, T. Wang, and Y. Liang. Sanger: A co-design framework for enabling sparse attention using reconfigurable architecture. In MICRO, 2021.
  • [64] P. Micikevicius, S. Narang, J. Alben, G. F. Diamos, E. Elsen, et al. Mixed precision training. In ICLR, 2018.
  • [65] A. Mishra, J. A. Latorre, J. Pool, D. Stosic, D. Stosic, et al. Accelerating sparse deep neural networks. arXiv preprint arXiv:2104.08378, 2021.
  • [66] D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley, M. Patwary, et al. Efficient large-scale language model training on gpu clusters using megatron-lm. In SC, 2021.
  • [67] Y. Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o (1/k^ 2). In Doklady ANSSSR, volume 269, pages 543–547, 1983.
  • [68] S. Pal, J. Beaumont, D.-H. Park, A. Amarnath, S. Feng, et al. Outerspace: An outer product based sparse matrix multiplication accelerator. In HPCA, 2018.
  • [69] Z. Pan, P. Chen, H. He, J. Liu, J. Cai, and B. Zhuang. Mesa: A memory-saving training framework for transformers. arXiv preprint arXiv:2111.11124, 2021.
  • [70] M. Paul, S. Ganguli, and G. K. Dziugaite. Deep learning on a data diet: Finding important examples early in training. In NeurIPS, 2021.
  • [71] Z. Qu, L. Liu, F. Tu, Z. Chen, Y. Ding, and Y. Xie. Dota: detect and omit weak attentions for scalable transformer acceleration. In ASPLOS, 2022.
  • [72] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
  • [73] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. Zero: Memory optimizations toward training trillion parameter models. In SC, 2020.
  • [74] J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In KDD, 2020.
  • [75] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, et al. Zero-offload: Democratizing billion-scale model training. In USENIX ATC, 2021.
  • [76] R. Schwartz, J. Dodge, N. A. Smith, and O. Etzioni. Green ai. Communications of the ACM, 2020.
  • [77] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
  • [78] S. Smith, M. Patwary, B. Norick, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990, 2022.
  • [79] D. So, Q. Le, and C. Liang. The evolved transformer. In ICML, pages 5877–5886, 2019.
  • [80] K. Song, X. Tan, T. Qin, J. Lu, and T.-Y. Liu. Mass: Masked sequence to sequence pre-training for language generation. In ICML, 2019.
  • [81] B. Sorscher, R. Geirhos, S. Shekhar, S. Ganguli, and A. S. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In NeurIPS, 2022.
  • [82] E. Strubell, A. Ganesh, and A. McCallum. Energy and policy considerations for deep learning in nlp. In ACM, 2019.
  • [83] R. Taori, I. Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca, 2023.
  • [84] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. In ICML, 2021.
  • [85] H. Touvron, M. Cord, A. Sablayrolles, G. Synnaeve, and H. Jégou. Going deeper with image transformers. In ICCV, 2021.
  • [86] H. Touvron, T. Lavril, G. Izacard, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  • [87] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. In NeurIPS, 2017.
  • [88] C. Wang, G. Zhang, and R. Grosse. Picking winning tickets before training by preserving gradient flow. In ICLR, 2020.
  • [89] N. Wang, J. Choi, D. Brand, C.-Y. Chen, and K. Gopalakrishnan. Training deep neural networks with 8-bit floating point numbers. In NeurIPS, 2018.
  • [90] H. You, C. Li, P. Xu, Y. Fu, Y. Wang, X. Chen, R. G. Baraniuk, et al. Drawing early-bird tickets: Towards more efficient training of deep networks. In ICLR, 2020.
  • [91] H. You, Z. Sun, H. Shi, Z. Yu, Y. Zhao, Y. Zhang, C. Li, B. Li, and Y. Lin. Vitcod: Vision transformer acceleration via dedicated algorithm and accelerator co-design. In HPCA, 2023.
  • [92] Y. You, I. Gitman, and B. Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv:1708.03888, 2017.
  • [93] Y. You, J. Li, S. Reddi, J. Hseu, S. Kumar, S. Bhojanapalli, X. Song, J. Demmel, K. Keutzer, and C.-J. Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In ICLR, 2019.
  • [94] E. B. Zaken, Y. Goldberg, and S. Ravfogel. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. In ACL, 2022.
  • [95] H. Zhang, Y. N. Dauphin, and T. Ma. Fixup initialization: Residual learning without normalization. In ICLR, 2019.
  • [96] J. Zhang, S. P. Karimireddy, A. Veit, S. Kim, et al. Why adam beats sgd for attention models. In OpenReview, 2019.
  • [97] M. Zhang and Y. He. Accelerating training of transformer-based language models with progressive layer dropping. In NeurIPS, 2020.
  • [98] S. Zhou, Y. Wu, Z. Ni, X. Zhou, H. Wen, and Y. Zou. Dorefa-net: Training low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606.06160, 2016.