优化 LLM 的速度与内存

GPT3/4、Falcon和Llama等大型语言模型（LLM）在处理以人为中心任务方面的能力正在快速提升，使其成为现代知识密集型行业中不可或缺的工具。然而，将这些模型部署到实际任务中仍然充满挑战：

要展现出接近人类的文本理解与生成能力，LLM 目前需要由数十亿个参数组成（参见 Kaplan et al、Wei et. al）。这相应地放大了推理过程中的内存需求。
在许多实际任务中，LLM 需要被提供大量的上下文信息。这就要求模型在推理阶段具备处理超长输入序列的能力。

这些挑战的核心在于增强 LLM 的计算和内存能力，尤其是在处理大规模输入序列时。

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

更低精度：研究表明，以 8 位和 4 位等降低的数值精度运行，可以在模型性能没有显著下降的情况下获得计算优势。
Flash Attention：Flash Attention 是注意力算法的一种变体，它不仅提供了更节省内存的方法，还通过优化 GPU 内存利用率实现了更高的效率。
架构创新： 鉴于 LLM 在推理阶段始终以相同方式部署，即基于长输入上下文的自回归文本生成，业界已提出多种专用模型架构以实现更高效的推理。其中，模型架构方面最重要的进展包括 Alibi、Rotary embeddings、Multi-Query Attention (MQA) 和 Grouped-Query-Attention (GQA)。

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

1. 较低精度

要理解大语言模型（LLM）的内存需求，最好将其视为一组权重矩阵和向量，而将文本输入视为一系列向量。在下文中，weights 一词将用于指代所有模型权重矩阵和向量。

在撰写本指南时，LLM 至少包含数十亿个参数。每个参数都由一个小数表示，例如 4.5689，通常以 float32、bfloat16 或 float16 格式存储。这使得我们可以轻松计算将 LLM 加载到内存中所需的内存量：

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

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

加载具有 X 十亿参数的模型权重，在 bfloat16/float16 精度下大约需要 2 * X GB 的 VRAM

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

举几个例子，说明以 bfloat16 格式加载模型大致需要多少 VRAM：

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 的 VRAM。前文列出的大多数模型仅加载就需要超过 80GB 的显存，因此必然需要张量并行和/或流水线并行。

🤗 Transformers 现已支持在各自的配置类中包含 base_tp_plan 的受支持模型的张量并行。有关张量并行的更多信息，请参见此处。此外，如果您有兴趣以支持张量并行的方式编写模型，不妨查看 text-generation-inference 库。

开箱即用地支持朴素的流水线并行。只需使用 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。让我们来试一试。

我们首先加载模型和分词器，然后将两者都传入 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.accelerator.memory.max_memory_allocated 来测量加速器内存的峰值分配量。

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

输出：

29.0260648727417

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

如今几乎所有模型都在 bfloat16 下训练，如果你的 GPU 支持 bfloat16，就没有理由以完整的 float32 精度运行模型。Float32 不会比训练模型时使用的精度带来更好的推理结果。

如果您不确定模型权重在 Hub 上是以哪种格式存储的，可以随时查看检查点配置中的 "dtype" 字段，例如此处。使用 from_pretrained(..., dtype=...) 加载模型时，建议将模型设置为与配置文件中相同的精度类型；但如果原始类型为 float32，则推理时可以使用 float16 或 bfloat16。

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

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.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 et al.）。如最近的 GPTQ 论文所示，模型甚至可以量化为 3 位或 2 位，且性能损失在可接受范围内 🤯。

不深入细节，量化方案旨在降低权重的精度，同时尽可能保持模型推理结果的准确性（即尽可能接近 bfloat16）。请注意，量化在文本生成任务中效果尤为显著，因为我们只关心选择最可能的下一个 token 集合，而并不真正在意下一个 token 的 logit 分布的具体数值。关键在于下一个 token 的 logit 分布大致保持不变，从而确保 argmax 或 topk 操作能得出相同的结果。

存在多种量化技术，此处不再详细展开，但总体而言，所有量化技术的工作原理如下：

1. 将所有权重量化到目标精度
1. 加载量化后的权重，并以 bfloat16 精度传入输入向量序列
1. 动态将权重反量化为 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.accelerator.max_memory_allocated())

输出：

15.219234466552734

显著减少！内存占用已降至略高于 15 GB，因此我们可以在 4090 等消费级 GPU 上运行该模型。我们观察到内存效率有了很好的提升，且模型输出几乎没有退化。不过，我们也能注意到推理过程中出现了轻微的减速。

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

del model
del pipe

flush()

让我们看看 4-bit 量化带来的峰值 GPU 内存消耗。将模型量化为 4-bit 可以使用与之前相同的 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.accelerator.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 显存从 32GB 降低至仅 15GB，而以 4 位精度运行该模型可进一步将所需的 GPU 显存降低至略高于 9GB。

