优化 LLM 的速度和内存

诸如 GPT3/4、Falcon 和 Llama 等大型语言模型（LLM）在处理以人为中心的任务方面的能力正在迅速提升，已成为现代知识密集型产业中的关键工具。然而，在实际任务中部署这些模型仍然充满挑战。

• 为了展现近乎人类的文本理解和生成能力，目前的 LLM 需要包含数十亿个参数（参见 Kaplan 等人，Wei 等人）。这相应地增加了推理时的内存需求。
• 在许多现实世界的任务中，需要向 LLM 提供大量的上下文信息。这就要求模型具备在推理过程中处理超长输入序列的能力。

这些挑战的核心在于如何增强 LLM 的计算和内存能力，特别是在处理广泛的输入序列时。

在本指南中，我们将介绍高效部署 LLM 的有效技术。

1. 降低精度：研究表明，以降低的数值精度（即 8 位和 4 位）运行模型，可以在不显著降低模型性能的情况下获得计算优势。

2. Flash Attention：Flash Attention 是一种注意力算法的变体，它不仅提供了一种更节省内存的方法，而且由于优化了 GPU 内存利用率，还实现了更高的效率。

3. 架构创新：考虑到 LLM 在推理时通常以相同的方式部署，即长输入上下文下的自回归文本生成，人们提出了专门的模型架构来允许更高效的推理。在此方面，模型架构最重要的进展是 Alibi、旋转嵌入（Rotary embeddings）、多查询注意力（MQA） 和 分组查询注意力（GQA）。

在本指南中，我们将从张量的角度对自回归生成进行分析。我们将深入探讨采用降低精度的利弊，全面探讨最新的注意力算法，并讨论改进的 LLM 架构。在此过程中，我们将运行实际示例来展示每一项功能改进。

1. 降低精度

LLM 的内存需求可以通过将 LLM 视为一组权重矩阵和向量，并将文本输入视为向量序列来最好地理解。下文中，“权重”一词将用于表示所有的模型权重矩阵和向量。

在撰写本指南时，LLM 至少包含几十亿个参数。每个参数由一个十进制数组成，例如 4.5689，通常以 float32、bfloat16 或 float16 格式存储。这使我们可以轻松计算将 LLM 加载到内存中的内存需求。

在 float32 精度下，加载拥有 X 十亿参数的模型需要大约 4 * X GB 的显存。

如今，模型很少在完整的 float32 精度下进行训练，通常在 bfloat16 精度下训练，或者较少使用 float16 精度。因此，经验法则变为：

在 bfloat16/float16 精度下，加载拥有 X 十亿参数的模型需要大约 2 * X GB 的显存。

对于较短的文本输入（少于 1024 个 token），推理的内存需求很大程度上由加载权重的内存需求决定。因此，目前我们假设推理的内存需求等于将模型加载到 GPU 显存中的内存需求。

为了给出一些例子，说明以 bfloat16 加载模型大约需要多少显存：

• GPT3 需要 2 * 175 GB = 350 GB 显存
• Bloom 需要 2 * 176 GB = 352 GB 显存
• Llama-2-70b 需要 2 * 70 GB = 140 GB 显存
• Falcon-40b 需要 2 * 40 GB = 80 GB 显存
• MPT-30b 需要 2 * 30 GB = 60 GB 显存
• bigcode/starcoder 需要 2 * 15.5 = 31 GB 显存

截至本文撰写时，市场上最大的 GPU 芯片是 A100 和 H100，提供 80GB 的显存。上面列出的大多数模型仅加载就需要超过 80GB 的空间，因此必须使用 张量并行（tensor parallelism） 和/或 流水线并行（pipeline parallelism）。

🤗 Transformers 现在支持对具有 base_tp_plan 配置类的受支持模型进行张量并行。在此 了解有关张量并行的更多信息。此外，如果您有兴趣以张量并行友好的方式编写模型，请查看 text-generation-inference 库。

