优化大型语言模型（LLM）的推理过程及挑战解决方案

原文：优化大型语言模型（LLM）的推理过程及挑战解决方案 | ATYUN.COM 官网-人工智能教程资讯全方位服务平台

2023年11月20日 由 daydream 发表4940

通过堆叠transformer层来创建大型模型可以在各种语言任务上实现更好的准确性、小样本学习能力，甚至接近人类的新兴能力。然而，训练这些基础模型的成本很高，在推理过程中可能需要大量的存储和计算资源（持续成本）。当今最流行的大型语言模型（LLM）在规模上可以达到数千亿到数百亿个参数的大小，并且根据使用情况，可能需要处理长输入（或上下文），这也可能增加成本。

本文讨论了LLM推理中最紧迫的挑战，以及一些实际解决方案。读者应该具有对transformer架构和注意机制的基本了解。

了解LLM推理

大多数流行的解码器LLM（例如GPT-3）都是预训练的，采用因果建模目标，基本上是下一个词的预测器。这些LLM将一系列标记作为输入，并自回归地生成下一个标记，直到满足停止条件（例如标记数量限制或停用词列表）或生成特殊标记标记生成结束。这个过程包括两个阶段：预填充阶段和解码阶段。

请注意，标记是模型处理的语言的原子部分。一个标记大约等于四个英文字符。在输入模型之前，将自然语言中的所有输入转换为标记。

预填充阶段或处理输入

在预填充阶段，LLM处理输入标记以计算中间状态（键和值），用于生成“第一个”新标记。每个新标记依赖于所有先前的标记，但由于已知输入的完整范围，在高层次上，这是一个高度并行化的矩阵-矩阵操作。它可以充分利用GPU。

解码阶段或生成输出

在解码阶段，LLM自回归地一次生成一个输出标记，直到满足停止条件。每个顺序输出标记都需要知道之前迭代的全部输出状态（键和值）。这类似于矩阵-向量操作，与预填充阶段相比，GPU的计算能力利用率不高。数据（权重、键、值、激活）从内存传输到GPU的速度主导延迟，而不是计算实际发生的速度。换句话说，这是一个受内存限制的操作。

本文介绍的许多推理挑战和相应的解决方案涉及解码阶段的优化：高效的注意力模块、有效管理键和值等。

不同的LLM可能使用不同的分词器，因此，比较它们之间的输出标记可能并不直观。当比较推理吞吐量时，即使两个LLM的每秒输出的标记数量相似，如果它们使用不同的分词器，它们也可能并不等效。这是因为相应的标记可能表示不同数量的字符。

批处理

提高GPU利用率和吞吐量最简单的方法是使用批处理。由于多个请求使用相同的模型，权重的内存成本会分散。将较大的批次一次性传输到GPU以进行全体处理将更充分地利用可用的计算资源。

然而，批处理大小只能增加到一定限度，超过这个限度可能会导致内存溢出。要更好地理解为什么会发生这种情况，需要查看键-值（KV）缓存和LLM内存需求。

传统的批处理（也称为静态批处理）是次优的。这是因为对于批处理中的每个请求，LLM可能会生成不同数量的完成标记，因此它们具有不同的执行时间。因此，所有批处理中的请求必须等到最长的请求完成后才能开始执行，而如果生成长度差异较大，则可能会加剧此情况。有一些方法可以缓解这个问题，例如在飞行中的批处理，稍后将对此进行讨论。

键值缓存

解码阶段的常见优化之一是KV缓存。解码阶段在每个时间步生成一个标记，但是每个标记都依赖于所有先前标记的键和值张量（包括在预填充时计算的输入标记的KV张量以及在当前时间步计算的任何新KV张量）。

为了避免在每个时间步上为所有标记重新计算所有这些张量，可以将它们缓存在GPU内存中。每次迭代时，当计算出新元素时，它们只是添加到正在运行的缓存中，以便在下一次迭代中使用。在一些实现中，模型的每一层都有一个KV缓存。

LLM内存需求

GPU上LLM的内存需求的两个主要贡献因素是模型权重和KV缓存。

• 模型权重：内存被模型参数占用。例如，加载了70亿个参数的模型（例如Llama 2 7B），以16位精度（FP16或BF16）加载，则大约需要7B * sizeof（FP16）〜= 14 GB的内存。

