优化LLM的速度和内存

像GPT3/4、Falcon 和 Llama 这样的大型语言模型 (LLM) 在处理以人类为中心的任务方面正在迅速发展，并已成为现代知识型产业的重要工具。然而，在实际任务中部署这些模型仍然具有挑战性。

• 为了展现接近人类的文本理解和生成能力，LLM 目前需要由数十亿个参数组成 (参见 Kaplan et al, Wei et. al)。这必然会增加推理的内存需求。
• 在许多实际任务中，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 的 VRAM。

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

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

对于较短的文本输入（小于 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，提供 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。让我们试一试。

我们首先加载模型和分词器，然后将两者传递给 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.）。甚至可以将模型量化到 3 位或 2 位，性能损失在可接受范围内，如最近的 GPTQ 论文 🤯 所示。

不深入过多细节，量化方案旨在降低权重的精度，同时尽量保持模型的推理结果尽可能准确（即尽可能接近 bfloat16）。请注意，量化对于文本生成尤其有效，因为我们只关心选择 *最有可能的下一个 token 集*，而不关心下一个 token *logit* 分布的确切值。重要的是下一个 token *logit* 分布大致保持不变，这样 `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，因此可以在像 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, 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 从 32GB 减少到仅 15GB，而以 4 位精度运行模型则将所需的 GPU VRAM 进一步减少到仅 9GB 多一点。

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

有关量化的更多信息以及如何量化模型以使所需的 GPU VRAM 少于 4 位，我们建议查阅 `AutoGPTQ` 实现。

总而言之，重要的是要记住模型量化以提高内存效率为代价，有时也会牺牲准确性和推理时间。

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

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

2. Flash Attention

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

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

让我们仔细看看。计算自注意力层输出$\mathbf{O}$的公式，对于长度为$\mathbf{X}$的输入是：$\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$我们将需要将近 1TB 来存储$\mathbf{QK}^T$矩阵。

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

随着 LLM 在文本理解和生成方面的改进，它们被应用于日益复杂的任务。模型曾经处理几句话的翻译或摘要，现在可以管理整个页面，这需要处理大量输入的能力。

我们如何摆脱对长输入长度的过高内存要求？我们需要一种新的计算自注意力机制的方法，它能消除$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。该算法在数学上给出相同的输出，并且速度更快，内存效率更高。

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

我们的 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 自注意力，而后者又能够使用 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的位置嵌入

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

每个单词标记都被赋予一个概率质量，表示它关注所有其他单词标记的程度，因此与其他所有单词标记关联起来。例如，单词“love”以 5% 的概率关注“Hello”，以 30% 的概率关注“I”，以 65% 的概率关注它自己。

一个基于自注意力但没有位置嵌入的 LLM 在理解文本输入彼此之间的位置时会遇到很大困难。这是因为由$\mathbf{QK}^T$将每个词元与$O(1)$计算中的所有其他词元关联起来，无论它们之间的相对位置距离如何。因此，对于没有位置嵌入的 LLM，每个词元似乎与其他所有词元具有相同的距离，例如，区分“Hello I love you”和“You love I hello”将非常具有挑战性。$O(1)$计算，无论它们彼此之间的相对位置距离如何。因此，对于没有位置嵌入的LLM，每个token似乎与其他所有token具有相同的距离，例如，区分“Hello I love you”和“You love I hello”将非常具有挑战性。

为了让 LLM 理解句子顺序，需要额外的“提示”，通常以“位置编码”（或“位置嵌入”）的形式应用。位置编码将每个标记的位置编码为 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）中的主要方法，但后来发现这些位置编码存在一些问题。

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

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

• 旋转位置嵌入（RoPE）
• ALiBi

RoPE和ALiBi都认为，最好直接在自注意力算法中向LLM提示句子顺序，因为词元正是在那里相互关联的。更具体地说，句子顺序应该通过修改$\mathbf{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中，例如：

• Falcon
• Llama
• PaLM

作为替代方案，ALiBi 提出了一种更简单的相对位置编码方案。输入词元之间的相对距离，以一个由预定义值 m 缩放的负整数形式，在 softmax 计算之前直接添加到 matrix 的每个查询-键条目中。$\mathbf{QK}^T$矩阵 right before the softmax computation.

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

ALiBi目前被用于许多最重要的LLM中，例如：

• MPT
• BLOOM

RoPE和ALiBi位置编码都可以外推到训练时未见的输入长度，尽管已表明ALiBi的开箱即用外推效果远优于RoPE。对于ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度即可。对于RoPE，如果保持与训练时相同的$\theta$当输入文本长度远超训练时所见长度时，这将导致糟糕的结果，参见Press 等人。然而，社区已经发现了一些有效的技巧来调整$\theta$，从而使得RoPE位置嵌入在推断的文本输入序列中也能很好地工作（参见此处）。

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

• 文本输入的位置提示应直接提供给$QK^T$自注意力层的矩阵
• 大型语言模型（LLM）应被鼓励学习位置编码彼此之间恒定的相对距离
• 文本输入词元距离越远，其查询-值概率越低。RoPE和ALiBi都会降低彼此相距较远的词元的查询-键概率。RoPE通过增加查询-键向量之间的角度来减小它们的向量积。ALiBi通过向向量积添加大的负数来实现。

总之，旨在处理大文本输入任务的LLM，最好使用RoPE和ALiBi等相对位置嵌入进行训练。此外请注意，即使一个带有RoPE和ALiBi的LLM仅以固定长度进行训练，例如$N_1 = 2048$，在实践中仍可用于处理远大于$N_1$的文本输入，例如$N_2 = 8192 > N_1$通过外推位置嵌入。

3.2 键值缓存

LLM 的自回归文本生成通过迭代地输入一个序列，采样下一个词元，将下一个词元附加到输入序列，并持续这样做，直到 LLM 产生一个表示生成已完成的词元。

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

我们来快速运行一段代码，以展示自回归在实践中是如何工作的。我们将简单地通过`torch.argmax`获取最有可能的下一个词元。

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

如我们所见，每次我们将刚采样的词元添加到文本输入词元中，都会增加其长度。

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

因此，词元从不依赖于先前的词元，更具体地说，$\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` 调用来检索键值缓存，然后将其与当前词元一起传递。

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

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

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

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

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

