原文链接： developer.nvidia.com/blog/master...
作者：By Shashank Verma and Neal Vaidya
通过堆叠transformer层来创建大模型可以提高模型的准确性，使模型具有少样本学习能力，并且在各种语言任务中，这类模型可以表现出与人类类似的应对能力。这些基础模型的训练成本很高，而且在推理过程中会占用大量的内存与计算。目前最流行的大型语言模型(llm)可以达到数百亿到数千亿个参数，并且根据使用场景，可能需要读取长文本输入，这也会增加费用。例如RAG的管道需要将大量的信息放入模型的输入中，这极大地增加了LLM的工作量。
理解大模型推理
大多数流行的只有解码器的llm(例如GPT-3)都是在因果建模目标上进行预训练的，本质上是作为下一个单词的预测器。这些llm将一系列token作为输入，并自动回归地生成后续的token，直到它们满足停止条件(例如，提供一个停用词列表或者限制生成token的数量）。这个过程包括两个阶段:预填充阶段和解码阶段。
预填充或处理输入
在预填充阶段，LLM处理输入的token用于计算第一个输出token的中间状态（键值对），每个新的token取决于所有先前的token，但由于已知输入的全部范围，实际上，这是一个高度并行化的矩阵-矩阵操作。它有效地提高了 GPU 的利用率。
解码阶段或生成输出
在解码阶段，LLM（大型语言模型）自回归地逐个生成输出token，直到满足停止条件为止。每个顺序输出token都需要知道所有先前迭代的输出状态（键和值）。这类似于一个矩阵-向量操作，与预填充阶段相比，它没有充分利用 GPU 的计算能力。数据（权重、键、值、激活）从内存传输到 GPU 的速度造成了延迟。
这篇文章中涉及的许多推理挑战和相应的解决方案都涉及优化这个解码阶段：高效的注意力模块，有效管理键和值等。
不同的LLM可能使用不同的分词器，因此，在它们之间比较输出token可能并不直接。在比较推理吞吐量时，即使两个LLM每秒输出相似的token数，如果它们使用不同的分词器，它们可能并不等效。这是因为相同的token可能代表不同数量的字符。
批处理
通过批处理是提高 GPU 利用率和吞吐量的最简单方法。由于多个请求使用相同的模型，权重的内存成本被分摊。将较大的批次一次性传输到 GPU 进行处理将更充分地利用可用的计算资源。
然而，批次大小只能增加到一定限制，此时可能会导致内存溢出。要更好地理解为什么会发生这种情况，需要查看键-值（KV）缓存和LLM内存需求。
传统的批处理（也称为静态批处理）并不是最好的办法。这是因为在批处理中的每个请求中，LLM 可能会生成不同数量的完成token，它们具有不同的执行时间。因此，批处理中的所有请求必须等到最长的请求完成。有一些方法可以减轻这种情况，例如飞行批处理（In-flight batching），这将在后面讨论。
kv缓存
解码阶段的一个常见优化是KV缓存。解码阶段在每个时间步生成单个token，但每个token依赖于所有先前token的键和值张量(包括在预填充时计算的输入token的Kv张量，以及直到当前时间步计算的所有新token的Kv张量)。
为了避免每次都重新计算张量，会将他们缓存在显存中，每次迭代时，当新的元素被计算好了，就会把他们添加到正在运行的缓存中，用于下次迭代。有的时候，模型的每一层都有一个kv 缓存。
大模型的内存要求
实际上，显存的主要组成是kvcache和模型权重: 一个7B的大语言模型在16位的精度下需要14GB显存。
在批处理中，批处理中每个请求的 KV 缓存仍必须单独分配，因此会占用大量内存。下式描述了 KV 缓存的大小，适用于当今大多数常见的 LLM 架构。
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> 每个 t o k e n 的 k v c a c h e 占用 = 2 ∗ 模型层数 ∗ ( 变换器的隐藏尺寸 ) ∗ 精度所占用位数 每个token的kvcache占用 = 2 * 模型层数 * (变换器的隐藏尺寸) * 精度所占用位数 </math>每个token的kvcache占用=2∗模型层数∗(变换器的隐藏尺寸)∗精度所占用位数
有效管理 KV 缓存是一项具有挑战性的工作。随着批量大小和序列长度的线性增长，内存需求也会迅速增加。因此，它限制了可提供的吞吐量，并给长上下文输入带来了挑战。这就是本文章中介绍的几种优化方法背后的动机。
利用模型并行化拓展大模型
减少模型权重对每台设备内存占用的一种方法是将模型分布到多个 GPU 上。分散内存和计算足迹可以运行更大的模型或更大的输入批次。模型并行化是对模型进行训练或推理的必要条件，因为模型需要的内存比单个设备上可用的内存要多，而且要使训练时间和推理结果（延迟或吞吐量）适合某些用例。根据模型权重的分割方式，有几种并行化模型的方法。
需要注意的是，数据并行也是一种经常被提及的技术，与下面列出的其他技术具有相同的背景。在这种情况下，模型的权重被复制到多个设备上，输入的（全局）批次大小被分割到每个设备的微批次中。它通过处理更大的批次来减少整体执行时间。不过，这是一种训练时间优化，在推理过程中意义不大。
流水线并行
流水线并行涉及将模型（垂直）分块，每个分块包括在单独设备上执行的层的子集。图 2a 是四路流水线并行的示意图，其中模型被顺序分割，所有层的四分之一子集在每个设备上执行。一个设备上的一组操作的输出被传递到下一个设备，由下一个设备继续执行下一个分块。
这种方法的主要局限是，由于处理的顺序性，一些设备或层在等待前一层的输出（激活、梯度）时可能会处于空闲状态。这就导致了前向和后向通道的低效率或 "流水线气泡"。在图 2b 中，白色的空白区域是大的流水线气泡，这些气泡具有较差的流水线并行性，其中的设备处于闲置状态，未得到充分利用。
如图 2c 所示，微批处理可以在一定程度上缓解这一问题。输入的全局批量大小被分割成子批量，子批量逐个处理，梯度在最后累积。请注意，和 分别表示微批次设备上的前向和后向通道。这种方法缩小了管道气泡的大小，但并没有完全消除气泡。
张量并行
张量并行涉及将模型的各个层分片（横向）为较小的独立计算块，这些计算块可在不同设备上执行。注意力区块和多层感知器（MLP）层是可以利用张量并行的变压器的主要组成部分。在多头注意力块中，每个头或每组头都可以分配到不同的设备上，这样它们就可以独立并行计算。
图 3a 显示了双向张量并行在双层 MLP 上的示例，每一层用一个圆形方框表示。在第一层中，权重矩阵分为 <math xmlns="http://www.w3.org/1998/Math/MathML"> A 1 A_1 </math>A1和 <math xmlns="http://www.w3.org/1998/Math/MathML"> A 2 A_2 </math>A2 。 <math xmlns="http://www.w3.org/1998/Math/MathML"> X A 1 XA_1 </math>XA1和 <math xmlns="http://www.w3.org/1998/Math/MathML"> x A 2 xA_2 </math>xA2的计算可以在两台不同设备上的同一批输入上独立执行。这就有效地将每个设备上存储权重所需的内存减半。还原运算g将输出合并送到第二层。
图 3b 是自我注意层双向张量并行的一个例子。多个注意力头在本质上是并行的，可以在不同设备之间分割。
序列并行
张量并行有其局限性，因为它要求将层划分为独立的、可管理的块。它不适用于 LayerNorm 和 Dropout 等操作，这些操作需要在张量并行组中复制。虽然 LayerNorm 和 Dropout 的计算成本很低，但它们确实需要大量冗余内存来存储激活。
如减少大型变换器模型中的激活重计算一文所示，这些操作在整个输入序列中是独立的，这些操作可以沿着 "序列维度 "进行分割，从而提高内存效率。这就是所谓的序列并行性。
模型并行化技术并不是互斥性的，可以结合使用。它们可以帮助扩展和减少 LLM 的每个 GPU 内存占用，但也有专门针对注意力模块的优化技术。
优化注意力机制
模型通过查询在所有的键中寻找匹配的信息，并根据匹配程度（即权重）从对应的值中提取出相关信息。通过这种方式，模型可以实现对输入数据中不同部分的有选择性的关注，从而在处理长序列、复杂结构或者需要忽略无关信息的情况下取得更好的效果。
多头注意力
作为对SDPA的增强，通过多次并行执行 Q、K 和 V 矩阵的不同学习投影来执行注意力层，使模型能够共同处理来自不同位置的不同表征子空间的信息。这些子空间是独立学习的，使模型能够更深入地了解输入中的不同位置。
如图 5 所示，我们将多个并行注意力运算的输出连接起来，并对其进行线性投影，以将其组合起来。每个并行注意力层称为"头部"，这种方法称为多头注意力 (MHA).
在原始工作中中，每个注意力头都在模型的简化维度上操作(例如 <math xmlns="http://www.w3.org/1998/Math/MathML"> d m o d e l / 8 d_{model}/8 </math>dmodel/8)。这使得计算成本与单头注意力相似。
多查询注意力
多查询注意力(MQA)是对多头注意力(MHA)的一种推理优化,正如Fast Transformer Decoding中提出的那样,它在多个注意力头之间共享键和值。查询向量仍然像以前一样投射多次。
虽然MQA和MHA的计算量是相同的,但从内存中读取的数据量(键、值)只有之前的一小部分。当受内存带宽限制时,这可以提高计算利用率。它还减小了内存中KV缓存的大小,从而为更大的批量尺寸腾出空间。
减少键值头的数量会带来一定的准确性下降。此外,需要在推理时利用这一优化的模型,必须在启用MQA的情况下进行训练(或至少使用约5%的训练量进行微调)。
组合查询注意力
GQA 在MHA和MQA之间达到平衡，通过将键和值投射到几个查询头组中实现。在每个组内，他的行为类似于多查询注意力，图6 显示 多头注意力有多个键值头(左), 组合查询注意力(中间)有多于一个的键值头,但少于查询头的数量,这在内存需求和模型质量之间达到了平衡。多查询注意力(右)只有一个键值头,以帮助节省内存。
最初使用 MHA 训练的模型，可以使用 GQA 进行 "升级训练"，只需使用原来训练计算量的一小部分。它们可以达到接近 MHA 的质量，同时保持接近 MQA 的计算效率。Llama 2 70B 就是一个利用 GQA 的模型实例。
MQA 和 GQA 等优化技术通过减少键值头的存储数量，有助于减少 KV 缓存所需的内存。但 KV 缓存的管理仍可能存在效率低下的问题。与注意力模块本身的优化不同，下一节将介绍一种更高效的 KV 缓存管理技术。
Flash attention
优化注意力机制的另一种方法是修改某些计算的顺序,以更好地利用GPU的内存层次结构。神经网络通常是按层描述的,大多数实现也是按这种方式布局的,一次对输入数据执行一种计算。这并不总是能达到最佳性能,因为在已经被带入内存层次结构的更高、更高性能级别的值上进行更多计算可能会更有益。
在实际计算过程中将多个层融合在一起,可以最小化GPU需要读取和写入其内存的次数,并将需要相同数据的计算组合在一起,即使它们是神经网络中不同层的一部分。
一种非常流行的融合方式是FlashAttention,这是一种 I/O 感知的精确注意力算法,如《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》中所详述的。精确注意力意味着它在数学上与标准多头注意力完全相同,因此可以无需修改地替换到现有的模型架构甚至已经训练好的模型中。
I/O 感知意味着它在融合操作时考虑了一些之前讨论过的内存移动成本。特别是,FlashAttention使用"切片"来完全计算和一次性写出最终矩阵的一小部分,而不是分步骤对整个矩阵进行部分计算,并在此过程中写出中间值。
图 7 显示了 40 GB GPU 上的 FlashAttention 计算模式和内存层次结构。右图显示了通过融合和重新排序注意力机制的不同组件所带来的相对速度提升。