• KV缓存：内存被缓存自注意张量占用，以避免冗余计算。

使用批处理时，批处理中每个请求的KV缓存仍然必须单独分配，并且可能具有较大的内存占用。下面的公式描述了针对当今大多数常见LLM体系结构适用的每个标记的KV缓存大小。

每个标记的KV缓存大小（以字节为单位）= 2 （num_layers）（num_heads * dim_head）*精度（以字节为单位）

前面的2个因素用于K和V矩阵。通常，（num_heads * dim_head）的值与transformer的隐藏大小（或模型的维度d_model）相同。这些模型属性通常在模型配置文件或相关配置文件中找到。

这个内存大小对于输入序列的每个标记，在输入批次中都是必需的。假设使用半精度，KV缓存的总大小由下面的公式给出。

KV缓存的总大小（以字节为单位）=（batch_size）（sequence_length） 2 （num_layers）（hidden_size）*sizeof（FP16）

例如，对于一个使用16位精度的Llama 2 7B模型和批量大小为1的情况，KV缓存的大小将是1 * 4096 * 2 * 32 * 4096 * 2字节，约为2 GB。

高效管理这个KV缓存是一个具有挑战性的任务。随着批处理大小和序列长度的线性增长，内存需求可能快速增加。因此，它限制了可以提供的吞吐量，并对长上下文输入提出了挑战。这是本文介绍的几种优化背后的动机。

使用模型并行化扩展LLM

减少模型权重在每个设备的内存占用的一种方法是将模型分布在多个GPU上。分布内存和计算占用可以运行更大的模型或更大的输入批处理。模型并行化是训练或推理需要超过单个设备可用内存的模型的必要手段，并使训练时间和推理度量（延迟或吞吐量）适用于某些用例。基于模型权重如何分割，可以采用多种模型并行化方式。

请注意，数据并行处理也是在讨论的其他方法中经常提到的技术。在数据并行处理中，模型的权重会复制到多个设备上，并且输入的全局批次大小会分片到每个设备上的微批次中。它通过处理更大的批次来降低整体执行时间。然而，这是一种训练时的优化，在推理过程中不那么重要。

流水线并行

流水线并行将模型（垂直）分成块，其中每个块包含在单独的设备上执行的一组层。图2a是四路流水线并行的示例，其中模型被顺序分割，每个设备上执行四分之一的层次子集。一个设备上的一组操作的输出传递到下一个设备，后者继续执行后续块。F_n和B_n分别表示设备n上的前向和后向传递。由于处理是顺序的，某些设备或层可能在等待前一层的输出（激活、梯度）时保持空闲，导致前向和后向传递中存在低效率或“流水线气泡”。在图2b中，空白的区域是具有天真流水线并行的大流水线气泡，其中设备保持闲置和未充分利用。

微批处理可以在一定程度上缓解这个问题，如图2c所示。全局输入批次大小被分割成子批次，一个接一个地进行处理，并在最后累积梯度。请注意，F_{n,m}和B_{n,m}分别表示设备n上的前向和后向传递，其中微批次m处理。这种方法缩小了流水线气泡的大小，但不能完全消除它们。

张量并行

张量并行将模型的各个层分割为更小、独立的计算块，可以在不同设备上并行执行。注意力块和多层感知器（MLP）层是可以利用张量并行性的transformer的主要组件。在多头注意力块中，每个头或一组头可以分配给不同的设备，因此它们可以独立并行计算。

图3a显示了在一个具有两层MLP的两路张量并行性的示例，其中每个层由一个圆角框表示。在第一层中，权重矩阵A被分割为A_1和A_2。计算XA_1和XA_2可以在两个不同设备上独立执行。减少了在每个设备上存储权重的内存需求。g将第二层的输出组合起来。

图3b是自注意力层的双向张量并行的示例。多个注意力头天然是并行的，并可以分割到不同的设备上。

序列并行

张量并行有一些局限性，因为它要求将层划分为相互独立、可管理的块。它不适用于LayerNorm和Dropout等操作，而是将它们复制到张量并行组中。虽然LayerNorm和Dropout的计算开销很小，但它们需要大量的内存来存储（冗余的）激活值。

如《减少大型Transformer模型中的激活重新计算》所示，这些操作在输入序列中是相互独立的，可以沿着“序列维度”对它们进行分隔，使它们更节省内存。这被称为序列并行。

