优化大型语言模型的速度和内存

大型语言模型（LLM），例如 GPT3/4、Falcon 和 Llama，在处理以人为中心的任务的能力方面正在迅速发展，已成为现代知识型行业中必不可少的工具。然而，在现实世界中部署这些模型仍然具有挑战性。

• 为了展现接近人类的文本理解和生成能力，当前的大型语言模型需要由数十亿个参数组成（参见 Kaplan 等人，Wei 等人）。这相应地增加了推理的内存需求。
• 在许多现实世界的任务中，大型语言模型需要获得广泛的上下文信息。这需要模型能够在推理过程中管理非常长的输入序列。

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

在本指南中，我们将介绍一些有效的大型语言模型高效部署技术。

1. 降低精度：研究表明，以降低的数值精度（即 8 位和 4 位）运行可以实现计算优势，而不会导致模型性能显著下降。

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

3. 架构创新：考虑到大型语言模型在推理过程中始终以相同的方式部署，即使用长输入上下文进行自回归文本生成，因此有人提出了专门的模型架构，可以实现更有效的推理。此处模型架构方面最重要的进步是 Alibi、旋转嵌入、多查询注意力 (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 个 token），推理的内存需求在很大程度上受加载权重的内存需求支配。因此，目前，让我们假设推理的内存需求等于将模型加载到 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，它们提供 80 GB 的 VRAM。上面列出的大多数模型仅加载就需要超过 80 GB，因此必须使用 张量并行 和/或 流水线并行。

🤗 Transformers 不开箱即用地支持张量并行，因为它要求模型架构以特定的方式编写。如果您有兴趣以支持张量并行的方式编写模型，可以参考 文本生成推理库。

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

如果您有 8 个 80 GB 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。让我们试一试。

我们首先加载模型和分词器，然后将两者传递给 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() 的实用程序方法。

from accelerate.utils import release_memory
# ...

release_memory(model)

现在，如果您的 GPU 没有 32 GB 的 VRAM 会怎样？研究发现，模型权重可以量化为 8 位或 4 位，而性能不会显著下降（参见 Dettmers 等人）。如最近的 GPTQ 论文 中所示，模型甚至可以量化为 3 位或 2 位，性能损失是可以接受的 🤯。

无需过多细节，量化方案旨在降低权重的精度，同时尽量保持模型的推理结果尽可能准确（即尽可能接近 bfloat16）。请注意，量化对文本生成特别有效，因为我们只关心选择“最有可能的下一个 token 集”，并不真正关心下一个 token 对数几率分布的精确值。重要的是，下一个 token 对数几率分布保持大致相同，以便 argmax 或 topk 操作得到相同的结果。

存在各种量化技术，我们在此不做详细讨论，但总的来说，所有量化技术的工作原理如下：

• 1. 将所有权重量化为目标精度
• 2. 加载量化后的权重，并以 bfloat16 精度传递输入向量序列
• 3. 动态地将权重量化为 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，因此可以在消费级 GPU（例如 4090）上运行此模型。我们看到了内存效率的显著提升，并且模型输出几乎没有下降。但是，我们也注意到推理过程中略有减慢。

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

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 亿个参数的模型来说，这真的不算多。

虽然我们在这里看到模型的准确性下降非常小，但在实践中，4 位量化通常会导致与 8 位量化或完整的 bfloat16 推理不同的结果。用户可以自行尝试。

还要注意，与 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) 至关重要，因为它使模型能够理解输入标记之间的上下文关系。但是，自注意力层的峰值 GPU 内存消耗随着输入标记数量（也称为序列长度）的增加而呈二次方增长，在下面我们将此数量表示为$N$。虽然对于较短的输入序列（最多 1000 个输入标记）来说，这并不明显，但对于较长的输入序列（大约 16000 个输入标记）来说，它就会成为一个严重的问题。

让我们仔细看看。计算输出的公式$\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}$

