烦人的内存墙
1 小时前 / 阅读约16分钟
来源:36kr
训练AI模型计算成本激增,模型参数规模快速增长,硬件峰值浮点运算能力提升,但DRAM和互连带宽增长滞后。内存成为大型AI模型主要瓶颈,需重新思考模型训练、部署、设计及硬件设计。

前所未有的无监督训练数据的可用性,以及神经网络的扩展规律,导致用于服务/训练低层逻辑模型(LLM)的模型规模和计算需求出现了前所未有的激增。然而,主要的性能瓶颈正日益转移到内存带宽上。

过去20年,服务器硬件的峰值浮点运算能力(FLOPS)以每两年3倍的速度增长,超过了DRAM和互连带宽的增长速度,后两者分别仅以每两年1.6倍和1.4倍的速度增长。这种差距使得内存而非计算成为人工智能应用(尤其是服务应用)的主要瓶颈。

本文分析了编码器和解码器Transformer模型,并展示了内存带宽如何成为解码器模型的主要瓶颈。我们提出重新设计模型架构、训练和部署策略,以克服这一内存限制。

引言

近年来,训练大型语言模型 (LLM) 所需的计算量以每两年 750 倍的速度增长。这种指数级增长趋势是人工智能加速器发展的主要驱动力,这些加速器致力于提升硬件的峰值计算能力,但往往以牺牲其他部分(例如内存层次结构)的简化为代价。

然而,这些趋势忽略了训练和服务人工智能模型过程中一个新兴的挑战:内存和通信瓶颈。事实上,许多人工智能应用的瓶颈并非计算能力,而是芯片内部/芯片间以及与人工智能加速器之间/芯片间的通信。这并非新现象,过去的一些研究已经观察到并警告了这个问题。最早对此的观察可以追溯到 1990 年,当时 Ousterhout 在分析影响操作系统性能的因素后得出以下结论:

“第一个与硬件相关的问题是内存带宽:基准测试表明,内存带宽跟不上 CPU 的速度……如果未来的机器内存带宽没有显著提高,某些类型的应用程序可能会受到内存性能的限制。”

1995 年晚些时候,William Wulf 和 Sally Mckee 进一步呼应了这一预测,并创造了“内存墙”(memory wall)一词。他们的论证遵循一个简单而精妙的逻辑。完成一项操作所需的时间取决于我们执行算术运算的速度以及我们向硬件运算单元提供数据的速度。¹在最简单的情况下,数据要么在缓存中可用,要么需要从 DRAM 中获取。基于此假设,即使 80% 的数据可以直接在缓存中获取,只有 20% 的数据需要从 DRAM 中获取,如果从 DRAM 中获取这 20% 的cache-miss数据需要超过 5 个时钟周期,那么完成操作所需的时间将完全受限于 DRAM。这意味着,无论硬件每秒执行的算术运算速度有多快,问题都将完全受限于 DRAM 带宽。

他们预测,计算速度的提升速度与数据获取速度的提升速度出现分歧,将会导致“内存墙”问题。基于此,他们得出以下结论:

“两者的性能都在呈指数级增长,但微处理器的性能增长指数远大于DRAM。而且,不同指数级增长之间的差异也在呈指数级增长。”

后来的几篇文献也报道了类似的观察结果。

在这项工作中,我们通过研究更新的数据重新审视了这一趋势,尤其关注用于训练人工智能模型的硬件,以及用于训练/运行这些模型的计算特性。30 年过去了,上述观察和预测依然完全正确。尽管内存技术取得了诸多创新,但趋势表明,“内存墙”正日益成为一系列人工智能任务的主要瓶颈。

我们首先分析自 1998 年 Yann Lecun 在 MNIST 数据集上训练出著名的 Lenet-5 模型以来,服务器级 AI 硬件的峰值计算能力发生了怎样的变化。我们可以看到,过去 20 年间,硬件的峰值计算能力增长了 6 万倍,而 DRAM 的峰值计算能力仅增长了 100 倍,互连带宽也仅增长了 30 倍。