模型并行性技术不是排他性的，可以结合使用。它们可以帮助扩展和减少 LLM 的每 GPU 内存占用，但也有专门针对注意力模块的优化技术。

多头注意力

作为对SDPA的增强，通过使用Q，K和V矩阵的不同、学习到的投影，同时并行地执行多次注意力层，使模型能够共同处理来自不同表示子空间的信息。这些子空间是独立学习的，为模型提供了对输入中不同位置的更丰富的理解。

如图5所示，多个并行注意力操作的输出被串联并线性投影以进行合并。每个并行注意力层称为一个“头”，这种方法被称为多头注意力（MHA）。

在原始论文中，当使用八个并行注意力头时，每个注意力头在模型的降维维度上（例如d_{model}/8）。这样可以使计算成本与单头注意力相似。

多查询注意力

MHA的一个推理优化方法称为多查询注意力（MQA），正如Fast Transformer Decoding中所建议的，它共享多个注意力头之间的键和值。查询向量仍然通过多次投影进行操作，与之前一样。

尽管MQA中的计算量与MHA相同，但从内存中读取的数据量（键、值）是原来的一部分。当受内存带宽限制时，这可以提高计算利用率。它还减小了内存中KV缓存的大小，为更大的批量大小腾出了空间。

降低了键值头数量会导致潜在的精度降低。此外，需要在推理中使用此优化的模型需要使用MQA进行训练（或至少使用大约5%的训练量进行微调）。

分组查询注意力

分组查询注意力（GQA）在MHA和MQA之间取得平衡，通过将键和值投影到一些查询头的几个组中实现（图6）。在每个组内，它的行为类似于多查询注意力。

图6显示，多头注意力具有多个键值头（左）。分组查询注意力（中）具有多个键值头，但比查询头的数量少，这是内存需求和模型质量之间的平衡。多查询注意力（右）具有一个键值头，有助于节省内存。

通过使用GQA对使用MHA的模型进行“上训练”，只需使用原始训练计算能力的一小部分即可。它们可以获得接近MHA的质量，并保持更接近MQA的计算效率。Llama270B是一个利用GQA的模型的示例。

像MQA和GQA这样的优化技术有助于通过减少存储的键和值头的数量来减小KV缓存所需的内存。但在管理KV缓存方面可能仍存在一些效率低下的问题。接下来的部分将介绍一种更高效的KV缓存管理技术。

闪存注意力

优化注意力机制的另一种方法是修改某些计算的顺序，以更好地利用GPU内存层次结构。神经网络通常以层为单位进行描述，并且大多数实现也是按照这种方式进行布局，即依次对输入数据执行一种计算。但这并不总是导致最佳性能，因为在较高的、性能更好的内存层次中进行更多的计算，并对已经带入这些层次的更高效的内存中的值进行分组，可能是有益的。

在实际计算中，将多个层级合并在一起，可以最大限度地减少GPU从内存中读取和写入的次数，并将需要相同数据的计算分组在一起，即使它们是神经网络中不同层级的一部分。

其中一个非常流行的合并方法是闪存注意力（FlashAttention），它是一种基于I/O感知的精确注意力算法。精确注意力意味着它在数学上与标准的多头注意力完全相同（可以为多查询和分组查询注意力提供变体），因此可以将其替换到现有的模型架构或已经训练好的模型中，而无需进行修改。

I/O感知是指在将操作合并在一起时，考虑了先前讨论的一些内存移动成本。特别是，FlashAttention使用“分块”来一次完全计算和写入最终矩阵的一小部分，而不是在多个步骤中对整个矩阵进行部分计算，并在其中写出中间值。

图7显示了切片的FlashAttention计算模式和40 GB GPU上的内存层次结构。右侧的图表显示了从合并和重新排序Attention机制的不同组件中获得的相对加速比。

有效管理具有分页的KV缓存

有时，KV缓存被静态地“过度配置”，以适应可能的最大输入大小（支持的序列长度），因为输入大小是不可预测的。例如，如果模型支持的最大序列长度为2048，则无论请求的输入大小和生成的输出大小如何，都会在内存中保留一个大小为2048的保留空间。此空间可能是连续分配的，并且通常很大部分未使用，导致了内存浪费或碎片化。此保留空间与请求的存活时间绑定。

