Transformers 4.37 中文文档（十二）（5）

Transformers 4.37 中文文档（十二）（4）https://developer.aliyun.com/article/1564915

2. 闪光关注

今天表现最佳的 LLMs 基本上共享相同的基本架构，包括前馈层、激活层、层归一化层，以及最关键的自注意力层。

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

让我们仔细看看。计算自注意力层对于长度为N N N 的输入X \mathbf{X} X 的输出O \mathbf{O} O  的公式是：O=Attn(X)=V×Softmax(QKT) with Q=WqX,V=WvX,K=WkX \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} O=Attn(X)=V×Softmax(QKT) with  Q=WqX,V=WvX,K=WkX X=(x1,...xN) \mathbf{X} = (\mathbf{x}1, … \mathbf{x}{N}) X=(x1,…xN) 是注意力层的输入序列。投影Q \mathbf{Q} Q 和K \mathbf{K} K 将分别包含N N N 个向量，导致QKT \mathbf{QK}^T QKT 的大小为N2 N² N2 。

LLMs 通常具有多个注意力头，因此可以并行进行多个自注意力计算。假设 LLM 有 40 个注意力头并且以 bfloat16  精度运行，我们可以计算存储QKT \mathbf{QK^T} QKT 矩阵所需的内存为40∗2∗N2 40 * 2 * N² 40∗2∗N2  字节。对于N=1000 N=1000 N=1000，只需要大约 50MB 的 VRAM，然而，对于N=16000 N=16000  N=16000，我们将需要 19GB 的 VRAM，而对于N=100,000 N=100,000 N=100,000，我们将需要近 1TB 的  VRAM 来存储QKT \mathbf{QK}^T QKT 矩阵。

长话短说，对于大型输入上下文来说，默认的自注意力算法很快变得内存消耗过高。

随着 LLMs 在文本理解和生成方面的改进，它们被应用于越来越复杂的任务。虽然模型曾经处理几句话的翻译或总结，现在它们可以处理整页的内容，需要处理广泛的输入长度。

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