始终应该使用键值缓存，因为它可以带来相同的结果，并且对于较长的输入序列，显著加快速度。在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. 第一次，键值缓存为空，输入提示为“用户：法国有多少人口？”模型自回归生成文本“法国大约有7500万人口”，并在每个解码步骤增加键值缓存。
2. 第二次，输入提示是“用户：法国有多少人口？\n助理：法国大约有7500万人口。\n用户：德国有多少人口？”。由于缓存的存在，前两句的所有键值向量都已计算。因此，输入提示只包含“用户：德国有多少人口？”。在处理缩短的输入提示时，其计算出的键值向量会与第一次解码的键值缓存连接起来。然后，第二个助理的回答“德国大约有8100万居民”将使用包含“用户：法国有多少人口？\n助理：法国大约有7500万人口。\n用户：德国有多少人口？”编码的键值向量的键值缓存进行自回归生成。

此处应注意两点

1. 对于部署在聊天中的大型语言模型（LLM），保留所有上下文至关重要，以便LLM理解对话的所有先前上下文。例如，对于上述例子，当用户询问“德国有多少人口”时，LLM需要理解用户指的是人口。
2. 键值缓存对于聊天非常有用，因为它允许我们持续增长编码的聊天历史，而无需从头重新编码聊天历史（例如，在使用编码器-解码器架构时就会出现这种情况）。

在 `transformers` 中，当 `return_dict_in_generate=True` 被传入（除了默认的 `use_cache=True` 之外）时，`generate` 调用将返回 `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` 的键值缓存需要存储的浮点值数量。浮点值数量是序列长度的两倍，再乘以注意力头数量，再乘以注意力头维度，最后乘以层数。对于我们的 LLM，在假定输入序列长度为 16000 的情况下，计算结果为：

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$在所有注意力头中必须相同，这意味着我们只需在缓存中存储一对键值投影，而不是`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 分组查询注意力（GQA）

分组查询注意力，由 Google 的 Ainslie 等人提出，发现使用 MQA 通常会导致质量下降，而普通的多键值头投影则不然。该论文认为，通过不那么大幅度地减少查询头投影权重，可以保持更好的模型性能。它建议使用 `n < n_head` 个键值投影权重，而不是仅使用单个键值投影权重。通过将 `n` 选择为远小于 `n_head` 的值，例如 2、4 或 8，几乎可以保留 MQA 的所有内存和速度收益，同时牺牲更少的模型容量，从而可以说性能下降更少。

此外，GQA 的作者发现，现有模型检查点可以通过仅使用原始预训练计算量的 5% 进行升级训练，以具备 GQA 架构。尽管原始预训练计算量的 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 等加速器只会变得更快并允许更多的内存，但我们仍然应该始终确保使用最佳可用算法和架构，以获得最大的效益 🤗