朴素流水线并行（Naive pipeline parallelism）是开箱即用的。为此，只需使用 device="auto" 加载模型，它就会自动将不同的层放置在可用的 GPU 上，如此处所述。但请注意，虽然非常有效，但这种朴素流水线并行并不能解决 GPU 空闲的问题。为此，需要更高级的流水线并行，如此处所述。

如果您可以使用 8 x 80GB A100 节点，则可以按如下方式加载 BLOOM。

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

通过使用 device_map="auto"，注意力层将平均分配到所有可用的 GPU 上。

在本指南中，我们将使用 bigcode/octocoder，因为它可以在单个 40 GB A100 GPU 芯片上运行。请注意，我们将要应用的所有内存和速度优化，同样适用于需要模型并行或张量并行的模型。

由于模型是以 bfloat16 精度加载的，根据我们上面的经验法则，我们预计运行 bigcode/octocoder 推理的内存需求约为 31 GB 显存。让我们试一试。

我们首先加载模型和分词器（tokenizer），然后将它们传递给 Transformers 的 pipeline 对象。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

输出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

太棒了，我们现在可以直接使用结果将字节转换为千兆字节。

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

让我们调用 torch.cuda.memory.max_memory_allocated 来测量峰值 GPU 内存分配。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

29.0260648727417

与我们粗略的计算非常接近！我们可以看到数字并不完全准确，因为从字节到千字节需要乘以 1024 而不是 1000。因此，这个粗略的公式也可以理解为“最多 X GB”的计算。请注意，如果我们尝试以完整的 float32 精度运行模型，将需要高达 64 GB 的显存。

如今几乎所有模型都是用 bfloat16 训练的，如果 您的 GPU 支持 bfloat16，则没有理由以完整的 float32 精度运行模型。Float32 不会比训练模型时使用的精度提供更好的推理结果。

如果您不确定模型权重在 Hub 上以何种格式存储，您可以随时查看检查点的配置（config）下的 "dtype"，例如此处。建议在加载时使用 from_pretrained(..., dtype=...) 将模型设置为与配置中写入的精度类型相同的精度，除非原始类型是 float32，在这种情况下，推理时可以使用 float16 或 bfloat16。

让我们定义一个 flush(...) 函数来释放所有分配的内存，以便我们能够准确测量峰值分配的 GPU 内存。

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.cuda.empty_cache()
  torch.cuda.reset_peak_memory_stats()

现在让我们为下一个实验调用它。

flush()

在 Accelerate 库中，您还可以使用名为 release_memory 的设备无关实用方法，它考虑了 XPU、MLU、NPU、MPS 等各种硬件后端。

from accelerate.utils import release_memory
# ...

release_memory(model)

那么，如果您的 GPU 没有 32 GB 的显存怎么办？研究发现，模型权重可以量化为 8 位或 4 位，而不会显着降低性能（参见 Dettmers 等人）。模型甚至可以量化为 3 位或 2 位，且性能损失在可接受范围内，如最近的 GPTQ 论文所示 🤯。

在不深入细节的情况下，量化方案旨在降低权重的精度，同时尽可能保持模型推理结果的准确性（即尽可能接近 bfloat16）。请注意，量化对于文本生成特别有效，因为我们只关心选择“最可能出现的下一个 token 集合”，而不真正在乎下一个 token 概率分布的确切值。唯一重要的是下一个 token 的概率分布保持大致不变，以便 argmax 或 topk 操作给出相同的结果。

存在各种量化技术，我们不在这里详细讨论，但通常所有的量化技术工作原理如下：

• 1. 将所有权重量化为目标精度
• 2. 加载量化后的权重，并以 bfloat16 精度传入向量输入序列
• 3. 动态地将权重反量化（dequantize）为 bfloat16，以在 bfloat16 精度下与其输入向量一起执行计算