简而言之，Flash Attention 将V×Softmax(QKT\mathbf{V} \times  \text{Softmax}(\mathbf{QK}^TV×Softmax(QKT)计算分开，而是通过迭代多个 softmax  计算步骤来计算输出的较小块：Oi←sija∗Oi+sijb∗Vj×Softmax(QKi,jT) for multiple i,j  iterations \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} Oi←sija∗Oi+sijb∗Vj×Softmax(QKi,jT) for multiple i,j iterations

其中sija s^a_{ij}和sijb s^b_{ij}是一些需要为每个i i和j j重新计算的 softmax 归一化统计量。

请注意，整个 Flash Attention 有点复杂，在这里大大简化了，因为深入讨论超出了本指南的范围。读者可以查看写得很好的Flash Attention 论文以获取更多详细信息。

这里的主要要点是：

通过跟踪 softmax 归一化统计量，并使用一些智能数学，Flash Attention 给出了与默认自注意力层数值相同的输出，而内存成本仅随N N线性增加。

从公式来看，人们直觉地会说 Flash Attention  必须比默认的自注意力公式慢得多，因为需要进行更多的计算。事实上，与普通注意力相比，Flash Attention 需要更多的 FLOPs，因为  softmax 归一化统计量必须不断重新计算（如果感兴趣，请参阅论文获取更多详细信息）

然而，与默认注意力相比，Flash Attention 在推理速度上要快得多，这是因为它能够显著减少对 GPU（VRAM）较慢、高带宽内存的需求，而是专注于更快的片上内存（SRAM）。

基本上，Flash Attention 确保所有中间写入和读取操作都可以使用快速的片上SRAM 内存来完成，而无需访问较慢的 VRAM 内存来计算输出向量O \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 的内存节省。我们附加原始文本提示"问题：请用 Python 编写一个将字节转换为千兆字节的函数。\n\n 答案：在这里"

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 个标记的截止。这并不奇怪，因为我们为演示目的重复了系统提示十次，从而提示模型重复自己。

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

让我们测量 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 的架构，使其对需要长文本输入的任务最有效和高效，例如：

• 检索增强问答，
  
• 总结，
  
• 聊天
  

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

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

• 位置嵌入
  
• 键-值缓存
  

让我们更详细地讨论每个组件

3.1 改进 LLM 的位置嵌入

自注意力将每个标记与其他标记相关联。例如，文本输入序列的 Softmax(QKT)矩阵*“Hello”, “I”, “love”, “you”*可能如下所示：

每个单词标记都被赋予一个概率质量，用于关注所有其他单词标记，因此与所有其他单词标记相关联。例如，单词*“love”关注单词“Hello”的概率为 5%，关注“I”*的概率为 30%，自身的概率为 65%。

基于自注意力的 LLM，但没有位置嵌入，将在理解文本输入之间的位置方面遇到很大困难。这是因为由QKT \mathbf{QK}^T QKT  计算的概率分数将每个单词标记与其他单词标记在O(1) O(1) O(1) 计算中相关联，而不考虑它们之间的相对位置距离。因此，对于没有位置嵌入的  LLM，每个标记似乎与所有其他标记具有相同的距离，例如，区分“你好 我爱你”和“你爱我 你好”将会非常具有挑战性。

为了让 LLM 理解句子顺序，需要额外的提示，通常以位置编码（也称为位置嵌入）的形式应用。位置编码将每个标记的位置编码为 LLM 可以利用的数值表示，以更好地理解句子顺序。

注意力机制就是你所需要的论文的作者们引入了正弦位置嵌入P=p1,…,pN  \mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N P=p1,…,pN 。其中每个向量pi  \mathbf{p}_i pi 是根据其位置i i i 计算的正弦函数。然后将位置编码简单地添加到输入序列向量中X^=x^1,…,x^N  \mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N X^=x¹,…,xN =x1+p1,…,xN+pN \mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N x1+p1,…,xN+pN 从而提示模型更好地学习句子顺序。

其他人（例如Devlin 等人）使用了学习的位置编码，而不是固定的位置嵌入，这些位置嵌入在训练期间进行学习。

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

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

最近，能够解决上述问题的相对位置嵌入变得更加流行，其中最著名的是：

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

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

不详细讨论，RoPE指出位置信息可以被编码到查询-键对中，例如：qi \mathbf{q}_i 和xj  \mathbf{x}_j ，通过将每个向量旋转一个角度θ∗i \theta * i 和θ∗j \theta * j ，其中i,j i, j  描述每个向量的句子位置：q^iTx^j=qiTRθ,i−jxj. \mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j  = \mathbf{{q}}i^T \mathbf{R}{\theta, i -j} \mathbf{{x}}j. q^iTxj=qiTRθ,i−jxj。 Rθ,i−j \mathbf{R}{\theta, i - j} Rθ,i−j代表一个旋转矩阵。θ \theta θ在训练过程中不会被学习，而是设置为一个预定义的值，该值取决于训练过程中的最大输入序列长度。

通过这样做，qi \mathbf{q}_i 和qj \mathbf{q}_j 之间的概率分数只有在i≠j i \ne j 时才会受到影响，并且仅取决于相对距离i−j i - j ，而不考虑每个向量的具体位置i i 和j j 。

RoPE被用在当今一些最重要的 LLM 中，例如：

• 猎鹰
  
• 大羊驼
  
• PaLM
  

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

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

ALiBi 在当今一些最重要的 LLM 中使用，例如：

• MPT
  
• BLOOM
  

RoPE 和 ALiBi 位置编码都可以外推到训练中未见过的输入长度，然而已经证明相对于 RoPE，外推对于 ALiBi 来说更容易。对于 ALiBi，只需增加下三角位置矩阵的值以匹配输入序列的长度。对于 RoPE，保持训练期间使用的相同θ \theta θ，在传递比训练期间看到的文本输入长得多的文本输入时会导致结果不佳，参见 Press 等人。然而，社区已经发现了一些有效的技巧，可以调整θ \theta θ，从而使 RoPE 位置嵌入适用于外推的文本输入序列（参见这里）。

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

• 关于文本输入的位置提示应直接提供给自注意力层的QKT QK^T QKT 矩阵
• 
  
• LLM 应该被激励学习常数 相对 距离位置编码之间的关系
• 
  
• 文本输入令牌之间的距离越远，它们的查询-值概率就越低。RoPE 和 ALiBi 都降低了远离彼此的令牌的查询-键概率。RoPE 通过增加查询-键向量之间的角度来减少它们的向量积。ALiBi 通过向向量积添加大的负数

总之，用于处理大文本输入的任务的 LLM 最好使用相对位置嵌入进行训练，例如 RoPE 和 ALiBi。还要注意，即使一个带有 RoPE 和  ALiBi 的 LLM 只在固定长度的数据上进行了训练，比如N1=2048 N_1 = 2048  N1=2048，它仍然可以在实践中用于比N1 N_1 N1更大的文本输入，比如N2=8192>N1 N_2 = 8192 >  N_1 N2=8192>N1，通过外推位置嵌入。

3.2 关键-值缓存

LLMs 的自回归文本生成通过迭代地输入一个序列，抽样下一个标记，将下一个标记附加到输入序列中，并继续这样做，直到 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']

正如我们所看到的，每次我们通过刚刚抽样的标记增加文本输入标记。

除了极少数例外，LLMs 是使用因果语言建模目标进行训练的，因此会屏蔽注意力分数的上三角矩阵 - 这就是为什么在上述两个图表中，注意力分数留空（也就是概率为 0）。关于因果语言建模的快速回顾，您可以参考Illustrated Self Attention blog。

因此，标记永远不依赖于先前的标记，更具体地说，qi \mathbf{q}_i qi向量永远不会与任何键、值向量kj,vj \mathbf{k}_j, \mathbf{v}j kj,vj相关联，如果j>i j > i j>i。相反，qi \mathbf{q}i qi只关注先前的键-值向量km

接下来，我们将告诉 LLM 利用键-值缓存，通过在每次前向传递中检索并转发它。在 Transformers 中，我们可以通过向forward调用传递use_cache标志来检索键-值缓存，然后可以将其与当前标记一起传递。

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

正如大家所看到的，当使用键-值缓存时，文本输入标记的长度不会增加，而是保持为单个输入向量。另一方面，键-值缓存的长度在每个解码步骤都会增加一个。

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

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

• 与计算完整的 QKT 矩阵相比，计算效率显著提高，因为进行的计算较少。这导致推理速度增加。
  
• 所需的最大内存并不是随着生成的令牌数量的平方增加，而是线性增加。
  

应该始终利用键值缓存，因为它会产生相同的结果，并且对于较长的输入序列会显著加快速度。当使用文本管道或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 会自回归解码两次：

1. 第一次，键值缓存为空，输入提示是"用户：法国有多少人口？"，模型会自回归生成文本`“法国大约有 7500 万人口”，同时在每个解码步骤中增加键值缓存。
   
2. 第二次输入提示是"用户：法国有多少人口？\n 助手：法国大约有 7500 万人口\n 用户：德国有多少人口？"。由于缓存的存在，前两个句子的所有键值向量已经计算完毕。因此，输入提示只包括"用户：德国有多少人口？"。在处理缩短的输入提示时，它的计算键值向量会与第一次解码的键值缓存连接起来。然后第二个助手的回答"德国大约有 8100 万居民"会根据编码的键值向量"用户：法国有多少人口？\n 助手：法国大约有 7500 万人口\n 用户：德国有多少人口？"进行自回归生成。
   

这里有两点需要注意：

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

在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

太好了，不需要额外的时间来重新计算注意力层的相同键和值！然而，有一个问题。虽然QKT \mathbf{QK}^T QKT  矩阵所需的峰值内存显著减少，但在内存中保持键值缓存可能会对长输入序列或多轮对话非常昂贵。请记住，键值缓存需要存储所有先前输入向量的键值向量xi,  for i∈{1,…,c−1} \mathbf{x}_i \text{, for } i \in {1, \ldots, c - 1}  xi, for i∈{1,…,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 亿个浮点值需要大约 15 GB 的内存，这大约是模型权重本身的一半！研究人员提出了两种方法，可以显著减少存储键值缓存的内存成本，这将在接下来的小节中探讨。

3.2.2 多查询注意力（MQA）

Multi-Query-Attention 是 Noam Shazeer 在Fast Transformer Decoding: One Write-Head is All You Need论文中提出的。正如标题所说，Noam 发现，可以使用一个单一的头值投影权重对，而不是使用n_head个键值投影权重，这个对在所有注意力头部之间共享，而不会显著降低模型的性能。

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

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

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

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

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 的作者发现，现有的模型检查点可以通过仅使用原始预训练计算量的 5%进行更新训练，以实现 GQA 架构。虽然原始预训练计算量的 5%仍然是一个巨大的数量，但 GQA 的更新训练使现有的检查点可以用于更长的输入序列。

GQA 是最近提出的，因此在撰写本笔记时采用的情况较少。GQA 最显著的应用是Llama-v2。

总之，强烈建议在 LLM 部署自回归解码并需要处理大型输入序列的情况下使用 GQA 或 MQA。

结论

研究界不断提出新的巧妙方法来加快越来越大的 LLM 的推理时间。例如，一个有前途的研究方向是推测解码，其中较小、更快的语言模型生成“简单标记”，只有 LLM 本身生成“困难标记”。更详细的内容超出了本笔记的范围，但可以在这篇不错的博客文章中阅读。

大型 LLM（如 GPT3/4、Llama-2-70b、Claude、PaLM）能够在Hugging Face Chat或 ChatGPT 等聊天界面中运行如此迅速，这在很大程度上要归功于上述精度、算法和架构的改进。未来，像 GPU、TPU 等加速器将会变得更快，允许更多的内存，但仍然应始终确保使用最佳的可用算法和架构，以获得最大的性价比🤗

使用的 LLM，因此我们可以将所需的内存消耗从 15 GB 减少到输入序列长度为 16000 时的不到 400 MB。

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

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

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 的作者发现，现有的模型检查点可以通过仅使用原始预训练计算量的 5%进行更新训练，以实现 GQA 架构。虽然原始预训练计算量的 5%仍然是一个巨大的数量，但 GQA 的更新训练使现有的检查点可以用于更长的输入序列。

GQA 是最近提出的，因此在撰写本笔记时采用的情况较少。GQA 最显著的应用是Llama-v2。

总之，强烈建议在 LLM 部署自回归解码并需要处理大型输入序列的情况下使用 GQA 或 MQA。

结论

研究界不断提出新的巧妙方法来加快越来越大的 LLM 的推理时间。例如，一个有前途的研究方向是推测解码，其中较小、更快的语言模型生成“简单标记”，只有 LLM 本身生成“困难标记”。更详细的内容超出了本笔记的范围，但可以在这篇不错的博客文章中阅读。

大型 LLM（如 GPT3/4、Llama-2-70b、Claude、PaLM）能够在Hugging Face Chat或 ChatGPT 等聊天界面中运行如此迅速，这在很大程度上要归功于上述精度、算法和架构的改进。未来，像 GPU、TPU 等加速器将会变得更快，允许更多的内存，但仍然应始终确保使用最佳的可用算法和架构，以获得最大的性价比🤗