内存墙问题涉及内存容量和带宽的限制,以及内存传输的延迟(延迟的改进难度甚至比带宽更大)。这涉及到不同层级的内存数据传输。例如,计算逻辑与片上内存之间的数据传输,计算逻辑与 DRAM 内存之间的数据传输,以及不同插槽上不同处理器之间的数据传输。在所有这些情况下,数据传输的容量和速度都显著落后于硬件的计算能力。

现在,如果我们研究近期人工智能模型(尤其是LLM)的发展趋势,我们会发现,受神经网络扩展定律的启发,从业者们以前所未有的速度扩展了训练这些模型所需的数据量、模型规模和计算资源。尽管在2018年至2022年期间,训练这些模型所需的计算量/浮点运算次数(FLOPs)增长了750倍/两年(见图2),但计算资源未必是瓶颈,尤其是在模型服务方面。

首先,LLM的规模在此期间以410倍/两年的速度增长,超过了单个芯片的可用内存。人们或许会希望,我们可以通过将训练/服务扩展到多个加速器来实现分布式内存并行,从而避免单个硬件内存容量和带宽的限制。然而,将工作分配到多个进程也可能面临内存墙问题:神经网络 (NN) 加速器之间数据传输的通信瓶颈,其速度甚至比片上数据传输更慢、效率更低。与单系统内存的情况类似,我们尚未能克服扩展网络带宽的技术挑战。

其次,即使模型能够装入单个芯片,芯片内部寄存器、L2 缓存、全局内存等之间的数据传输仍然日益成为瓶颈。得益于专用计算单元(例如 Tensor 核心)的最新进展,大量计算的算术运算可以在几个周期内完成。因此,为了始终保持这些算术单元的利用率,需要快速地向其提供大量数据,而这正是芯片内存带宽成为瓶颈的地方。

如图 1 所示,过去 20 年间,服务器硬件的峰值浮点运算能力 (FLOPS) 以每两年 3.0 倍的速度增长,远超 DRAM 和互连带宽的增长速度(后两者分别仅以每两年 1.6 倍和 1.4 倍的速度增长)。这种差距使得内存而非计算能力日益成为瓶颈,即使在模型能够容纳于单个芯片的情况下也是如此。

接下来,我们将对 Transformer 进行详细的案例研究,通过考察当今常用的模型,更好地展示浮点运算能力 (FLOPS)、内存操作 (MOP) 和端到端运行时间之间的相互作用。

案例研究

在本节中,我们首先概述Transformer推理的运行时特性和性能瓶颈。我们考察Transformer架构的两种不同变体:编码器架构(例如BERT),它并发处理所有tokens;以及解码器架构(例如GPT),它以自回归的方式运行,每次迭代处理并生成一个token。

A. 算术强度(Arithmetic Intensity)

衡量性能瓶颈的一种常用方法是计算仅计算 Transformer 编码器和仅计算解码器模型所需的总 FLOP 数。然而,单独使用这一指标可能会产生很大的误导。因此,我们需要研究相关操作的算术强度。算术强度是指从内存加载的每个字节可以执行的 FLOP 数。它可以通过将总 FLOP 数除以访问的总字节数(也称为 MOP,即内存操作数)来计算。

为了说明考虑算术强度的重要性,我们研究了 BERT-Base 和 BERT-Large 以及 GPT-2 。前两者是编码器模型,其推理涉及矩阵-矩阵运算;后者是解码器/自回归模型,其推理涉及重复的矩阵-向量乘法。

B. 性能分析

为了分析 Transformer 工作负载在通用硬件上的瓶颈,我们在 Intel Gold 6242 CPU 上对 Transformer 推理进行了性能分析。图 3 显示了不同序列长度下这些模型的总 FLOPs、MOPs、算术强度和最终延迟。显然,对于每种序列长度,GPT-2 的延迟都明显长于 BERT-Base 或 BERT-Large,即使 BERT-Base 和 GPT-2 的模型配置和端到端 FLOPs 基本相同(如图 3a 所示)。这是由于 GPT 的自回归推理固有的矩阵向量运算具有更高的内存操作量和更低的算术强度(见图 3c)。算术强度更高的模型在相同甚至更高的 FLOPs 下可以比算术强度更低的模型运行得更快。这清楚地表明,内存墙如何成为解码器模型(在小批量大小下)的主要瓶颈,而不是计算瓶颈。