4 位量化使模型能够在 RTX3090、V100 和 T4 等 GPU 上运行，这些 GPU 对大多数人来说都很容易获取。

如需了解更多有关量化的信息，并了解如何将模型量化为比 4-bit 占用更少的 GPU 显存，我们建议研究 GPT-QModel 的实现。

总而言之，请记住，模型量化是以牺牲精度（在某些情况下还包括推理时间）为代价来换取更高的内存效率的。

如果 GPU 显存不是您使用场景的限制因素，通常无需考虑量化。然而，许多 GPU 在不使用量化方法的情况下根本无法运行 LLM，在这种情况下，4 位和 8 位量化方案是非常有用的工具。

如需更详细的使用说明，我们强烈建议您参阅Transformers 量化文档。接下来，让我们看看如何通过采用更优的算法和改进的模型架构来提升计算和内存效率。

2. Flash Attention

当今表现最佳的 LLM 在架构上大同小异，均由前馈层、激活层、层归一化层以及最为关键的自注意力层组成。

自注意力层是大语言模型（LLM）的核心，它使模型能够理解输入 token 之间的上下文关系。然而，自注意力层的 GPU 峰值内存消耗 呈二次方增长 ，其计算和内存复杂度均随输入 token 数量（也称为 序列长度) ，下文用 $N$ . 表示）的增加而增加。对于较短的输入序列（最多 1000 个输入 token），这一点并不明显，但对于较长的输入序列（约 16000 个输入 token），这便成为一个严重的问题。

让我们更仔细地看一下。计算自注意力层输出 $\mathbf{O}$ 的公式，对于长度为 $\mathbf{X}$ 的输入 $N$ 如下： $\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 的 VRAM，但对于 $N=16000$ ，则需要 19 GB 的 VRAM，而对于 $N=100,000$ ，仅存储 $\mathbf{QK}^T$ 矩阵就需要近 1TB 的内存。

简而言之，默认的自注意力算法在处理大型输入上下文时，其内存开销会迅速变得难以承受。

随着大语言模型在文本理解和生成方面的能力不断提升，它们被应用于越来越复杂的任务。过去，模型只需处理几句话的翻译或摘要，而现在则需要处理整页内容，这就要求模型具备处理超长输入的能力。

我们如何消除长输入序列带来的高昂内存需求？我们需要一种新的自注意力机制计算方法，以消除 $\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}$

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

请注意，完整的 Flash Attention 实现要复杂得多，本指南为了控制篇幅，对其进行了大幅简化。如需了解更多细节，建议读者参阅撰写详尽的 Flash Attention 论文。

这里的主要要点是：

通过跟踪 softmax 归一化统计量并运用一些巧妙的数学方法，Flash Attention 能够 在数值上完全一致地 输出与默认自注意力层相同的结果，其内存开销仅随 $N$ .

从公式来看，人们会直觉地认为 Flash Attention 比默认的自注意力公式慢得多，因为它需要执行更多的计算。事实上，与常规注意力相比，Flash Attention 确实需要更多的 FLOPs，因为 softmax 归一化统计量必须不断重新计算（如有兴趣，请参阅论文了解更多细节）

然而，与默认的注意力机制相比，Flash Attention 在推理阶段的速度要快得多，这得益于它能够显著降低对 GPU 上较慢的高带宽内存（VRAM）的需求，转而利用速度更快的片上内存（SRAM）。

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

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

3. 架构创新

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

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

现在，让我们探讨如何修改 LLM 的架构，使其在处理需要长文本输入的任务时最为高效，例如：

检索增强问答，
摘要，
聊天

请注意，chat 不仅要求 LLM 能够处理长文本输入，还要求 LLM 能够高效地处理用户与助手之间的多轮对话（例如 ChatGPT）。

模型训练完成后，其基础 LLM 架构很难再更改，因此提前考虑 LLM 的任务并据此优化模型架构至关重要。模型架构中有两个重要组件，在处理长输入序列时往往会迅速成为内存和/或性能瓶颈。

位置嵌入
键值缓存

让我们更详细地介绍每个组件

3.1 改进 LLM 的位置嵌入

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

每个词元（word token）都会分配一个概率质量，用于关注所有其他词元，从而与所有其他词元建立关联。例如，单词 “love” 以 5% 的概率关注单词 “Hello”，以 30% 的概率关注 “I”，并以 65% 的概率关注自身。

基于自注意力的 LLM 如果没有位置嵌入，将很难理解文本输入中各个词之间的相对位置关系。这是因为通过 $\mathbf{QK}^T$ 计算出的概率分数将每个词元与其他词元相关联，且 $O(1)$ 的计算复杂度为 e.g. ，与它们之间的相对位置距离无关。因此，对于没有位置嵌入的 LLM 而言，每个词元看起来与其他所有词元的距离都是相同的， “例如， 区分 ““Hello I love you” 和

