优化LLM的速度和内存

大型语言模型 (LLM)，如 GPT3/4、Falcon 和 Llama，在处理以人为中心的任务方面正在迅速进步，并将自己确立为现代知识型产业中必不可少的工具。然而，在实际任务中部署这些模型仍然具有挑战性

为了展现接近人类的文本理解和生成能力，LLM 目前需要由数十亿参数组成（参见 Kaplan et al, Wei et. al）。这因此增加了推理的内存需求。
在许多实际任务中，LLM 需要被赋予广泛的上下文信息。这使得模型需要在推理过程中具备管理非常长的输入序列的能力。

这些挑战的关键在于增强 LLM 的计算和内存能力，尤其是在处理扩展的输入序列时。

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

降低精度： 研究表明，以降低的数值精度（即 8 位和 4 位）运行可以在不显著降低模型性能的情况下实现计算优势。
Flash Attention： Flash Attention 是注意力算法的一种变体，它不仅提供了一种更节省内存的方法，而且由于优化了 GPU 内存利用率，还实现了更高的效率。
架构创新： 考虑到 LLM 在推理期间始终以相同的方式部署，即使用长输入上下文的自回归文本生成，因此已经提出了专门的模型架构，以实现更高效的推理。模型架构方面最重要的进展是 Alibi、Rotary embeddings、多查询注意力 (MQA) 和分组查询注意力 (GQA)。

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

1. 降低精度

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

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

加载具有 X 十亿参数的模型的权重大约需要 4*X GB 的 VRAM（float32 精度）

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

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

对于较短的文本输入（少于 1024 个 tokens），推理的内存需求很大程度上受加载权重的内存需求支配。因此，现在，让我们假设推理的内存需求等于将模型加载到 GPU VRAM 中的内存需求。

为了给出一些在 bfloat16 中加载模型大致需要多少 VRAM 的示例

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

在编写本文档时，市场上最大的 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 VRAM。让我们试一试。

我们首先加载模型和 tokenizer，然后将两者传递给 Transformers 的 pipeline 对象。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_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.max_memory_allocated 来测量峰值 GPU 内存分配。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

29.0260648727417

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

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