简而言之，这意味着 输入-权重矩阵 乘法（其中$X$是 输入，$W$是权重矩阵，$Y$是输出$$Y = X * W$$

更改为$$Y = X * \text{dequantize}(W)$$

对于每个矩阵乘法。随着输入通过网络图，反量化和重量化操作会依次对所有权重矩阵执行。

因此，使用量化权重时，推理时间通常不会减少，反而会增加。理论够多了，让我们试一试！要使用 Transformers 量化权重，您需要确保已安装 bitsandbytes 库。

!pip install bitsandbytes

然后，我们只需在 from_pretrained 中添加一个 load_in_8bit=True 标志，即可以 8 位量化加载模型。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

现在，让我们再次运行示例并测量内存使用情况。

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

输出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

不错，我们得到了与之前相同的结果，准确性没有损失！让我们看看这次使用了多少内存。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

15.219234466552734

少了很多！降到了 15 GB 多一点，因此我们可以在 4090 这样的消费级 GPU 上运行该模型。我们看到了非常好的内存效率提升，且模型输出几乎没有退化。然而，我们也可以注意到推理过程中稍微有些变慢。

我们删除模型并再次刷新内存。

del model
del pipe

flush()

让我们看看 4 位量化带来的峰值 GPU 内存消耗。将模型量化为 4 位可以使用与之前相同的 API —— 这次只需传入 load_in_4bit=True 而不是 load_in_8bit=True。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

输出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

我们看到的输出文本与之前几乎相同 —— 只是代码片段前少了 python。让我们看看需要多少内存。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

9.543574333190918

只有 9.5GB！对于一个超过 150 亿参数的模型来说，这确实不多。

虽然我们在这里看到模型的准确性几乎没有退化，但在实践中，4 位量化通常会导致与 8 位量化或完整 bfloat16 推理不同的结果。这需要用户自行尝试。

还要注意，此处的推理再次比 8 位量化慢了一些，这是因为 4 位量化使用了更激进的量化方法，导致$\text{quantize}$和$\text{dequantize}$在推理过程中花费的时间更长。

del model
del pipe

flush()

总的来说，我们看到以 8 位精度运行 OctoCoder 将所需的 GPU 显存从 32G 减少到仅 15GB，而以 4 位精度运行模型进一步将所需的 GPU 显存减少到 9GB 多一点。

4 位量化允许模型在大多数人比较容易获取的 RTX3090、V100 和 T4 等 GPU 上运行。

有关量化的更多信息，以及如何将模型量化到比 4 位更低显存需求，我们建议查看 GPT-QModel 的实现。

总之，请记住，模型量化是以准确性和（在某些情况下）推理时间为代价来换取更高的内存效率。

如果 GPU 内存对您的用例没有限制，通常无需考虑量化。但是，许多 GPU 无法在没有量化方法的情况下运行 LLM，在这种情况下，4 位和 8 位量化方案是非常有用的工具。

有关更详细的使用信息，我们强烈建议查看 Transformers 量化文档。接下来，让我们看看如何通过使用更好的算法和改进的模型架构来提高计算和内存效率。

2. Flash Attention

当今表现最好的 LLM 共享大致相同的基本架构，该架构由前馈层、激活层、层归一化层，以及最关键的自注意力层组成。

自注意力层对于大语言模型 (LLM) 至关重要，因为它们使模型能够理解输入 token 之间的上下文关系。然而，自注意力层的峰值 GPU 内存消耗随着输入 token 的数量（我们在下文中称之为 序列长度，记为 N）在计算和内存复杂度上呈 二次方 增长。$N$. 虽然这在较短的输入序列（最多 1000 个输入 token）中并不明显，但对于较长的输入序列（大约 16000 个输入 token）来说，它就成了一个严重的问题。

让我们仔细看看。计算长度为$\mathbf{O}$的输入$\mathbf{X}$的自注意力层输出$N$的 KL 估计为$$\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$是注意力层的输入序列。投影$\mathbf{Q}$和$\mathbf{K}$将各自包含$N$个向量，从而产生$\mathbf{QK}^T$大小为$N^2$ .