如图8所示，显示了由于过度配置和KV缓存的低效管理导致的内存浪费和碎片化的示例。1）“保留”表示为将来使用保留的内存，该内存在请求持续时间内保留。2）“内部碎片化”发生，因为很难预测生成的长度，因此会过度配置内存以适应最大序列长度。3）“外部碎片化”表示由于批次中的请求需要不同的预分配大小而导致的低效。

受操作系统中的分页的启发，PagedAttention算法使得可以将连续的键和值存储在内存中的非连续空间中。它将每个请求的KV缓存分为表示固定数量的令牌的块，这些令牌可以非连续存储。

在注意力计算过程中，使用一个块表按需获取这些块。当生成新的令牌时，将进行新的块分配。这些块的大小是固定的，消除了由于不同请求需要不同分配而引起的低效性。这大大减少了内存浪费，增加了批量大小（从而提高吞吐量）。

模型优化技术

到目前为止，我们已经讨论了不同方式的LLM消耗内存，一些建议的方法可以跨多个GPU分布内存，并对注意机制和KV缓存进行优化。同样，还有一些模型优化技术可以通过对模型权重本身进行修改来减少每个GPU上的存储器使用。此外，GPU还具有用于加速对这些修改值的操作的专用硬件，从而为模型提供更大的加速。

量化

量化是减少模型权重和激活的精度的过程。大多数模型使用32位或16位的精度进行训练，其中每个参数和激活元素占据32位或16位的存储器空间，即单精度浮点数。然而，大多数深度学习模型可以使用八位甚至更少的位数有效表示。

图9显示了一种量化方法的分布前后值。在这种情况下，通过四舍五入丢失了一些精度，并通过剪切丢失了一些动态范围，使得值可以以更小的格式表示。

减少模型的精度可以带来多个好处。如果模型在内存中占用的空间较小，您可以在相同的硬件上放置更大的模型。量化还意味着您可以在相同的带宽上传输更多的参数，这有助于加速受带宽限制的模型。

对于LLM，有许多不同的量化技术，涉及降低激活、权重或两者精度。量化权重要容易得多，因为它们在训练后是固定的。然而，这可能会浪费一些性能，因为激活仍然在更高的精度上保持。GPU没有用于乘法INT8和FP16数字的专用硬件，因此权重必须再转换为更高精度以进行实际操作。

也可以量化激活，即transformer块和网络层的输入，但这带来了自己的挑战。激活向量通常包含异常值，有效地增加了它们的动态范围，使得以比权重更低的精度表示这些值更具挑战性。

一种选择是通过通过模型传递代表性数据集找出那些异常值可能出现在哪里，并选择以较高精度表示某些激活，在其余激活中复用权重的动态范围。

稀疏性

与量化类似，研究表明许多深度学习模型对修剪是鲁棒的，也就是用0替换接近0的某些值。稀疏矩阵是其中许多元素为0的矩阵。这些矩阵可以用比完整的稠密矩阵更小的空间进行压缩形式来表示。

特别是GPU具有针对某种结构化稀疏性的硬件加速，其中四个值中有两个是用零表示的。稀疏表示可以与量化相结合，以实现更大的执行加速。找到以稀疏格式表示大型语言模型的最佳方法仍然是一个活跃的研究领域，并为将来改善推理速度提供了有希望的方向。

蒸馏

缩小模型大小的另一种方法是通过一种称为蒸馏的过程将其知识转移给更小的模型。该过程涉及训练一个较小的模型（称为学生模型）来模仿一个较大的模型（称为教师模型）的行为。

成功的蒸馏模型示例包括DistilBERT，它将一个BERT模型压缩了40%，同时保留了97%的语言理解能力，并且速度更快了60%。

虽然在LLM中的蒸馏是一个活跃的研究领域，但这种一般方法最初是在《Distilling the Knowledge in a Neural Network》中描述的用于神经网络的。

• 学生网络被训练成反映较大教师网络的性能，使用一个损失函数来度量它们输出之间的差异。这个目标是在可能包含将学生的输出与真实标签进行匹配的原始损失函数之外的。
• 匹配的教师输出可以是最后一层（称为logits）或中间层激活。