为了让 LLM 理解句子顺序，需要额外的线索，通常以位置编码（也称为位置嵌入）的形式应用。位置编码将每个 token 的位置编码为数值表示，使 LLM 能够利用这些信息更好地理解句子顺序。

《 Attention Is All You Need 》一文的作者引入了正弦位置嵌入 $\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 编码句子顺序的主要方法，但人们发现这些位置编码存在一些问题：

正弦位置编码和可学习位置编码都属于绝对位置编码， i.e. 即为每个位置 ID 编码一个唯一的嵌入向量： $0, \ldots, N$ . 如 Huang 等人和 Su 等人, 所示，绝对位置编码会导致 LLM 在处理长文本输入时性能下降。对于长文本输入，如果模型能够学习输入 token 之间的相对位置距离，而不是它们的绝对位置，则更有优势。
在使用学习得到的位置嵌入时，LLM 必须在固定的输入长度 $N$ 上进行训练，这使得模型难以外推到比训练时更长的输入长度。

最近，能够解决上述问题的相对位置嵌入变得越来越流行，其中最值得注意的是：

两者 RoPE 和 ALiBi 都主张，最好在自注意力算法中直接提示 LLM 句子顺序，因为正是在这里，词元（token）之间建立了关联关系。更具体地说，句子顺序应通过修改 $\mathbf{QK}^T$ 计算来提示。

不深入细节， RoPE 指出位置信息可以编码到查询-键向量对中， e.g. $\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 所使用，例如：

作为替代方案， ALiBi 提出了一种更简单的相对位置编码方案。输入 token 之间的相对距离会以一个预定义值 m 缩放后的负整数形式，添加到 $\mathbf{QK}^T$ 矩阵中，就在 softmax 计算之前。

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

ALiBi 被当今多个最重要的 LLM 所采用，例如：

两者 RoPE 和 ALiBi 位置编码都可以外推到训练期间未见过的输入长度，而研究表明， ALiBi 在开箱即用的情况下，其外推效果远优于 RoPE. 对于 ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度即可。对于 RoPE, θ $\theta$ 在训练时使用的 c.f Press et al.. 然而，社区发现了一些有效的方法来调整 RoPE 位置嵌入，使其能够很好地处理外推的文本输入序列（参见此处).

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

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

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

3.2 键值缓存

使用 LLM 进行自回归文本生成的过程是：迭代地输入一个序列，采样下一个 token，将该 token 追加到输入序列中，并持续执行此操作，直到 LLM 生成一个表示生成已结束的 token。

请参阅 Transformer 文本生成教程，以获取关于自回归生成工作原理的更直观说明。

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

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

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 追加到文本输入 token 中。

除极少数例外，LLM 均使用因果语言建模目标进行训练，因此会对注意力得分矩阵的上三角部分进行掩码处理——这就是为什么上述两个图中注意力得分被留空（即概率为 0）。如需快速回顾因果语言建模，请参阅 Illustrated Self Attention 博客。

因此，token 从不依赖于后续的 token，更具体地说， $\mathbf{q}_i$ 向量从不与任何键或值向量 $\mathbf{k}_j, \mathbf{v}_j$ 建立关联（当 $j > i$ 时）。相反， $\mathbf{q}_i$ 仅关注之前的键值向量 $\mathbf{k}${m < i}, \mathbf{v}${m < i} \text{ , for } m \in {0, \ldots i - 1}$。为了减少不必要的计算，因此可以缓存每一层在所有先前时间步的键值向量。

接下来，我们将指示 LLM 通过在每个前向传播中检索并转发键值缓存来利用该缓存。在 Transformers 中，我们可以在调用 forward 时传入 use_cache 标志来检索键值缓存，然后将其与当前 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(model.device)

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']

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

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

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

与计算完整的 $\mathbf{QK}^T$ 矩阵相比，计算量大幅减少，从而显著提升了计算效率，并加快了推理速度。
所需最大内存不会随生成 token 数量的增加而呈二次方增长，而是仅呈线性增长。

应当始终使用键值缓存，因为它能产生相同的结果，并显著加速较长输入序列的处理。在使用文本流水线或 generate 方法时，Transformers 默认启用键值缓存。我们有一篇专门介绍缓存的指南，见此处。

请注意，尽管我们建议使用键值缓存，但使用它们时，您的 LLM 输出可能会有细微差异。这是矩阵乘法内核本身的特性——您可以在此处阅读更多相关内容 here。

3.2.1 多轮对话

键值缓存（key-value cache）对于需要多次自回归解码的应用（例如聊天）尤其有用。下面来看一个示例。

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 执行了两次自回归解码：