的 LLM 通常有多个注意力头，从而并行进行多个自注意力计算。假设该 LLM 有 40 个注意力头并以 bfloat16 精度运行，我们可以计算存储$\mathbf{QK^T}$矩阵所需的内存为$40 * 2 * N^2$字节。对于$N=1000$仅需约 50 MB 显存，然而对于$N=16000$我们将需要 19 GB 显存，而对于$N=100,000$我们几乎需要1TB的内存才能存储这些$\mathbf{QK}^T$矩阵。

简而言之，对于大型输入上下文，默认的自注意力算法很快会变得内存消耗过高，令人望而却步。

随着大语言模型（LLM）在文本理解和生成能力上的提升，它们被应用于日益复杂的任务中。模型曾经只能处理几个句子的翻译或摘要，现在它们可以处理整个页面，这要求模型具备处理超长输入的能力。

我们如何摆脱大型输入长度带来的过高内存需求？我们需要一种计算自注意力机制的新方法，从而摆脱$\mathbf{QK}^T$矩阵。 Tri Dao 等人 正是开发了这样一种新算法，并将其称为 Flash Attention。

简单来说，Flash Attention 将$\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$的计算过程进行拆解，通过多次迭代 softmax 计算步骤来计算输出的较小块。$$\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$$sijb​∗Vj​×Softmax(QKi,jT​) for multiple i,j iterations

，其中$s^a_{ij}$和$s^b_{ij}$是一些 softmax 归一化统计量，需要针对每个$i$和$j$ .

请注意，整个 Flash Attention 的实际实现更为复杂，这里进行了极大的简化，因为深入探讨超出了本指南的范围。欢迎读者查阅撰写精良的 Flash Attention 论文以获取更多细节。

这里的主要结论是：

通过追踪 softmax 归一化统计量并使用巧妙的数学方法，Flash Attention 能够产生与默认自注意力层数值完全一致的输出，且其内存成本仅随输入长度呈线性增长。$N$ .

观察公式，人们凭直觉可能会认为 Flash Attention 相比默认自注意力公式一定慢得多，因为需要进行更多的计算。确实，由于必须不断重新计算 softmax 归一化统计量，Flash Attention 相比普通注意力机制需要更多的 FLOPs（浮点运算次数）（如果感兴趣，请参阅论文了解更多详情）。

然而，Flash Attention 在推理过程中要快得多，这归功于它能够显著减少对 GPU 较慢的高带宽内存（VRAM）的需求，转而专注于速度更快的片上内存（SRAM）。

从本质上讲，Flash Attention 确保了所有的中间写入和读取操作都可以使用快速的片上 SRAM 内存完成，而不必为了计算输出向量而去访问较慢的 VRAM 内存。$\mathbf{O}$ .

在实践中，如果 Flash Attention 可用，目前绝对没有任何理由**不**使用它。该算法产生的输出在数学上是相同的，而且速度更快、内存效率更高。

3. 架构创新

到目前为止，我们已经通过以下方式研究了计算和内存效率的提升：

• 将权重转换为低精度格式
• 用更节省内存和计算效率更高的版本替换自注意力算法

现在，让我们探讨如何更改大语言模型的架构，使其对于需要长文本输入的任务（例如：

• 检索增强问答（RAG），
• 摘要生成，
• 聊天

请注意，聊天不仅需要 LLM 处理长文本输入，还要求 LLM 能够高效处理用户和助手之间的来回对话（如 ChatGPT）。

一旦训练完成，基础 LLM 架构就很难改变，因此预先考虑 LLM 的任务并据此优化模型架构非常重要。模型架构中有两个重要组件会迅速成为处理长输入序列时的内存和/或性能瓶颈：

• 位置嵌入（Positional embeddings）
• 键值缓存（Key-value cache）

