AI模型训练瓶颈：内存带宽的挑战与解决方案

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

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

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

引言

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

然而，这些趋势忽略了训练和服务人工智能模型过程中一个新兴的挑战：内存和通信瓶颈。事实上，许多人工智能应用的瓶颈并非计算能力，而是芯片内部/芯片间以及与人工智能加速器之间/芯片间的通信。这种瓶颈被称为“内存墙”。

最早对此的观察可以追溯到 1990 年，当时Ousterhout在分析影响操作系统性能的因素后得出以下结论：“第一个与硬件相关的问题是内存带宽。”

1995年晚些时候，William Wulf和Sally Mckee进一步呼应了这一预测，并创造了“内存墙”（memory wall）一词。他们论证了数据获取速度对完成操作所需时间的影响，并预测了计算速度提升与数据获取速度提升之间的分歧将引发“内存墙”问题。

这项研究重新审视了这一趋势，关注用于训练人工智能模型的硬件及其计算特性。尽管内存技术取得了诸多创新，但趋势表明，“内存墙”正日益成为一系列人工智能任务的主要瓶颈。

通过分析自1998年以来服务器级AI硬件的峰值计算能力变化，我们发现硬件的峰值计算能力增长了6万倍，而DRAM的峰值计算能力仅增长了100倍，互连带宽也仅增长了30倍。

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

案例研究

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

A. 算术强度（Arithmetic Intensity）

衡量性能瓶颈的一种方法是计算仅Transformer编码器和解码器模型所需的总FLOP数。然而，单独使用这一指标可能会产生误导。因此，我们需要研究相关操作的算术强度。算术强度是指从内存加载的每个字节可以执行的FLOP数。

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

B. 性能分析

为了分析Transformer工作负载在通用硬件上的瓶颈，我们在Intel Gold 6242 CPU上对Transformer推理进行了性能分析。图3显示了不同序列长度下这些模型的总FLOPs、MOPs、算术强度和最终延迟。显然，对于每种序列长度，GPT-2的延迟都明显长于BERT-Base或BERT-Large。

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

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

为了继续创新并突破内存瓶颈，我们需要重新思考人工智能模型的设计、训练方法以及部署策略。

A. 高效训练算法

训练神经网络模型的主要挑战之一是需要进行大量的超参数调优。这包括确定学习率、退火策略、收敛所需的迭代次数等。这会给训练最先进的模型带来额外开销。一种很有前景的解决方法是使用二阶随机优化方法。

B. 高效部署

部署最新的最先进模型或大型推荐系统极具挑战性。一个很有前景的解决方案是压缩这些模型以进行推理。

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

同时提升芯片的内存带宽和峰值计算能力面临着根本性的挑战。然而，牺牲峰值计算能力可以换取更好的计算/带宽权衡。我们可以设想一种介于CPU和人工智能加速器之间的替代架构。

结论

近年来，训练最先进的自然语言处理（NLP）Transformer模型的计算成本以惊人的速度增长。为了更直观地理解这些数字，我们可以对比一下：过去20年，硬件峰值浮点运算能力增长了6万倍，而同期DRAM和互连带宽的增长速度分别仅为100倍和30倍。

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