第一次，键值缓存为空，输入提示为 "User: How many people live in France?"，模型以自回归方式生成文本 "Roughly 75 million people live in France"，并在每个解码步骤中不断扩充键值缓存。
第二次输入提示为 "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?" 组成。在处理缩短后的输入提示时，其计算出的键值向量会被拼接到第一次解码的键值缓存中。随后，第二个 Assistant 的回答 "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?" 的编码键值向量。

这里有两点需要注意：

对于部署在聊天场景中的 LLM，保留所有上下文至关重要，以便 LLM 能够理解对话中所有先前的上下文。例如，对于上述示例，当用户询问 "And how many are in Germany" 时，LLM 需要理解用户指的是人口数量。
键值缓存（key-value cache）在对话场景中极为有用，因为它允许我们持续扩展已编码的对话历史，而无需每次都从头重新编码整个对话历史（例如，在使用编码器-解码器架构时就需要这样做）。

在 transformers 中，当传入 return_dict_in_generate=True 时，generate 调用除了默认的 use_cache=True 之外，还会返回 past_key_values。请注意，目前还无法通过 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 的键值缓存中需要存储的浮点值数量。浮点值的数量等于序列长度乘以注意力头数乘以注意力头维度再乘以层数，最后乘以 2。假设输入序列长度为 16000，针对我们的 LLM 进行计算，结果为：

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

输出：

7864320000

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

3.2.2 多查询注意力（MQA）

Multi-Query-Attention 由 Noam Shazeer 在其论文 Fast Transformer Decoding: One Write-Head is All You Need 中提出。正如标题所示，Noam 发现，与其使用 n_head 组 key-value 投影权重，不如使用一组在所有注意力头之间共享的单一 key-value 投影权重，而不会导致模型性能显著下降。

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

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

除了节省内存，MQA 还能提升计算效率，具体说明如下。在自回归解码过程中，需要反复加载大型键值向量，并将其与当前的键值向量对拼接后，输入到 $\mathbf{q}_c\mathbf{K}^T$ 计算中。对于自回归解码，这种持续加载所需的内存带宽可能成为严重的时间瓶颈。通过减小键值向量的规模，可以减少需要访问的内存量，从而缓解内存带宽瓶颈。更多详情请参阅 Noam 的论文.

这里需要理解的关键点是，将键值注意力头数量减少到 1 只有在使用了键值缓存的情况下才有意义。在没有键值缓存的情况下，模型单次前向传播的峰值内存占用保持不变，因为每个注意力头仍然拥有唯一的查询向量，因此每个注意力头仍然具有不同的 $\mathbf{QK}^T$ 矩阵。

MQA 在社区中得到了广泛采用，目前已被许多最流行的 LLM 使用：

此外，本 notebook 中使用的检查点 - bigcode/octocoder - 采用了 MQA。

3.2.3 分组查询注意力（GQA）

Google 的 Ainslie 等人提出的 Grouped-Query-Attention 发现，与使用标准的 multi-key-value head 投影相比，使用 MQA 往往会导致质量下降。该论文认为，通过不那么激进地减少 query head 投影权重的数量，可以保留更多的模型性能。与其仅使用单个 key-value 投影权重，不如使用 n < n_head 个 key-value 投影权重。通过将 n 选择为远小于 n_head 的值，例如 2、4 或 8，几乎可以保留 MQA 带来的所有内存和速度优势，同时牺牲更少的模型容量，从而可以说性能损失也更小。

此外，GQA 的作者发现，现有的模型检查点可以通过 uptraining 转换为 GQA 架构，所需的计算量仅为原始预训练计算量的 5%。尽管 5% 的原始预训练计算量仍然可能非常庞大，但 GQA uptraining 使现有检查点能够支持更长的输入序列。

GQA 是最近才提出的，因此在本笔记本撰写时，其采用率还比较低。 GQA 最显著的应用是 Llama-v2。

综上所述，如果 LLM 采用自回归解码进行部署，并且需要处理大型输入序列（例如聊天场景），则强烈建议使用 GQA 或 MQA。

结论

研究社区不断提出新颖巧妙的方法来加速规模日益增大的 LLM 的推理过程。例如，一个颇具前景的研究方向是投机解码，其中“简单 token”由更小、更快的语言模型生成，而只有“困难 token”才由 LLM 本身生成。更深入的细节超出了本 notebook 的范围，但可以在这篇精彩的博文中阅读。

GPT3/4、Llama-2-70b、Claude、PaLM 等大规模 LLM 之所以能在 Hugging Face Chat 或 ChatGPT 等聊天界面中快速运行，很大程度上得益于上述在精度、算法和架构方面的改进。未来，GPU、TPU 等加速器只会变得更快并支持更大的内存，但无论如何，都应始终确保使用当前可用的最佳算法和架构，以获得最高的性价比 🤗

Update on GitHub