在生产环境中优化LLM

注意：此博客文章也可在Transformers上作为文档页面获取。

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、Multi-Query Attention (MQA) 和 Grouped-Query-Attention (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，提供80GB显存。上述大多数模型仅加载就需要超过80GB，因此必然需要张量并行和/或流水线并行。

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

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

如果您可以使用8块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，因为它可以在单个40GB A100 GPU设备芯片上运行。请注意，我们接下来将应用的所有内存和速度优化同样适用于需要模型并行或张量并行的模型。

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

我们首先加载模型和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的显存。

现在几乎所有模型都以 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显存怎么办？研究发现，模型权重可以量化为8位或4位，而不会显著降低性能（参见Dettmers et al.）。如最近的GPTQ论文所示，模型甚至可以量化到3位或2位，且性能损失可接受🤯。

这里不做过多细节讨论，量化方案旨在降低权重的精度，同时尽量保持模型的推理结果尽可能准确（即尽可能接近bfloat16）。请注意，量化对于文本生成特别有效，因为我们只关心选择*最有可能的下一个token集合*，而不太关心下一个token*对数*分布的确切值。重要的是，下一个token*对数*分布大致保持不变，以便argmax或topk操作给出相同的结果。

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

• 1. 将所有权重量化到目标精度
• 2. 加载量化后的权重，并以bfloat16精度传递输入向量序列
• 3. 动态地将权重反量化为bfloat16，以便与bfloat16精度的输入向量进行计算
• 4. 在与输入计算后，再次将权重量化到目标精度。

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

$$Y = X * W$$

改变为

$$Y = X * \text{dequantize}(W); \text{quantize}(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

显著减少！我们现在只剩下略高于15GB，因此可以在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亿参数的模型来说，这确实不多。

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

另请注意，与8位量化相比，这里的推理速度再次稍慢，这是由于4位量化使用了更激进的量化方法，导致$\text{quantize}$和$\text{dequantize}$在推理过程中花费的时间更长。

del model
del pipe

flush()

总而言之，我们发现以8位精度运行OctoCoder将所需的GPU显存从32G GPU显存减少到仅15GB，而以4位精度运行模型则将所需的GPU显存进一步减少到略高于9GB。

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

有关量化的更多信息以及如何将模型量化到比4位更少的GPU显存，我们建议查阅`AutoGPTQ`实现。

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

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

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

2. Flash Attention：向前迈进的一大步

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

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

我们再仔细看看。自注意力层对于长度为$N$的输入$\mathbf{X}$计算输出$\mathbf{O}$的公式为

$$\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$，仅需要约50MB显存，但是对于$N=16000$，我们将需要19GB显存，而对于$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。为此，我们将模型转换为BetterTransformers，从而启用 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()

3. LLM架构背后的科学：长文本输入和聊天的策略选择

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

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

现在，让我们看看如何更改LLM的架构，使其最有效地处理需要长文本输入的任务，例如

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

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

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

• 位置嵌入
• 键值缓存

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

3.1 改进LLM的位置嵌入

自注意力将每个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理解句子顺序，需要额外的*线索*，通常以*位置编码*（或*位置嵌入*）的形式应用。位置编码将每个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 et al.）使用了学习型位置编码，其中位置嵌入$\mathbf{P}$在训练期间学习。

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

1. 正弦和学习型位置嵌入都是绝对位置嵌入，即为每个位置ID编码一个唯一的嵌入：$0, \ldots, N$。正如Huang et al.和Su et al.所展示的，绝对位置嵌入会导致LLM在处理长文本输入时性能不佳。对于长文本输入，如果模型学习了输入token之间相对位置距离而不是它们的绝对位置，将更有优势。
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 提出了一种更简单的相对位置编码方案。输入 token 之间的相对距离以负整数形式，乘以预定义值 m，添加到 $\mathbf{QK}^T$ 矩阵的每个查询-键条目中，紧接着 softmax 计算。

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

ALiBi 被用于当今最重要的多个大型语言模型（LLM）中，例如：

• MPT
• BLOOM

RoPE 和 ALiBi 位置编码都可以外推到训练期间未见的输入长度，尽管研究表明 ALiBi 的外推效果比 RoPE 更好。对于 ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度。对于 RoPE，如果使用训练期间相同的 $\theta$ 值，当输入文本长度远超训练时，会导致糟糕的结果，参见 Press et al.。然而，社区已经发现了一些有效的技巧来调整 $\theta$，从而使 RoPE 位置嵌入在外推文本输入序列中也能表现良好（参见此处）。

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

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

总之，旨在处理大型文本输入任务的 LLM 最好使用相对位置嵌入进行训练，例如 RoPE 和 ALiBi。还要注意，即使 LLM 使用 RoPE 和 ALiBi 并且仅在固定长度（例如 $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 都使用因果语言建模目标进行训练，因此会掩盖注意力分数矩阵的右上三角——这就是上面两个图中注意力分数留空（即概率为 0）的原因。有关因果语言建模的快速回顾，您可以参考图解自注意力博客。

因此，token 从不依赖未来的 token，更具体地说，如果 $j > i$，则 $\mathbf{q}_i$ 向量永远不会与任何键值向量 $\mathbf{k}_j, \mathbf{v}_j$ 建立关系。相反，$\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)
  # past_key_values are a tuple (one for each Transformer layer) of tuples (one for the keys, one for the values)
  # cached keys and values each are of shape (batch_size, num_heads, sequence_length, embed_size_per_head)
  # hence let's print how many cached keys and values we have for the first Transformer layer
  print("number of cached keys of the first Transformer layer", len(past_key_values[0][0][0,0,:,:]))
  print("number of cached values of the first Transformer layer", len(past_key_values[0][1][0,0,:,:]))
  
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

输出:

shape of input_ids torch.Size([1, 1])
number of cached keys of the first Transformer layer: 20
number of cached values of the first Transformer layer: 20
shape of input_ids torch.Size([1, 1])
number of cached keys of the first Transformer layer: 21
number of cached values of the first Transformer layer: 21
shape of input_ids torch.Size([1, 1])
number of cached keys of the first Transformer layer: 22
number of cached values of the first Transformer layer: 22
shape of input_ids torch.Size([1, 1])
number of cached keys of the first Transformer layer: 23
number of cached values of the first Transformer layer: 23
shape of input_ids torch.Size([1, 1])
number of cached keys of the first Transformer layer: 24
number of cached values of the first Transformer layer: 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 方法时，键值缓存默认启用。

需要注意的是，键值缓存对于需要处理大量文本输入的应用程序（例如聊天）尤其有用。让我们看一个例子。

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: 法国有多少人口？"，模型自回归生成文本 "大约有 7500 万人居住在法国"，同时在每个解码步骤增加键值缓存。
• 2. 第二次，输入提示是 "用户：法国有多少人口？\n助手：大约有 7500 万人居住在法国\n用户：那德国有多少人？"。得益于缓存，前两句话的所有键值向量都已计算。因此，输入提示只包含 "用户：那德国有多少人？"。在处理缩短的输入提示时，它计算出的键值向量会连接到第一次解码的键值缓存中。助手第二次的回答 "德国大约有 8100 万居民" 随后通过包含 "用户：法国有多少人口？\n助手：大约有 7500 万人居住在法国\n用户：那德国有多少人？" 编码的键值向量的键值缓存进行自回归生成。

这里有两点需要注意：

1. 对于部署在聊天中的 LLM 来说，保持所有上下文至关重要，这样 LLM 才能理解对话的所有先前上下文。例如，对于上面的例子，LLM 需要理解用户在询问 "那德国有多少人" 时指的是人口。
2. 键值缓存对于聊天非常有用，因为它允许我们持续增长编码的聊天历史，而不必从头开始重新编码聊天历史（例如，在使用编码器-解码器架构时就会出现这种情况）。

然而，这里有一个问题。虽然 $\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 亿个浮点值大约需要 15 GB 内存，这大约是模型权重本身的一半！研究人员提出了两种方法，可以显著降低存储键值缓存的内存成本：

1. 多查询注意力（MQA）

多查询注意力（Multi-Query-Attention, MQA）由 Noam Shazeer 在其论文《快速 Transformer 解码：一个写头足矣》（Fast Transformer Decoding: One Write-Head is All You Need）中提出。正如标题所示，Noam 发现，可以使用一组单一的键值投影权重，并在所有注意力头之间共享，而不会显著降低模型性能。

通过使用一个单一的头值投影权重对，键值向量 $\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{Q}\mathbf{K}}^T$ 矩阵。

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

• Falcon
• PaLM
• MPT
• BLOOM

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

2. 分组查询注意力（GQA）

分组查询注意力（Grouped-Query-Attention，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 的推理时间。例如，一个有前景的研究方向是推测解码，其中“容易的 token”由较小、较快的语言模型生成，只有“困难的 token”才由 LLM 本身生成。深入细节超出了本笔记本的范围，但可以在这篇精彩的博客文章中阅读更多内容。

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