让我们更详细地了解每个组件：

3.1 改进 LLM 的位置嵌入

自注意力机制将每个 token 与其他 token 建立关联。举例来说，文本输入序列 “Hello”, “I”, “love”, “you” 的$\text{Softmax}(\mathbf{QK}^T)$矩阵可能如下所示：

每个词 token 都会被赋予一个概率质量，用于衡量它对所有其他词 token 的关注程度，从而与其他词 token 建立联系。例如，单词 “love” 对单词 “Hello” 的关注度为 5%，对 “I” 的关注度为 30%，对其自身的关注度为 65%。

一个基于自注意力机制但没有位置嵌入的 LLM，在理解文本输入中各部分之间的位置关系方面会有很大困难。这是因为通过$\mathbf{QK}^T$计算出的概率分数，将每个词 token 与其他所有 token 在$O(1)$的计算下关联起来，而不管它们彼此之间的相对位置距离如何。因此，对于没有位置嵌入的 LLM，每个 token 似乎与所有其他 token 的距离都相同。例如，要区分 “Hello I love you” 和 “You love I hello” 将非常具有挑战性。

为了让 LLM 理解句子顺序，需要额外的 提示（cue），这通常以 位置编码（positional encodings）（也称为 位置嵌入（positional embeddings））的形式实现。位置编码将每个 token 的位置信息编码为一种数值表达，LLM 可以利用它来更好地理解句子顺序。

Attention Is All You Need 论文的作者引入了正弦位置嵌入（sinusoidal positional embeddings）$\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$，其中每个向量$\mathbf{p}_i$都是根据其位置作为正弦函数计算得出的。$i$然后，位置编码只需简单地添加到输入序列向量中$\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$从而提示模型更好地学习句子顺序。

除了使用固定位置嵌入外，其他人（如 Devlin 等人）使用了学习型位置编码，其中位置嵌入$\mathbf{P}$是在训练过程中学习到的。

正弦和学习型位置嵌入曾是 LLM 中编码句子顺序的主要方法，但研究发现这些位置编码存在一些问题：

1. 正弦和学习型位置嵌入都是绝对位置嵌入，即为每个位置 id 编码一个独特的嵌入$0, \ldots, N$。正如 Huang 等人 和 Su 等人 所证明的那样，绝对位置嵌入会导致 LLM 在处理长文本输入时表现不佳。对于长文本输入，如果模型能学习输入 token 彼此之间的相对位置距离，而不是绝对位置，会更有优势。
2. 当使用学习型位置嵌入时，LLM 必须在固定的输入长度 $N$ 上进行训练，这使得将模型推断（extrapolate）到比训练长度更长的输入上变得非常困难。

最近，能够解决上述问题的相对位置嵌入变得更加流行，其中最著名的是：

• 旋转位置嵌入（Rotary Position Embedding, RoPE）
• ALiBi

RoPE 和 ALiBi 都认为，最好直接在自注意力算法中向 LLM 提示句子顺序，因为正是在那里，词 token 彼此建立了关联。更具体地说，应该通过修改$\mathbf{QK}^T$计算来提示句子顺序。

无需赘述太多细节，RoPE 指出位置信息可以编码到查询-键（query-key）对中，例如$\mathbf{q}_i$和$\mathbf{x}_j$通过将每个向量旋转一个角度$\theta * i$和$\theta * j$，其中$i, j$描述了每个向量在句子中的位置。$$\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$$

$\mathbf{R}_{\theta, i - j}$从而代表一个旋转矩阵。$\theta$在训练期间是不会被学习的，而是设置为一个预定义的值，该值取决于训练期间的最大输入序列长度。

通过这样做，两者之间的概率得分$\mathbf{q}_i$和$\mathbf{q}_j$仅在以下情况下受到影响：$i \ne j$并且仅取决于相对距离$i - j$而与每个向量的具体位置无关。$i$和$j$ .