大型语言模型 (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$矩阵。

简而言之，对于大型输入上下文，默认的自注意力算法很快就会变得非常占用内存。

随着大型语言模型在文本理解和生成方面的改进，它们被应用于越来越复杂的任务。虽然模型曾经处理几个句子的翻译或总结，但现在它们可以管理整个页面，需要能够处理更长的输入长度。

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

重新计算。

请注意，整个Flash Attention稍微复杂一些，这里进行了大幅简化，因为深入探讨超出了本指南的范围。我们邀请读者查看写得很好的Flash Attention论文以获取更多详细信息。

这里的要点是$N$ .

通过跟踪softmax归一化统计数据并使用一些巧妙的数学方法，Flash Attention与默认的自注意力层相比，给出了数值上完全相同的输出，并且内存成本仅随

线性增长。从公式来看，人们会直观地认为，与默认的自注意力公式相比，Flash Attention必须慢得多，因为需要进行更多的计算。实际上，与普通注意力相比，Flash Attention确实需要更多的FLOPs，因为softmax归一化统计数据必须不断重新计算（如有兴趣，请参阅论文了解更多详细信息）。

然而，与默认注意力相比，Flash Attention在推理方面快得多，这得益于它能够显著降低对GPU（VRAM）较慢、高带宽内存的需求，转而专注于更快的片上内存（SRAM）。$\mathbf{O}$ .

本质上，Flash Attention确保所有中间写入和读取操作都可以使用快速的片上SRAM内存来完成，而不是必须访问较慢的VRAM内存来计算输出向量。

在实践中，如果可以使用Flash Attention，目前绝对没有理由不使用它。该算法在数学上给出相同的输出，并且更快且更节省内存。

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

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

-----
"""

我们的OctoCoder模型现在接收一个明显更长的输入提示，其中包含一个所谓的系统提示。系统提示用于引导LLM成为一个更好的助手，该助手针对用户的任务量身定制。在下文中，我们使用一个系统提示，使OctoCoder成为一个更好的代码助手。

long_prompt = 10 * system_prompt + prompt

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

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)

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

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

现在让我们像以前一样运行模型，但不使用Flash Attention，并测量GPU内存峰值需求和推理时间。

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

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

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

37.668193340301514

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

我们可以看到，GPU内存峰值需求现在比一开始高得多，这主要是由于更长的输入序列造成的。此外，生成现在需要一分钟多一点的时间。

flush()

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

model.to_bettertransformer()

为了进行比较，让我们运行相同的函数，但改为启用Flash Attention。为此，我们将模型转换为BetterTransformer，并通过这样做启用PyTorch的SDPA自注意力，后者能够使用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

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

我们得到了与之前完全相同的结果，但由于Flash Attention，我们可以观察到非常显著的加速。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出:

32.617331981658936

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

我们几乎回到了最初的29GB GPU内存峰值，就像一开始一样。

flush()

我们可以观察到，与一开始传递短输入序列相比，当使用Flash Attention传递非常长的输入序列时，我们只使用了大约100MB的GPU内存。

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

3. 架构创新

• 到目前为止，我们已经研究了如何通过以下方法提高计算和内存效率：
• 将权重转换为更低精度的格式

用更节省内存和计算的版本替换自注意力算法

• 现在让我们看看如何改变LLM的架构，使其最有效率地执行需要长文本输入的任务，例如：
• 检索增强型问答
• 摘要

聊天

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

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

键值缓存

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

3.1 改善LLM的位置嵌入自注意力将每个token与其他token联系起来。例如，$\text{Softmax}(\mathbf{QK}^T)$

文本输入序列“Hello”, “I”, “love”, “you”的矩阵可能如下所示

每个单词token都被赋予一个概率质量，它关注所有其他单词token，因此与所有其他单词token建立联系。例如，单词“love”以5%的概率关注单词“Hello”，以30%的概率关注“I”，并以65%的概率关注自身。$\mathbf{QK}^T$基于自注意力但没有位置嵌入的LLM在理解文本输入彼此的位置方面会遇到很大困难。这是因为计算的概率分数将每个单词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}_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}$在训练期间学习。

正弦和学习到的位置嵌入曾经是将句子顺序编码到大型语言模型中的主要方法，但发现了一些与这些位置编码相关的问题

1. 正弦和学习到的位置嵌入都是绝对位置嵌入，即为每个位置 ID 编码唯一的嵌入$0, \ldots, N$。正如 Huang 等人 和 Su 等人 所示，绝对位置嵌入会导致大型语言模型在长文本输入上的性能下降。对于长文本输入，如果模型学习输入标记彼此之间的相对位置距离而不是其绝对位置，则更有利。
2. 当使用学习到的位置嵌入时，大型语言模型必须在固定输入长度上进行训练$N$，这使得很难外推到比训练时更长的输入长度。

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

• 旋转位置嵌入 (RoPE)
• ALiBi

两者 RoPE 和 ALiBi 都认为最好直接在自注意力算法中提示大型语言模型关于句子顺序，因为在那里单词标记彼此关联。更具体地说，句子顺序应通过修改$\mathbf{QK}^T$计算。

在不深入过多细节的情况下，RoPE 指出位置信息可以编码到查询-键对中，例如$\mathbf{q}_i$和$\mathbf{x}_j$通过将每个向量旋转一个角度$\theta * i$和$\theta * j$分别与$i, j$描述每个向量的句子位置

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

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

• Falcon
• Llama
• PaLM

作为替代方案，ALiBi 提出了一个更简单的相对位置编码方案。输入标记彼此之间的相对距离被添加为一个负整数，该负整数乘以一个预定义的值m，添加到矩阵的每个查询-键条目中，$\mathbf{QK}^T$在softmax计算之前。

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

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

• MPT
• BLOOM

无论是RoPE 还是ALiBi 位置编码都可以外推到训练期间未见过的输入长度，而据表明，与RoPE 相比，ALiBi 的外推效果开箱即用地更好。对于ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度即可。对于RoPE，保持相同$\theta$在训练期间使用会导致在传递比训练期间看到的文本输入长得多的文本输入时产生较差的结果，c.f Press 等人。但是，社区已经找到了一些有效的技巧来适应$\theta$，从而允许RoPE 位置嵌入在扩展的文本输入序列中也能很好地工作（参见此处）。

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

• 关于文本输入的位置线索应直接提供给$QK^T$自注意力层的矩阵
• LLM 应该被激励去学习一个恒定的相对距离位置编码彼此之间的关系
• 文本输入标记彼此越远，其查询-值概率越低。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来增加文本输入token的数量。

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

因此，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中，我们可以通过将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的长度不会增加，而保持为单个输入向量。另一方面，键值缓存的长度在每个解码步骤中都会增加一个。

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

使用键值缓存有两个优点

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

应该始终使用键值缓存，因为它会导致相同的结果，并且对于较长的输入序列可以显著加快速度。Transformers在使用文本管道或generate方法时默认启用键值缓存。我们有一个专门介绍缓存的完整指南此处。

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调用除了默认的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的键值缓存需要存储的浮点数数量。浮点数的数量等于序列长度的两倍乘以注意力头的数量乘以注意力头维度乘以层的数量。对于假设输入序列长度为16000的LLM计算得到的结果是：

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

输出:

7864320000

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

3.2.2 多查询注意力 (MQA)

多查询注意力 由 Noam Shazeer 在其论文《快速Transformer解码：一个写头就够了》中提出。正如标题所说，Noam发现，与其使用n_head个键值投影权重，不如使用一个跨所有注意力头共享的键值投影权重对，而不会显著降低模型性能。

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

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

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

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

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

• Falcon
• PaLM
• MPT
• BLOOM

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

3.2.3 分组查询注意力 (GQA)

分组查询注意力 由来自谷歌的Ainslie等人提出，他们发现与使用普通的多个键值头投影相比，使用MQA通常会导致质量下降。该论文认为，通过减少查询头投影权重的数量，可以保留更多的模型性能。与其只使用一个键值投影权重，不如使用n < n_head个键值投影权重。通过将n选择为一个显著小于n_head的值，例如2、4或8，可以保留MQA几乎所有内存和速度方面的优势，同时牺牲更少的模型容量，从而可能牺牲更少的性能。

此外，GQA 的作者发现，现有的模型检查点可以通过微调来获得 GQA 架构，所需的预训练计算量仅为原始预训练计算量的 5%。虽然 5% 的原始预训练计算量仍然可能是一个巨大的数字，但 GQA 的微调使得现有的检查点能够用于更长的输入序列。

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

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

结论

研究界不断提出新的、巧妙的方法来加快越来越大的 LLM 的推理时间。例如，一个很有前景的研究方向是推测解码，其中“简单标记”由更小、更快的语言模型生成，只有“困难标记”由 LLM 本身生成。详细介绍不在本笔记本的范围之内，但可以在此优秀的博文中阅读。

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