突破内存瓶颈的有希望的解决方案

“指数增长不可能永远持续下去”[,即使对于大型超大规模公司而言,以 410 倍/2 年的速度延迟指数级扩展也并非长久之计。此外,计算能力和带宽能力之间的差距日益扩大,这将很快使训练大型模型变得极具挑战性,因为成本将呈指数级增长。

为了继续创新并突破内存瓶颈,我们需要重新思考人工智能模型的设计。这里存在几个问题。

首先,当前设计人工智能模型的方法大多是临时性的,或者只涉及非常简单的扩展规则。例如,最近出现的大型 Transformer 模型大多只是对最初 BERT 模型中提出的几乎相同的基础架构进行缩放。

其次,我们需要设计更高效的人工智能模型训练方法。目前的神经网络需要大量的训练数据和数十万次的迭代才能学习,效率非常低下。有人可能会指出,这与人脑的学习方式也不同,人脑通常每个概念/类别只需要极少的例子。

第三,当前的优化和训练方法需要大量的超参数调优(例如学习率、动量等),这通常会导致数百次的试错,才能找到合适的超参数设置来成功训练模型。因此,图 2 (b) 中报告的训练成本只是实际开销的下限,而实际成本通常要高得多。

第四,最先进模型的庞大规模使得它们难以部署用于推理。这不仅限于 GPT-3 等模型。事实上,部署超大规模公司使用的大型推荐系统也是一项重大挑战。

最后,硬件加速器的设计主要集中在提高峰值计算能力,而对改善内存密集型工作负载的关注相对较少。这使得训练大型模型以及探索替代模型(例如图神经网络)都变得困难,因为图神经网络通常受限于带宽,无法有效地利用当前的加速器。

以上所有问题都是机器学习领域的根本性难题。本文将简要讨论近期针对后三项问题的研究(包括我们自身的研究)。

A. 高效训练算法

训练神经网络模型的主要挑战之一是需要进行大量的超参数调优。这包括确定学习率、退火策略、收敛所需的迭代次数等等。这会给训练最先进的模型带来(显著的)额外开销。许多此类问题源于用于训练的一阶随机梯度下降(SGD)方法。虽然SGD的变体易于实现,但它们对超参数调优的鲁棒性较差,并且对于未知最佳超参数集的新模型,很难进行调优。一种很有前景的解决方法是使用二阶随机优化方法。这些方法通常对超参数调优更具鲁棒性,并且可以达到最先进的水平。然而,当前方法的内存占用量是其他方法的3-4倍,这需要加以解决。微软的 Zero 框架在这方面展现出良好的前景,它通过移除/分片冗余的优化状态变量,展示了如何在内存容量不变的情况下训练规模扩大 8 倍的模型。如果能够解决这些高阶方法的开销问题,那么它们就能显著降低训练大型模型的总成本。

另一种很有前景的方法是减少内存占用并提高优化算法的数据局部性,但代价是需要执行更多计算。数值线性代数中一个显著的例子是通信避免算法族。神经网络训练中内存优化的一个例子是重物化,即在前向传播过程中只存储/检查点一部分激活值,而不是保存所有激活值。如图4所示,这减少了特征图的内存占用。其余的激活值可以在需要时重新计算。尽管这会增加计算量,但只需增加20%的计算量,就可以显著减少高达5倍的内存占用。这还可以让开发者在单芯片内存上训练大型模型,而无需使用分布式训练。分布式训练通常难以设置(大型超大规模数据中心公司除外),而且对于非专业开发者来说也难以调试。

有趣的是,传统趋势表明,新的神经网络模型架构的开发是基于研究人员在单个芯片内可访问的资源,而不是使用复杂的分布式内存方法(见图 4)。当然,也有许多反例,例如大型超标量计算机拥有专门的团队来支持研究人员部署大型模型,但当我们考虑整个社区时,这类例子就比较有限了。事实上,即使是最新的低级内存模型,也常常需要投入大量精力来压缩模型,使其能够适应单个系统,从而使更广泛的研究人员群体能够访问这些模型。

另一个重要的解决方案是设计对低精度训练具有鲁棒性的优化算法。事实上,人工智能加速器领域的一项重大突破是使用半精度 (FP16) 运算,而非单精度运算 。这使得硬件计算能力提升了 10 倍以上。然而,在不降低精度的情况下,利用当前的优化方法进一步降低精度(从半精度到 INT8)一直是一个挑战。一个近期很有前景的趋势是混合使用 FP8 和 FP16(甚至最近的 FP4)。该领域的算法创新无疑将使我们能够更有效地利用硬件,并允许芯片的更多区域用于提升内存性能(这通常被称为内存间隙惩罚 )。

B. 高效部署

部署最新的最先进模型或大型推荐系统极具挑战性,因为它们需要分布式内存部署来进行推理。一个很有前景的解决方案是压缩这些模型以进行推理,可以通过降低精度(即量化)、移除(即剪枝)其冗余参数或设计小型语言模型来实现。

第一种方法是量化,它可以应用于训练和/或推理阶段。虽然将训练精度降低到远低于 FP16 非常具有挑战性,但可以使用超低精度进行推理。利用现有方法,将推理精度量化到 INT4 相对容易,且对准确率的影响极小。这可以使模型占用空间和延迟最多降低 8 倍。然而,使用低于 INT4 的精度进行推理更具挑战性,目前这是一个非常活跃的研究领域。

第二种方法是剪枝,它可以完全移除/剪枝模型中的冗余参数。利用现有方法,可以剪枝高达 30% 的结构化稀疏神经元,以及高达 80% 的非结构化稀疏神经元,且对准确率的影响极小。然而,突破这个限制非常具有挑战性,并且通常会导致严重的准确率下降。解决这个问题仍是一个开放性问题。

第三种方法,即小型语言模型,有望开辟全新的领域,并推动人工智能的广泛应用。有趣的是,自 2017 年 Transformer 模型问世以来,用于语言学习模型 (LLM) 的模型本身并没有改变。目前行之有效的方法是扩展数据规模和模型大小,这促成了这些模型的“涌现能力” 。然而,近期关于小型语言模型的研究已展现出令人鼓舞的成果。如果一个模型能够完全集成到芯片上,那么速度提升和能耗降低都将达到几个数量级。

C. 重新思考人工智能加速器的设计

同时提升芯片的内存带宽和峰值计算能力面临着根本性的挑战。然而,牺牲峰值计算能力可以换取更好的计算/带宽权衡。事实上,CPU架构已经集成了一个经过充分优化的缓存层次结构。这就是为什么CPU在带宽受限问题上比GPU性能更优的原因。

这类问题包括大型推荐问题。然而,目前CPU面临的主要挑战是其峰值计算能力(即FLOPS)比GPU或TPU等人工智能加速器低一个数量级。造成这种情况的原因之一是,人工智能加速器的设计主要目标是实现最大的峰值计算能力。这通常需要移除缓存层次结构等组件,以增加计算逻辑。我们可以设想一种介于这两种极端情况之间的替代架构,最好具有更高效的缓存,更重要的是,具有更高容量的DRAM(可能是一个具有不同带宽的DRAM层次结构)。后者对于缓解分布式内存通信瓶颈可能非常有帮助。

结论

近年来,训练最先进的自然语言处理(NLP)Transformer模型的计算成本以每两年750倍的速度增长,模型参数规模以每两年410倍的速度增长。相比之下,硬件峰值浮点运算能力(FLOPS)以每两年3.0倍的速度增长,而DRAM和互连带宽的增长速度却日益滞后,分别以每两年1.6倍和1.4倍的速度增长。

为了更直观地理解这些数字,我们可以对比一下:过去20年,硬件峰值浮点运算能力增长了6万倍,而同期DRAM和互连带宽的增长速度分别仅为100倍和30倍。按照这种趋势,内存——尤其是芯片内/芯片间内存传输——将很快成为大型人工智能模型的主要瓶颈。

因此,我们需要重新思考人工智能模型的训练、部署和设计,以及如何设计人工智能硬件来应对日益严峻的内存墙挑战。