RoPE 被用于当今多个最重要的 LLM 中，例如

• Falcon
• Llama
• PaLM

作为替代方案，ALiBi 提出了一种简单得多的相对位置编码方案。输入 token 之间的相对距离被添加为通过预定义值 m 缩放后的负整数，添加到每个 query-key 条目中$\mathbf{QK}^T$在 softmax 计算之前的矩阵中。

正如 ALiBi 论文所示，这种简单的相对位置编码允许模型即使在非常长的文本输入序列下也能保持高性能。

ALiBi 被用于当今多个最重要的 LLM 中，例如

• MPT
• BLOOM

RoPE 和 ALiBi 位置编码都可以外推到训练期间未见过的输入长度，而事实证明，与 RoPE 相比，ALiBi 的外推效果要好得多。对于 ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度即可。对于 RoPE，保留与$\theta$训练期间所用相同的参数会导致在处理比训练时所见文本输入长得多的输入时效果不佳，详见 Press 等人的论文。然而，社区已经发现了一些有效的技巧来适配 $\theta$，从而使 RoPE 位置嵌入能够很好地适用于外推的文本输入序列（参见此处）。

RoPE 和 ALiBi 都是相对位置嵌入，它们在训练期间不会被学习，而是基于以下直觉：

• 关于文本输入的位置线索应该直接提供给$\mathbf{QK}^T$自注意力层的矩阵。
• LLM 应该被激励去学习一种恒定的相对距离位置编码。
• 文本输入 token 之间距离越远，其 query-value 概率就越低。RoPE 和 ALiBi 都会降低相距较远的 token 的 query-key 概率。RoPE 通过增大 query-key 向量之间的角度来减小它们的向量积，从而降低概率。ALiBi 则通过向向量积添加大的负数来降低概率。

总之，旨在部署在需要处理大量文本输入的任务中的 LLM，最好使用相对位置嵌入（如 RoPE 和 ALiBi）进行训练。还要注意，即使使用 RoPE 和 ALiBi 的 LLM 仅在固定的长度上进行了训练，例如$N_1 = 2048$它仍然可以在实践中用于远大于 $N_1$ 的文本输入，例如 $N_2 = 8192 > N_1$，通过对位置嵌入进行外推。

3.2 键值缓存 (Key-value cache)

使用 LLM 进行自回归文本生成的工作方式是：迭代输入一个序列，采样下一个 token，将该 token 追加到输入序列中，并持续此过程，直到 LLM 生成了一个标志着生成结束的 token。

请查看 Transformer 的文本生成教程，以获得关于自回归生成工作原理的更直观的解释。

让我们运行一个简短的代码片段来演示自回归在实践中是如何工作的。我们将简单地通过 torch.argmax 获取可能性最高的下一个 token。

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

输出:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

正如我们所见，每次我们都会通过刚刚采样的 token 来增加文本输入。

除极少数例外，LLM 都是使用因果语言建模目标进行训练的，因此会屏蔽注意力得分矩阵的上三角部分——这就是为什么在上面的两个图中，注意力得分留白（即概率为 0）的原因。有关因果语言建模的快速回顾，您可以参考图解自注意力博客。

因此，token 永远不会依赖于后续的 token，更具体地说，$\mathbf{q}_i$向量永远不会与任何键 (key)、值 (value) 向量建立关系$\mathbf{k}_j, \mathbf{v}_j$如果 $j > i$。相反，$\mathbf{q}_i$仅关注之前的键值向量 $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , 对于 } m \in {0, \ldots i - 1}$。为了减少不必要的计算，可以缓存每一层之前所有时间步的键值向量。

在下文中，我们将指示 LLM 使用键值缓存，方法是在每次前向传递时检索并转发它。在 Transformers 中，我们可以通过将 use_cache 标志传递给 forward 调用来检索键值缓存，然后将其与当前 token 一起传递。

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

输出:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