如果您不确定模型权重在 Hub 上以哪种格式存储，您可以随时查看检查点的配置中的 "torch_dtype"，例如此处。建议在加载时使用 from_pretrained(..., torch_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 的 VRAM 怎么办？研究发现，模型权重可以量化为 8 位或 4 位，而性能损失不大（参见 Dettmers et al.）。正如最近的 GPTQ 论文 🤯 中所示，模型甚至可以量化为 3 位或 2 位，且性能损失可以接受。

在不深入太多细节的情况下，量化方案旨在降低权重的精度，同时尝试使模型的推理结果尽可能准确（又名尽可能接近 bfloat16）。请注意，量化对于文本生成特别有效，因为我们关心的只是选择最有可能的下一组 tokens，而不太关心下一个 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", 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", load_in_4bit=True, low_cpu_mem_usage=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 亿参数的模型来说，这真的不多。

虽然我们在这里的模型中看到精度下降非常小，但在实践中，与 8 位量化或完整的 bfloat16 推理相比，4 位量化通常会导致不同的结果。这取决于用户尝试。

另请注意，与 8 位量化相比，此处的推理再次稍慢，这是由于用于 4 位量化的更积极的量化方法导致 $\text{quantize}$ 和 $\text{dequantize}$ 在推理期间花费更长的时间。

del model
del pipe

flush()

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

4 位量化允许模型在 RTX3090、V100 和 T4 等 GPU 上运行，这些 GPU 对大多数人来说都非常容易获得。

有关量化的更多信息，以及了解如何量化模型以使其所需的 GPU VRAM 内存甚至低于 4 位，我们建议您研究 AutoGPTQ 实现。

总而言之，重要的是要记住，模型量化是以牺牲精度以及在某些情况下牺牲推理时间为代价来换取提高的内存效率。

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

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

2. Flash Attention

当今性能最佳的 LLM 或多或少共享相同的基本架构，该架构由前馈层、激活层、层归一化层以及最关键的自注意力层组成。

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

让我们仔细看看。计算输出的公式 $\mathbf{O}$ 的自注意力层，对于输入 $\mathbf{X}$ 长度为 $N$ 是 $O = Attn(X) = V \times Softmax(QK T) 其中 Q = W q X, V = W v X, K = W k X$ $X = (x 1, ... x N)$ 由此可见，X 是注意力层的输入序列。投影 $Q$ 和 $K$ 每一个都将由以下内容组成 $N$ 向量，从而得到 $QK T$ 大小为 $N 2$ .

通常，LLM 有多个注意力头，因此并行执行多个自注意力计算。假设 LLM 有 40 个注意力头，并在 bfloat16 精度下运行，我们可以计算存储以下内容所需的内存： $QK T$ 矩阵为 $40 * 2 * N 2$ 字节。对于 $N=1000$ 只需要大约 50 MB 的 VRAM，但是，对于 $N=16000$ 我们将需要 19 GB 的 VRAM，而对于 $N=100,000$ 我们将需要将近 1TB 才能存储 $QK T$ 矩阵。

长话短说，默认的自注意力算法对于大型输入上下文很快就会变得内存开销过大。

随着 LLM 在文本理解和生成方面的改进，它们被应用于越来越复杂的任务。虽然模型曾经处理几句话的翻译或摘要，但现在它们管理整页内容，需要处理大量输入长度的能力。

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

简而言之，Flash Attention 打破了 $V \times Softmax(QK T$ ) 计算，而是通过迭代多个 softmax 计算步骤来计算较小的输出块 $O i \leftarrow s a ij * O i + s b ij * V j \times Softmax(QK T i,j) 对于多个 i, j 迭代$

其中 $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 可用，则目前绝对没有理由不使用它。该算法在数学上给出了相同的输出，并且速度更快、内存效率更高。

让我们看一个实际的例子。

我们的 OctoCoder 模型现在收到了一个明显更长的输入提示，其中包括所谓的系统提示。系统提示用于引导 LLM 成为更好的助手，以更好地适应用户的任务。在下面，我们使用一个系统提示，使 OctoCoder 成为更好的编码助手。

system_prompt = """Below are a series of dialogues between various people and an AI technical assistant.
The assistant tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble but knowledgeable.
The assistant is happy to help with code questions and will do their best to understand exactly what is needed.
It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer.
That said, the assistant is practical really does its best, and doesn't let caution get too much in the way of being useful.

The Starcoder models are a series of 15.5B parameter models trained on 80+ programming languages from The Stack (v1.2) (excluding opt-out requests).
The model uses Multi Query Attention, was trained using the Fill-in-the-Middle objective, and with 8,192 tokens context window for a trillion tokens of heavily deduplicated data.

-----

Question: Write a function that takes two lists and returns a list that has alternating elements from each input list.

Answer: Sure. Here is a function that does that.

def alternating(list1, list2):
   results = []
   for i in range(len(list1)):
       results.append(list1[i])
       results.append(list2[i])
   return results

Question: Can you write some test cases for this function?

Answer: Sure, here are some tests.

assert alternating([10, 20, 30], [1, 2, 3]) == [10, 1, 20, 2, 30, 3]
assert alternating([True, False], [4, 5]) == [True, 4, False, 5]
assert alternating([], []) == []

Question: Modify the function so that it returns all input elements when the lists have uneven length. The elements from the longer list should be at the end.

Answer: Here is the modified function.

def alternating(list1, list2):
   results = []
   for i in range(min(len(list1), len(list2))):
       results.append(list1[i])
       results.append(list2[i])
   if len(list1) > len(list2):
       results.extend(list1[i+1:])
   else:
       results.extend(list2[i+1:])
   return results

-----
"""

为了演示目的，我们将系统提示重复了十次，以便输入长度足够长，以观察 Flash Attention 的内存节省效果。我们在后面附加了原始文本提示 "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"

long_prompt = 10 * system_prompt + prompt

我们再次以 bfloat16 精度实例化我们的模型。

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

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

现在让我们像以前一样在没有 Flash Attention 的情况下运行模型，并测量 GPU 内存峰值需求和推理时间。

import time

start_time = time.time()
result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

输出:

Generated in 10.96854019165039 seconds.
Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

我们获得了与之前相同的输出，但是这次，模型多次重复答案，直到达到 60 个 token 的截断为止。这并不奇怪，因为我们为了演示目的将系统提示重复了十次，因此提示模型重复自身。

注意，在实际应用中，系统提示不应重复十次 - 一次就足够了！

让我们测量 GPU 内存峰值需求。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

37.668193340301514

正如我们所见，GPU 内存峰值需求现在明显高于开始时，这主要是由于输入序列更长。并且生成时间现在也超过了一分钟。

我们调用 flush() 以释放 GPU 内存以进行下一次实验。

flush()

为了进行比较，让我们运行相同的函数，但启用 Flash Attention。为此，我们将模型转换为 BetterTransformer，并通过这样做启用 PyTorch 的 SDPA 自注意力，而 SDPA 自注意力又能够使用 Flash Attention。

model.to_bettertransformer()

现在我们运行与之前完全相同的代码片段，而 Transformers 将在后台使用 Flash Attention。

start_time = time.time()
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

输出:

Generated in 3.0211617946624756 seconds.
 Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

我们得到了与之前完全相同的结果，但由于 Flash Attention，我们可以观察到速度显着提升。

让我们最后一次测量内存消耗。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

32.617331981658936

我们几乎回到了最初的 29GB GPU 内存峰值。

我们可以观察到，与最初传递短输入序列相比，在使用 Flash Attention 传递非常长的输入序列时，我们仅多使用了大约 100MB 的 GPU 内存。

flush()

有关如何使用 Flash Attention 的更多信息，请查看此文档页面。

3. 架构创新

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

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

现在让我们研究一下如何更改 LLM 的架构，使其对于需要长文本输入的任务（例如）最有效和高效：

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

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

一旦训练完成，基本的 LLM 架构就很难更改，因此预先考虑 LLM 的任务并相应地优化模型的架构非常重要。模型架构中有两个重要的组件，对于大型输入序列，它们很快就会成为内存和/或性能瓶颈。

位置嵌入
键值缓存

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

3.1 改进 LLM 的位置嵌入

自注意力将每个 token 与彼此的 token 相关联。例如， $Softmax(QK T)$ 文本输入序列 “Hello”, “I”, “love”, “you” 的矩阵可能如下所示

每个单词 token 都被赋予一个概率质量，它根据该概率质量关注所有其他单词 token，因此与所有其他单词 token 建立联系。例如，单词 “love” 以 5% 的概率关注单词 “Hello”，以 30% 的概率关注 “I”，并以 65% 的概率关注自身。

基于自注意力但没有位置嵌入的 LLM 在理解文本输入彼此之间的位置方面会遇到很大的困难。这是因为由以下公式计算的概率分数： $QK T$ 在以下计算中将每个单词 token 与每个其他单词 token 相关联 $O(1)$ 计算，而不管它们彼此之间的相对位置距离如何。因此，对于没有位置嵌入的 LLM，每个 token 看起来都与所有其他 token 具有相同的距离，例如，区分 “Hello I love you” 和 “You love I hello” 将非常具有挑战性。

为了让 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 中的主要方法，但是发现了一些与这些位置编码相关的问题

正弦和学习的位置嵌入都是绝对位置嵌入，即为每个位置 ID 编码一个唯一的嵌入 $0, \ldots, N$ . 正如 Huang 等人和 Su 等人表明的那样，绝对位置嵌入会导致 LLM 在长文本输入中表现不佳。对于长文本输入，如果模型学习输入 tokens 彼此之间的相对位置距离，而不是它们的绝对位置，则会更有利。
当使用学习的位置嵌入时，LLM 必须在固定的输入长度上进行训练 $N$ ，这使得难以推断到比训练长度更长的输入长度。

最近，可以解决上述问题的相对位置嵌入变得越来越流行，最值得注意的是

RoPE 和 ALiBi 都认为，最好在自注意力算法中直接提示 LLM 关于句子顺序的信息，因为正是在那里，单词 tokens 彼此之间建立了关系。更具体地说，应该通过修改以下内容来提示句子顺序 $QK T$ 计算。

在不深入太多细节的情况下，RoPE 指出，位置信息可以编码到查询-键对中，例如 $\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 提出了一种更简单的相对位置编码方案。输入 tokens 彼此之间的相对距离作为一个负整数添加，并按预定义的值 m 缩放到 $QK T$ 矩阵的每个查询-键条目，就在 softmax 计算之前。

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

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

RoPE 和 ALiBi 位置编码都可以外推到训练期间未见的输入长度，而已有研究表明，对于 ALiBi 而言，开箱即用的外推效果比 RoPE 好得多。对于 ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度即可。对于 RoPE，保持与训练期间使用的相同的 $\theta$ 会导致在传递比训练期间看到的文本输入长得多的文本输入时，效果不佳，参见 Press 等人。然而，社区已经找到了一些有效的技巧来调整 $\theta$ ，从而使 RoPE 位置嵌入能够很好地用于外推的文本输入序列（参见此处）。

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

关于文本输入的位置提示应该直接提供给 $QK T$ 自注意力层的矩阵

应该激励 LLM 学习恒定的相对距离位置编码彼此之间

文本输入 tokens 彼此之间的距离越远，它们的查询-值概率就越低。 RoPE 和 ALiBi 都降低了彼此远离的 tokens 的查询-键概率。 RoPE 通过增加查询-键向量之间的角度来减小它们的向量积。 ALiBi 通过向向量积添加大的负数

总之，旨在部署在需要处理大型文本输入的任务中的 LLM 最好使用相对位置嵌入进行训练，例如 RoPE 和 ALiBi。另请注意，即使使用 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("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 增加文本输入 tokens。

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

因此，tokens 永远不依赖于先前的 tokens，更具体地说， $\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 中，我们可以通过将 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", len(past_key_values[0][0]))  # 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 的长度不会增加，而是保持为单个输入向量。另一方面，键值缓存的长度在每个解码步骤中都会增加一。

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

使用键值缓存有两个优点

计算效率显著提高，因为与计算完整的相比，执行的计算更少 $QK T$ 矩阵。这导致推理速度的提高
所需的最大内存不会随着生成的 token 数量呈二次方增长，而只会线性增长。

应该 *始终* 使用键值缓存，因为它会带来相同的结果，并显着加快较长输入序列的速度。当使用文本 pipeline 或 generate 方法时，Transformers 默认启用键值缓存。我们有一个专门介绍缓存的完整指南在此。

请注意，尽管我们建议使用键值缓存，但当您使用它们时，您的 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 运行了两次自回归解码

第一次，键值缓存为空，输入提示是 "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?"。在处理缩短的输入提示时，其计算的键值向量被连接到第一次解码的键值缓存中。然后，使用由 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?" 的编码键值向量组成的键值缓存，自回归地生成第二个 Assistant 的答案 "Germany has ca. 81 million inhabitants" 。

这里应该注意两件事

保留所有上下文对于部署在聊天中的 LLM 至关重要，这样 LLM 才能理解对话的所有先前上下文。例如，对于上面的示例，LLM 需要理解用户在询问 "And how many are in Germany" 时指的是人口。
键值缓存对于聊天非常有用，因为它允许我们持续增长编码的聊天历史记录，而不是必须从头开始重新编码聊天历史记录（例如，在使用编码器-解码器架构时就是这种情况）。

在 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

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

让我们计算一下对于我们之前使用的 LLM bigcode/octocoder，需要在键值缓存中存储的 float 值数量。float 值的数量等于 2 乘以序列长度，再乘以 attention head 的数量，再乘以 attention head 的维度，再乘以层数。在假设输入序列长度为 16000 的情况下，为我们的 LLM 计算得出

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

输出:

7864320000

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

3.2.2 多查询注意力 (MQA)

多查询注意力是在 Noam Shazeer 的论文 Fast Transformer Decoding: One Write-Head is All You Need 中提出的。正如标题所说，Noam 发现，与其使用 n_head 个键值投影权重，不如使用在所有 attention head 之间共享的单个 head-value 投影权重对，而模型的性能不会显着下降。

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

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

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

这里要理解的重要部分是，只有在使用键值缓存的情况下，将键值 attention head 的数量减少到 1 才有意义。在没有键值缓存的情况下，模型单次前向传递的峰值内存消耗保持不变，因为每个 attention head 仍然具有唯一的查询向量，因此每个 attention head 仍然具有不同的 $QK T$ 矩阵。

MQA 已被社区广泛采用，现在被许多最流行的 LLM 使用

此外，本 notebook 中使用的 checkpoint - bigcode/octocoder - 也使用了 MQA。

3.2.3 分组查询注意力 (GQA)

分组查询注意力是由 Google 的 Ainslie 等人提出的，他们发现与使用原始多键值 head 投影相比，使用 MQA 通常会导致质量下降。该论文认为，通过不太大幅度地减少查询 head 投影权重的数量，可以保持更多的模型性能。应该使用 n < n_head 个键值投影权重，而不是仅使用单个键值投影权重。通过将 n 选择为比 n_head 小得多的值，例如 2、4 或 8，几乎可以保留 MQA 的所有内存和速度增益，同时牺牲更少的模型容量，因此可以说性能损失更小。

此外，GQA 的作者发现，可以微调现有模型 checkpoint，使其具有 GQA 架构，而只需使用原始预训练计算量的 5%。虽然 5% 的原始预训练计算量仍然可能很大，但 GQA 微调允许现有 checkpoint 可用于更长的输入序列。

GQA 是最近才提出的，这就是为什么在编写本 notebook 时采用较少的原因。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 等加速器只会变得更快并允许更多内存，但尽管如此，仍应始终确保使用最佳可用算法和架构，以获得最大的收益 🤗

< > 更新在 GitHub 上