这篇文章最初发表在 NVIDIA 技术博客上。
通过堆叠 Transformer 层来创建大型模型，可以提高准确性、实现少量学习功能，并且在各种语言任务中实现近乎人类的性能。这些基础模型的训练成本高昂，而且在推理过程中可能会占用大量内存和计算资源（这是一种重复性成本）。目前最热门的大型语言模型 (LLM)可以达到数百亿到数千亿的参数规模，并且根据用例，可能需要处理长输入（或上下文），这也会增加费用。
本文讨论了大型语言模型（LLM）推理中最紧迫的挑战及其实用解决方案。建议读者对 Transformer 架构 和通用注意力机制有基本的了解。我们将在下一节重点讨论掌握 LLM 推理复杂性的重要性。
了解 LLM 推理
大多数热门的仅使用解码器的 LLM (例如 GPT-3)都基于因果关系建模目标进行了预训练，本质上是作为次词预测器。这些 LLM 接受一系列标记作为输入，并以自回归方式生成后续标记，直到它们满足停止标准(例如，要生成的标记数量的限制或停止词列表)，或直到它生成特殊的<end>标记标志着生成过程的结束。此过程涉及两个阶段：预填充阶段和解码阶段。
请注意，令牌是模型处理的语言的原子部分。一个令牌大约是四个英语字符。自然语言中的所有输入在输入模型之前都会转换为令牌。
预填充阶段或处理输入
在预填充阶段，LLM 处理输入令牌以计算中间状态(键和值)，这些状态用于生成"第一个"新令牌。每个新令牌都依赖于之前的所有令牌，但由于输入的完整范围已知，因此在高层次上，这是一个高度并行的矩阵矩阵运算。它实际上饱和了 GPU 利用率。
解码相位或生成输出
在解码阶段，LLM 一次自动回归生成一个输出令牌，直到满足停止条件为止。每个连续输出令牌都需要知道先前迭代的所有输出状态(键和值)。这就像矩阵向量运算，与预填充阶段相比，它未充分利用 GPU 计算能力。数据(权重、键、值、激活)从内存传输到 GPU 的速度主宰延迟，而不是计算实际发生的速度。换言之，这是一个受内存限制的运算。
本文介绍的许多推理挑战和相应解决方案都与此解码阶段的优化相关：高效注意力模块、有效管理键和值等。
不同的 LLM 可能会使用不同的 tokenizer，因此，比较它们之间的输出令牌可能并不简单。在比较推理吞吐量时，即使两个 LLM 的每秒输出令牌相似，但如果它们使用不同的 tokenizer，它们可能不相同。这是因为相应的令牌可能表示不同数量的字符。
批处理
提高 GPU 利用率和有效吞吐量的最简单方法是通过批处理。由于多个请求使用相同的模型，因此权重的内存成本分散。将更多的批量转移到 GPU 以一次性处理，将利用更多的可用计算。
但是，批量大小只能增加到一定的限制，此时它们可能会导致内存溢出。为了更好地理解发生这种情况的原因，需要查看键值 (KV) 缓存和 LLM 内存要求。
传统的批处理(也称为静态批处理)次优。这是因为，对于批量中的每个请求，LLM 可能会生成不同数量的完成令牌，随后它们的执行时间也不同。因此，批量中的所有请求必须等待最长的请求完成，而生成长度的巨大差异可能会使这种情况更加严重。有一些方法可以缓解这种情况，例如动态批处理，稍后将讨论这一点。
键值缓存
解码阶段的一个常见优化是 KV 缓存。解码阶段在每个时间步长生成单个令牌，但每个令牌都取决于之前所有令牌的键值张量(包括在预填充时计算的输入令牌 KV 张量，以及在当前时间步长之前计算的任何新 KV 张量)。
为避免在每个时间步长为所有令牌重新计算所有这些张量，我们可以将其缓存在 GPU 显存中。每次迭代计算新元时，只需将其添加到正在运行的缓存中，以便在下一次迭代中使用。在某些实现中，模型的每一层都有一个 KV 缓存。
图 1.键值缓存机制说明
LLM 内存需求
实际上，GPU LLM 内存需求的两个主要贡献者是模型权重和 KV 缓存。
- **模型权重：**模型参数占用内存。例如，具有 70 亿个参数（例如 Llama-2-7b），以 16 位精度（FP16 或 BF16）加载时，显存大小约为 70 亿 * 2 字节（FP16）= 14 GB。
- KV 缓存：内存中的自注意力张量被用作缓存，以避免重复计算。
在批处理时，批量中每个请求的 KV 缓存仍必须单独分配，并且可能会占用大量内存。以下公式描述了适用于当今大多数常见 LLM 架构的 KV 缓存的大小。
每个令牌的 KV 缓存大小（字节）= 2 * (num_layers) * (num_heads * dim_head) * precision_in_bytes
第一个系数为 2 解释了 K 和 V 矩阵。通常，(num_heads*dim_head) 的值与 Transformer 的 hidden_size (或模型维度 d_model)相同。这些模型属性通常可在模型卡或相关配置文件中找到。
输入序列中的每批输入中的每个令牌都需要此内存大小。假设半精度，KV 缓存的总大小由以下公式给出。
以字节为单位的 KV 缓存总大小 = (batch_size) + (sequence_length) * 2 * (num_layers) * (hidden_size) * sizeof(FP16)
例如，对于 16 位精度的 Lama 2 7B 模型，批量大小为 1，则 KV 缓存的大小为 1*4096*2*32*4096*2 字节，即小于 2 GB.
高效管理此 KV 缓存是一项艰巨的任务。随着批量大小和序列长度的线性增长，显存需求可以快速扩展。因此，它限制了可以提供的吞吐量，并对长上下文输入构成了挑战。这正是本文介绍的几项优化背后的原因。
使用模型并行扩展 LLM
减少模型权重在每个设备上的显存占用的一种方法是将模型分配到多个 GPU 上。通过扩展显存和计算占用，可以运行更大的模型或更大批量的输入。模型并行化是在需要比单个设备上可用的内存更多的模型上进行训练或推理的必要条件，并使训练时间和推理措施(延迟或吞吐量)适合某些用例。有几种方法可以根据模型权重的分割方式并行化模型。
请注意，数据并行性也是一种技术，通常与下面列出的其他技术在同一上下文中提及。在这种情况下，模型的权重在多个设备上复制，输入的(全局)批量大小在每个设备上被分解为微批量。它通过处理更多的批量来减少总体执行时间。但是，它是一种训练时间优化，在推理期间相关性较小。
管道并行
工作流并行涉及将模型(纵向)分片为多个数据块，其中每个数据块由在单独设备上执行的层子集组成。图 2a 展示了四路工作流并行，其中模型按顺序进行分区，并在每台设备上执行所有层的四分之一子集。一台设备上的一组操作的输出将传递给下一个设备，该设备将继续执行后续的数据块。和 分别指示设备 n 上的正向和反向传递。有效地将在每个设备上存储模型权重的内存需求调至四分之一。
此方法的主要局限性在于，由于处理的顺序性质，某些设备或层可能会在等待前几层的输出(激活、梯度)时保持空闲状态，从而导致前向和反向通道效率低下或出现"管道气泡"。在图 2b 中，白色空白区域是具有朴素管道并行性的大型管道气泡，其中设备处于空闲状态且未得到充分利用。
如图 2c 所示，微批处理可以在一定程度上减轻这种情况。输入的全局批量大小被拆分成子批量，逐个进行处理，并在结束时累加梯度。请注意，和分别指示设备上的正向和反向传递使用微批.此方法会缩小管道气泡的大小，但并不能完全消除它们。
图 2.四路管道并行图。来源：GPipe：通过微批量管道并行轻松扩展
张量并行度
张量并行性涉及将模型的各个层(水平)分片为较小的独立计算块，这些计算块可在不同的设备上执行。注意力块和多层感知器 (MLP) 层是可以利用张量并行性的 Transformer 的主要组件。在多头注意力块中，可以将每个头部或头部组分配给不同的设备，以便可以独立和并行计算。
图 3.多层感知器 (MLP) 和自注意力层中张量并行性的说明。来源：Megatron-LM：使用模型并行训练数十亿个参数语言模型
图 3a 展示了两层 MLP 上的双向张量并行性示例，其中每层均由一个圆框表示。在第一层中，权重矩阵分为 和.计算和可以在同一批次上独立执行(是输入的身份运算)在两台不同的设备上运行。这实际上将在每台设备上存储权重所需的内存需求减半。组合第二层中的输出。
图 3b 是自注意力层中双向张量并行性的示例。多个注意力头本质上是并行的，可以在设备之间分割。
序列并行度
Tensor 并行性存在局限性，因为它需要将层划分为独立的、可管理的块。它不适用于 LayerNorm 和 Dropout 等操作，而是在张量并行组中进行复制。尽管 LayerNorm 和 Dropout 的计算成本较低，但它们确实需要大量内存来存储（冗余的）激活。
如《减少大型 Transformer 模型中的激活重新计算》所示，这些操作独立于输入序列，并且可以沿着"序列维度"进行划分，从而提高内存效率。这称为序列并行性。
图 4.具有张量和序列并行性的 Transformer 层说明。来源：减少大型 Transformer 模型中的激活重新计算
模型并行技术并非独有，可以结合使用。它们可以帮助扩展和减少 LLM 的每个 GPU 显存占用，但也有专门针对注意力模块的优化技术。
优化注意力机制
扩展的点积注意力 (SDPA) 运算将查询和键值对映射到输出，详情请参阅《您只需集中注意力》。
多头注意力
作为对 SDPA 的增强，通过多次并行执行 Q、K 和 V 矩阵的不同学习投影来执行注意力层，使模型能够共同处理来自不同位置的不同表征子空间的信息。这些子空间是独立学习的，使模型能够更深入地了解输入中的不同位置。
如图 5 所示，我们将多个并行注意力运算的输出连接起来，并对其进行线性投影，以将其组合起来。每个并行注意力层称为"头部"，这种方法称为多头注意力 (MHA).
在原始作品中，每个注意力头都在模型的简化维度上操作(例如)。这使得计算成本与单头注意力相似。
图 5.缩放点积注意力(左)和多头注意力(右)的示例，即简单地并行显示多个 SDPA 头。来源：您只需集中注意力
多查询注意力
对 MHA 的一种推理优化，称为多查询注意力 (MQA)，正如在 快速 Transformer 解码 中所述，它在多个注意力头之间共享键和值。查询向量仍会像之前一样进行多次投影。
虽然 MQA 中完成的计算量与 MHA 相同，但从内存读取的数据量(键、值)只是之前的一小部分。受内存带宽限制时，这可以提高计算利用率。它还减少了内存中的 KV 缓存的大小，为更大的批量大小提供了空间。
减少键值头可能会导致准确性降低。此外，要利用这种优化的模型进行推理，需要进行训练（或至少经过微调，训练量约为原始训练量的 5%），同时需要启用 MQA。
分组查询注意力
分组查询注意力（GQA）通过将关键帧和值投射到几组查询头（图 6），在 MHA（多头注意力）和 MQA（多查询注意力）之间取得平衡。在每个组中，其行为类似于多查询注意力。
图 6 显示了多头注意力具有多个键值头(左)。分组查询注意力(中心)的键值头超过一个，但少于查询头的数量，这是内存需求和模型质量之间的平衡。多查询注意力(右)具有单个键值头以帮助节省内存。
图 6.不同注意力机制的比较。来源：GQA：从多头检查点训练通用多查询 Transformer 模型
最初使用 MHA 训练的模型，可以使用原始训练计算的一小部分通过 GQA 进行"上训练"。它们的质量接近 MHA，同时保持接近 MQA 的计算效率。Lama 2 70B 是一个利用 GQA 的模型示例。
MQA 和 GQA 等优化通过减少存储的键和值头的数量来帮助减少 KV 缓存所需的内存。管理此 KV 缓存的方式可能仍然效率低下。与优化注意力模块本身不同，下一节将介绍一种更高效的 KV 缓存管理技术。
闪光注意力
优化注意力机制的另一种方法是修改某些计算的顺序，以更好地利用 GPU 的内存层次结构。神经网络通常以层的形式进行描述，大多数实现也以这种方式进行布局，按顺序对输入数据进行一种计算。这并不总是能实现最佳性能，因为对已引入更高、性能更高的内存层次结构的值进行更多计算会很有好处。
在实际计算期间将多个层融合在一起，可以最大限度地减少 GPU 从内存读取和写入内存的次数，并将需要相同数据的计算分组，即使它们是神经网络中不同层的一部分。
一种非常流行的技术融合是 FlashAttention，这是一种 I/O 感知的精确注意力算法，详见FlashAttention：通过 I/O 感知实现快速且节省内存的精确注意力。所谓精确注意力，意味着它在数学上与标准的多头注意力（提供用于多查询和分组查询注意力的变体）相同，因此可以无缝集成到现有的模型架构中，甚至适用于未经修改的已训练模型。
I/O 感知表示在将运算融合在一起时，它会考虑到之前讨论的一些内存移动成本。特别是，FlashAttention 使用"平铺"来完全计算并一次性写出一小部分最终矩阵，而不是按步骤对整个矩阵进行部分计算，写出中间值。
图 7 显示了 40 GB GPU 上平铺的 FlashAttention 计算模式和内存层次结构。右侧图表显示了因融合和重新排序 Attention 机制的不同组件而导致的相对加速。
图 7.40 GB GPU 上平铺的 FlashAttention 计算模式和内存层次结构。来源：FlashAttention：通过 IO 感知实现快速且节省内存的精确注意力
通过分页有效管理 KV 缓存
由于输入大小不可预测，因此 KV 缓存有时会静态地"过度置备"，以考虑到尽可能大的输入(支持的序列长度)。例如，如果模型的支持最大序列长度为 2048，那么无论请求中的输入大小和生成的输出大小如何，都将在内存中保留大小为 2048 的空间。此空间可能会被连续分配，而且其中大部分仍然未使用，从而导致内存浪费或碎片化。在请求的整个生命周期中，此保留空间被占用。
图 8.过度调配和低效 KV 缓存管理导致的内存浪费和碎片说明。来源：借助 PagedAttention 为大型语言模型服务实现高效内存管理
受操作系统分页机制启发，PagedAttention 算法允许在内存中非连续空间存储连续的键和值。它将每个请求的 KV 缓存分割成代表固定数量令牌的块，并且这些令牌可以以非连续的方式进行存储。
在注意力计算期间，系统会根据需要使用记录帐户的块表来获取这些块。生成新令牌后，系统会进行新的块分配。这些块的大小是固定的，可以消除因挑战(例如需要不同分配的不同请求)而产生的低效问题。这显著限制了内存浪费，从而实现了更大的批量大小(从而提高吞吐量)。
模型优化技术
到目前为止，我们已经讨论了 LLM 消耗内存的不同方式、在多个不同 GPU 之间分配内存的一些方式，以及优化注意力机制和 KV 缓存。还有一些模型优化技术，可以通过修改模型权重本身来减少每个 GPU 上的内存使用。GPU 还有专用硬件来加速对这些修改值的操作，从而提高模型的速度。
量化
量化是降低模型权重和激活精度的过程。大多数模型都以 32 或 16 位精度进行训练，其中每个参数和激活元素占用 32 或 16 位内存(单精度浮点)。但是，大多数深度学习模型可以用每个值 8 位甚至更少的位来有效表示。
图 9 显示了一种可能的量化方法前后的值分布。在这种情况下，舍入会损失一些精度，裁剪会损失一些动态范围，从而以更小的格式表示值。
图 9.一种可能的量化方法前后的值分布
降低模型的精度可以产生一些好处。如果模型占用的内存空间较少，您可以在相同数量的硬件上拟合更大的模型。量化还意味着您可以在相同数量的带宽上传输更多参数，这有助于加速带宽受限的模型。
LLM 有许多不同的量化技术，涉及降低激活值、权重或两者的精度。量化权重更直接，因为它们在训练后是固定的。但是，由于激活值保持在更高的精度，因此这可能会在表格上保留一些性能。GPU 没有用于乘以 INT8 和 FP16 数字的专用硬件，因此必须将权重转换回更高的精度，以用于实际操作。
也可以量化 Transformer 块和网络层的激活、输入，但这也有其自身的挑战。激活向量通常包含离群值，从而有效地增加其动态范围，使得以低于权重的精度表示这些值更具挑战性。
一种方法是，通过在模型中传递代表性数据集，并选择以比其他模型更高的精度表示某些激活 (LLM.int8 ())，找出这些异常值可能出现的位置。另一种方法是借用易于量化的权重动态范围，并在激活中重复使用该范围。
稀疏
与量化类似，研究表明，许多深度学习模型在修剪或将接近 0 的某些值替换为 0 本身方面都很可靠。稀疏矩阵是许多元素为 0 的矩阵。它们可以用凝聚形式表示，与完整的密集矩阵相比，它占用的空间更小。
图 10.以压缩格式表示的稀疏矩阵，其中包含非零数据值及其相应的 2 位索引
尤其是 GPU，对于某些类型的结构化稀疏每 4 个值中就有 2 个由 0 表示。稀疏表示还可以与量化相结合，以实现更高的执行速度。寻找以稀疏格式表示大型语言模型的最佳方法仍然是一个活跃的研究领域，并为未来提高推理速度提供了一个大有希望的方向。
蒸
缩小模型大小的另一种方法是通过名为蒸此过程涉及训练一个较小的模型(称为学生)，以模拟较大模型(教师)的行为。
成功的模型压缩示例包括DistilBERT，它将 BERT 模型的大小压缩了 40%，同时保留了 97% 的语言理解能力，并且速度提升了 60%。
尽管在大型语言模型（LLM）中的知识蒸馏是一个活跃的研究领域，但更多相关信息可以在提取神经网络中的知识中找到：
- 学生网络经过训练，可以反映更大的教师网络的性能，使用测量其输出之间差异的损失函数。除了可能包括将学生的输出与真值标签相匹配的原始损失函数之外，还可以实现这一目标。
- 匹配的教师输出可以是最后一层(称为逻辑)或中间层激活。
图 11 展示了知识蒸的通用框架。教师的逻辑是学生使用蒸损失进行优化的软目标。其他蒸方法可能会使用其他损失度量来"提炼"教师的知识。
图 11.知识提炼的通用框架。来源：知识蒸：调查
另一种训练方法是使用由教师模型合成的数据对学生大语言模型 (LLM) 进行监督式训练，这在人类注释稀缺或不可用的情况下特别有用。逐步提炼！除了作为真实标签之外，还可以从教师 LLM 中提取推理过程。这些推理过程作为中间步骤，可以以数据高效的方式训练较小的学生 LLM。
需要注意的是，目前许多先进的 LLM 都有严格的许可证，禁止使用其输出来训练其他 LLM，因此很难找到合适的教师模型。
模型服务技术
模型执行通常受内存带宽限制，尤其是权重中的带宽限制。即使应用了前面介绍的所有模型优化，它仍然很可能受内存限制。因此，您希望在加载模型权重时尽可能多地处理这些权重。换言之，尝试并行执行操作。可以采取两种方法：
- **动态批处理：**涉及同时执行多个不同的请求。
- **推测推理：**涉及并行执行序列中的多个异步步骤以节省时间。
动态批处理
LLM 具有一些独特的执行特性，在实践中可能难以有效地批处理请求。单个模型可同时用于各种看起来非常不同的任务。从聊天机器人中的简单问答响应到文档摘要或生成长代码块，工作负载具有高度动态化的特点，输出大小变化了几个数量级。
这种通用性会使批处理请求并有效地并行执行这些请求(这是服务神经网络的常见优化)具有挑战性。这可能会导致一些请求比其他请求更早完成。
为了管理这些动态负载，许多 LLM 服务解决方案包括一种名为连续或动态批处理的优化调度技术。这利用了一个事实，即 LLM 的整个文本生成过程可以分解为模型上的多次执行迭代。
在运行中批处理时，服务器运行时无需等待整个批处理完成再转到下一组请求，而是会立即从批量中移除已完成的序列。然后，当其他请求仍处于运行中时，服务器会开始执行新请求。因此，运行中批处理可以在实际用例中大幅提高 GPU 的整体利用率。
推测推理
推测推理也称为推测采样、辅助生成或块级并行解码，是并行执行 LLM 的另一种方式。通常，GPT 风格的大型语言模型是按令牌生成文本标记的自回归模型。
生成的每个令牌都依赖于之前提供上下文的所有令牌。这意味着，在常规执行中，不可能并行从同一序列生成多个令牌 -- 您必须等待生成第 n 个令牌，然后才能生成 n+1.
图 12 展示了一个推理示例，其中草稿模型暂时预测多个并行验证或拒绝的未来步骤。在这种情况下，草稿中的前两个预测标记被接受，而最后一个标记被拒绝并删除，然后再继续生成。
图 12.推理示例。来源：深度自回归模型的块级并行解码
预测采样提供了一种变通方法。此方法的基本思路是使用一些"更便宜"的过程来生成包含多个令牌的草稿延续。然后，并行执行多个步骤的主"验证"模型，在需要执行步骤时将廉价草稿用作"预测"上下文。
如果验证模型生成的令牌与草稿相同，则您知道接受这些令牌用于输出。否则，您可以在第一个不匹配令牌之后丢弃所有内容，然后使用新草稿重复此过程。
如何生成草稿令牌有许多不同的选项，每种选项都有不同的权衡。您可以训练多个模型，或在单个预训练模型上微调多个头部，以预测未来的多个步骤的令牌。或者，您可以使用小型模型作为草稿模型，使用更大、功能更强大的模型作为验证器。
结束语
本文将概述许多最热门的解决方案，以帮助高效优化和服务 LLM，无论是在数据中心还是在 PC 的边缘。其中许多技术都经过优化，并通过 NVIDIA TensorRT-LLM 这是一个开源库，由 TensorRT 深度学习编译器、优化的内核、预处理和后处理步骤以及多 GPU/多节点通信基元组成，可在 NVIDIA GPU 上实现突破性性能。如需了解详情，请参阅 借助 NVIDIA TensorRT-LLM 优化大型语言模型的推理（现已公开发布）。
NVIDIA TensorRT-LLM 现在由 NVIDIA Triton 推理服务器 支持，使企业能够跨不同的 AI 框架、硬件加速器和部署模式同时为多个 AI 模型提供服务，并实现峰值吞吐量和更低延迟。
TensorRT-LLM 还支持 NVIDIA NeMo，它提供了一个端到端的云原生企业框架，供开发者构建、自定义和部署具有数十亿参数的生成式 AI 模型。欲了解更多信息，请开始使用 NeMo。