可以看出，使用键值缓存时，文本输入 token 的长度不会增加，而是保持为单个输入向量。另一方面，键值缓存的长度在每个解码步骤都会增加一个。

使用键值缓存意味着$\mathbf{QK}^T$本质上简化为$\mathbf{q}_c\mathbf{K}^T$，其中$\mathbf{q}_c$它是当前传递的输入 token 的 query 投影，该 token 始终只是一个单一向量。

使用键值缓存有两个优点：

• 计算效率显著提高，因为与计算完整的$\mathbf{QK}^T$矩阵相比，执行的计算量更少。这带来了推理速度的提升。
• 所需的最大内存不会随生成 token 的数量呈二次方增长，而仅呈线性增长。

始终应该利用键值缓存，因为它不仅能产生相同的结果，还能为长输入序列带来显著的速度提升。Transformers 在使用 text pipeline 或 generate 方法时，默认启用键值缓存。我们专门写了一份关于缓存的完整指南，请见此处。

注意，尽管我们建议使用键值缓存，但当您使用它们时，LLM 的输出可能会略有不同。这是矩阵乘法内核本身的特性——您可以在此处阅读更多相关信息。

3.2.1 多轮对话

键值缓存对于聊天等需要多次自回归解码的任务尤其有用。让我们看一个例子。

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

在此聊天中，LLM 运行了两次自回归解码：

1. 第一次，键值缓存为空，输入提示词是 "User: How many people live in France?"，模型自回归地生成文本 "Roughly 75 million people live in France"，并在每个解码步骤中增加键值缓存。
2. 第二次输入提示词是 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?"。由于缓存的存在，前两个句子的所有键值向量都已经计算完毕。因此，输入提示词仅包含 "User: And how many in Germany?"。在处理简短的输入提示词时，其计算出的键值向量会连接到第一次解码的键值缓存中。第二个助手的回答 "Germany has ca. 81 million inhabitants" 是在键值缓存的基础上自动生成的，该缓存包含了 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?" 的编码键值向量。

这里有两点需要注意：

1. 对于部署在聊天中的 LLM 来说，保留所有上下文至关重要，这样 LLM 才能理解对话的所有先前上下文。例如在上述示例中，LLM 需要理解当用户询问 "And how many are in Germany" 时，是指人口数量。
2. 键值缓存对于聊天极为有用，因为它允许我们持续增长编码后的聊天历史，而不必从头开始重新编码整个历史（例如在使用编码器-解码器架构时会出现这种情况）。

在 transformers 中，当传入 return_dict_in_generate=True 时，generate 调用将返回 past_key_values，此外还有默认的 use_cache=True。请注意，目前无法通过 pipeline 接口使用此功能。

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

输出:

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

太棒了，重新计算注意力层的相同键和值无需花费额外时间！然而有一个问题。虽然计算$\mathbf{QK}^T$矩阵所需的峰值内存显著减少了，但在长输入序列或多轮对话中，将键值缓存保留在内存中可能会变得非常昂贵。请记住，键值缓存需要存储之前所有输入向量的键值向量$\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$用于所有自注意力层和所有注意力头。

让我们计算一下我们之前使用的 LLM bigcode/octocoder 的键值缓存中需要存储的浮点数值数量。浮点数值的数量等于：两倍的序列长度 × 注意力头的数量 × 注意力头维度 × 层数。假设输入序列长度为 16000 进行计算：

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

输出:

7864320000

大约 80 亿个浮点数值！以 float16 精度存储 80 亿个浮点数值需要大约 15 GB 的 RAM，这大约是模型权重自身大小的一半！研究人员提出了两种可以显著降低存储键值缓存内存成本的方法，将在接下来的小节中探讨。

3.2.2 多查询注意力 (Multi-Query-Attention, MQA)