图11显示了知识蒸馏的一般框架。教师的logits是学生使用蒸馏损失为之进行优化的软目标。其他蒸馏方法可能使用其他的损失度量来从教师那里“蒸馏”知识。

蒸馏的另一种替代方法是使用由教师合成的数据来对学生LLM进行监督训练，这在人类注释稀缺或不可用的情况下特别有用。《Distilling Step by Step!》进一步从教师LLM中提取了理由，除了用作地面真实标签的标签外。这些理由作为中间推理步骤用于以数据高效的方式训练较小的学生LLM。

值得注意的是，今天的许多最先进的LLM都有限制性许可，禁止使用它们的输出来训练其他LLM，这使得寻找合适的教师模型变得具有挑战性。

模型服务技术

模型执行经常受到内存带宽限制的影响，特别是受到权重的带宽限制。即使应用了之前描述的所有模型优化措施，它仍然很可能受到内存限制。因此，在加载模型权重时，你希望尽可能多地利用它们。换句话说，尝试进行并行处理。可以采取两种方法：

• 动态批处理，在飞行中批处理涉及同时执行多个不同请求。
• 猜测推理（Speculative inference）涉及并行执行序列的多个不同步骤，以尝试节省时间。

在飞行中批处理

LLM具有一些独特的执行特性，可能使得在实践中难以有效地批处理请求。单个模型可以同时用于多种看起来非常不同的任务。从聊天机器人的简单问答响应到文档的摘要或长代码块的生成，工作负载是高度动态的，输出的大小相差几个数量级。

这种多样性可能使得分批处理请求并有效地并行执行成为一种具有挑战性的优化方法，但这可能导致一些请求比其他请求提前完成。

为了管理这些动态负载，许多LLM服务器解决方案包括一种优化的调度技术，称为连续或在飞行批处理。它利用了将LLM的整体文本生成过程分解为对模型的多次执行迭代的事实。

使用在飞行中批处理时，服务器运行时在整个批处理完成之前就开始将完成的序列从批次中删除。然后它开始执行新的请求，而其他请求仍在进行中。在实际使用案例中，在飞行中批处理可以大大增加GPU的整体利用率。

猜测推理

也称为猜测抽样、辅助生成或分块并行解码，猜测推理是一种不同的并行执行LLM的方法。通常，GPT-style大型语言模型是自回归模型，逐个令牌地生成文本。

生成的每个令牌都依赖于之前所有的令牌来提供上下文。这意味着在常规执行中，不可能并行地从同一个序列中生成多个令牌-你必须等待第n个令牌生成后才能生成n+1个。

图12显示了猜测推理的一个例子，其中草案模型暂时预测了多个未来步骤，并且可以并行验证或拒绝。在这种情况下，草案中的前两个预测令牌被接受，而最后一个在继续生成之前被拒绝并删除。

猜测抽样提供了一种解决方法。这种方法的基本思想是使用一些“更便宜”的过程生成几个令牌长的草稿连续输出。然后，在所需的执行步骤中使用廉价的草稿作为“猜测”上下文，同时并行地执行主要的“验证”模型。

如果验证模型生成与草稿相同的令牌，则可以确定接受这些令牌作为输出。否则，可以丢弃第一个不匹配令牌之后的所有内容，并使用新的草稿重复该过程。

有许多不同的选项可以用来生成草稿令牌，并且每种方法都有不同的权衡。可以训练多个模型或在单个预训练模型上微调多个头部，以预测多个步骤的令牌。或者可以使用一个小模型作为草稿模型，并使用一个更大、更能胜任的模型作为验证器。

结论

本文概述了许多最流行的解决方案，以帮助有效地优化和提供LLMs，无论是在数据中心还是在个人电脑上。通过使用包含TensorRT深度学习编译器、优化核心、预处理和后处理步骤以及多GPU/多节点通信原语的NVIDIA TensorRT-LLM，这些技术得到了优化和提供。

NVIDIA TensorRT-LLM现在受到NVIDIA Triton推理服务器的支持，使企业能够在不同的AI框架、硬件加速器和部署模型之间同时提供多个AI模型，以达到最高的吞吐量和最低的延迟。

TensorRT-LLM还为NVIDIA NeMo提供动力，后者为开发人员提供了一个端到端的云原生企业框架，用于构建、定制和部署拥有数十亿参数的生成型AI模型。