多查询注意力 (Multi-Query-Attention) 是在 Noam Shazeer 的论文《Fast Transformer Decoding: One Write-Head is All You Need》中提出的。正如标题所言，Noam 发现与其使用 n_head 个键值投影权重，不如使用一对跨所有注意力头共享的单头值投影权重，且不会导致模型性能明显下降。

通过使用一对单头值投影权重，键值向量$\mathbf{k}_i, \mathbf{v}_i$在所有注意力头中必须相同，这意味着我们只需要在缓存中存储 1 对键值投影，而不是 n_head 对。

由于大多数 LLM 使用 20 到 100 个注意力头，MQA 显著降低了键值缓存的内存消耗。对于本笔记本中使用的 LLM，在 16000 的输入序列长度下，我们可以将所需的内存消耗从 15 GB 降低到 400 MB 以下。

除了节省内存外，MQA 还提高了计算效率，如下所述。在自回归解码中，大型键值向量需要不断重新加载，与当前的键值向量对连接，然后馈入到$\mathbf{q}_c\mathbf{K}^T$计算中。对于自回归解码，不断重新加载所需的内存带宽可能成为严重的性能瓶颈。通过缩小键值向量的尺寸，所需访问的内存变少，从而减少了内存带宽瓶颈。有关更多详细信息，请查看 Noam 的论文。

这里需要理解的重要一点是，将键值注意力头的数量减少到 1 仅在使用了键值缓存时才有意义。如果不使用键值缓存，模型单次前向传递的峰值内存消耗保持不变，因为每个注意力头仍然拥有唯一的查询向量，所以每个注意力头仍然拥有不同的$\mathbf{QK}^T$矩阵。

MQA 已被社区广泛采用，现在被许多最受欢迎的 LLM 所使用。

• Falcon
• PaLM
• MPT
• BLOOM

此外，本笔记本中使用的检查点 bigcode/octocoder 也使用了 MQA。

3.2.3 分组查询注意力 (Grouped-Query-Attention, GQA)

分组查询注意力 (Grouped-Query-Attention) 是由 Google 的 Ainslie 等人提出的。他们发现与使用常规的多键值头投影相比，使用 MQA 往往会导致质量下降。该论文指出，通过不那么剧烈地减少查询头投影权重的数量，可以保留更多的模型性能。与其仅使用单个键值投影权重，不如使用 n < n_head 个键值投影权重。通过选择比 n_head 小得多的 n 值（例如 2、4 或 8），几乎可以保留 MQA 的所有内存和速度收益，同时牺牲更少的模型容量，从而可能牺牲更少的性能。

此外，GQA 的作者发现，现有的模型检查点可以进行升级训练 (uptraining)，使其拥有 GQA 架构，而所需的计算量仅为原始预训练计算量的 5%。虽然 5% 的原始预训练计算量仍然可能是一个巨大的数字，但 GQA 升级训练允许现有的检查点适用于更长的输入序列。

GQA 是最近才提出的，这就是为什么在撰写本笔记本时它的采用率较低。GQA 最显著的应用是 Llama-v2。

总之，如果 LLM 部署在自回归解码场景中，并且需要处理大型输入序列（例如聊天），强烈建议使用 GQA 或 MQA。

结论

研究社区不断提出新的、巧妙的方法来加速不断增大的 LLM 的推理时间。例如，一种很有前途的研究方向是推测解码 (speculative decoding)，其中“简单 token”由更小、更快的语言模型生成，而只有“困难 token”由 LLM 本身生成。深入探讨超出了本笔记本的范围，但可以在这篇精彩的博客文章中阅读更多内容。

像 GPT3/4、Llama-2-70b、Claude、PaLM 这样的大型 LLM 在 Hugging Face Chat 或 ChatGPT 等聊天界面中能运行得如此之快，很大程度上要归功于上述在精度、算法和架构方面的改进。展望未来，GPU、TPU 等加速器只会变得更快并提供更多的内存，但人们仍然应该始终确保使用最佳的算法和架构，以实现最佳的性价比 🤗