优化您的LLM生产

优化LLM生产

注意：这篇博客文章也可以作为Transformer的文档页面使用。

大型语言模型（LLMs），如GPT3/4、Falcon和LLama，在处理人类中心任务方面迅速取得进展，已经成为现代知识型产业中不可或缺的工具。然而，在实际任务中部署这些模型仍然具有挑战性：

• 为了展示接近人类的文本理解和生成能力，LLMs目前需要由数十亿参数组成（参见Kaplan等人，Wei等人）。这相应增加了推理的内存需求。
• 在许多实际任务中，LLMs需要提供广泛的上下文信息。这要求模型在推理过程中能够处理非常长的输入序列。

这些挑战的关键在于增强LLMs的计算和内存能力，特别是在处理广泛的输入序列时。

• LLM的优势：改变电子商务搜索
• 使用TensorFlow训练变分自编码器进行异常检测
• PyCharm vs. Spyder 选择合适的Python集成开发环境

在这篇博客文章中，我们将介绍目前最有效的技术来解决这些挑战，以实现高效的LLM部署：

1. 降低精度：研究表明，使用降低的数值精度，即8位和4位，可以在不显著降低模型性能的情况下获得计算优势。

2. 闪存注意力：闪存注意力是注意力算法的一种变体，不仅提供了更高效的内存使用方式，还通过优化GPU内存利用率实现了更高的效率。

3. 架构创新：考虑到LLMs在推理过程中始终以相同的方式部署，即使用长输入上下文进行自回归文本生成，已经提出了专门的模型架构，以实现更高效的推理。其中，最重要的模型架构进展包括Alibi、Rotary嵌入、多查询注意力（MQA）和分组查询注意力（GQA）。

在本笔记本中，我们将从张量的角度对自回归生成进行分析。我们深入探讨采用降低精度的利弊，全面探索最新的注意力算法，并讨论改进的LLM架构。在此过程中，我们运行实际示例展示每个功能改进。

1. 发挥降低精度的威力

通过将LLMs视为一组权重矩阵和向量，将文本输入视为一系列向量，可以更好地理解LLMs的内存需求。在接下来的内容中，权重的定义将用于表示所有模型权重矩阵和向量。

在撰写本文时，LLMs至少由数十亿个参数组成。每个参数由一个十进制数，例如4.5689组成，通常以float32、bfloat16或float16格式存储。这使我们可以轻松计算将LLM加载到内存所需的内存量：

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

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

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

对于较短的文本输入（少于1024个标记），推理的内存需求在很大程度上受到加载权重的内存需求的影响。因此，暂时假设推理的内存需求等于将模型加载到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，提供80GB的VRAM。大多数在之前列出的模型需要超过80GB的内存才能加载，因此必须使用张量并行和/或流水线并行。

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

原始的流水线并行可以直接使用。只需使用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()的实用方法

from accelerate.utils import release_memory
# ...

release_memory(model)

那么如果您的GPU没有32GB的VRAM怎么办？已经发现模型权重可以量化为8位或4位而不会显著降低性能（参见Dettmers等人）。模型可以量化为3位或2位，性能损失可接受，如最近的GPTQ论文所示🤯。

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

有各种各样的量化技术，我们不会在这里详细讨论，但通常，所有量化技术的工作原理如下：

• 1. 将所有权重量化为目标精度
• 1. 加载量化后的权重，并使用bfloat16精度传递输入向量序列
• 1. 动态将权重反量化为bfloat16，以在其输入向量的bfloat16精度下执行计算
• 1. 在与其输入进行计算后，再次将权重量化为目标精度。

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

Y=X∗W Y = X * W Y=X∗W

被改为

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

输出：

这是一个将字节转换为千兆字节的 Python 函数：\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\n该函数接受一个参数

我们几乎看到了与之前相同的输出文本 – 只是在代码片段之前缺少了python。让我们看看需要多少内存。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出：

9.543574333190918

只有9.5GB！对于一个超过150亿参数的模型来说，这真的不算多。

虽然我们在这里几乎没有看到模型准确性的明显降低，但在实际中，4位量化通常会导致与8位量化或完整的bfloat16推断不同的结果。这要由用户自己尝试。

还要注意，在这里推断速度比8位量化再次慢一些，这是由于4位量化采用了更激进的量化方法，在推断过程中量化和反量化所需时间更长。

del model
del pipe

flush()

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

4位量化使模型可以在RTX3090、V100和T4等GPU上运行，这对大多数人来说相当容易获取。

有关量化的更多信息以及如何将模型量化以需要更少的GPU VRAM内存的详细信息，我们建议查阅AutoGPTQ的实现。

总之，重要的是要记住，模型量化会以提高内存效率为代价，牺牲一定的准确性，并在某些情况下会导致推断时间延长。

如果GPU内存对于您的用例不是一个限制，通常无需考虑量化。然而，许多GPU无法在不量化的情况下运行LLM，因此，4位和8位量化方案是非常有用的工具。

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

今天表现最好的LLM几乎共享相同的基本架构，包括前馈层、激活层、层归一化层和最重要的自注意层。

自注意层对于大型语言模型（LLM）至关重要，因为它们使模型能够理解输入令牌之间的上下文关系。然而，自注意层的峰值GPU内存消耗与输入令牌数量（也称为序列长度）呈二次增长，我们在下文中用 N N N 表示。虽然对于较短的输入序列（最多1000个输入令牌）这几乎不可察觉，但对于较长的输入序列（大约16000个输入令牌）而言，这是一个严重的问题。

让我们仔细看一下。计算长度为 N N N 的输入 X \mathbf{X} X 的自注意力层输出 O \mathbf{O} O 的公式为：

O=Attn(X)=V×Softmax(QKT) 其中 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=Wq​X,V=Wv​X,K=Wk​X mathbfX=(x1,…xN) mathbf{X} = (\mathbf{x}_1, … \mathbf{x}_{N}) mathbfX=(x1​,…xN​) 是注意力层的输入序列。投影 Q \mathbf{Q} Q 和 K \mathbf{K} K 每个由 N N N 个向量组成，因此 QKT \mathbf{QK}^T QKT 的大小为 N2 N^2 N2 。

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

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

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

我们如何摆脱大型输入长度的过高内存需求？我们需要一种新的计算自注意力机制的方式，摒弃 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) 对于多个 i,j 迭代 \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​) 对于多个 i,j 迭代

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

请注意，整个 Flash Attention 稍微复杂，这里只是大致说明，深入讨论超出了本笔记本的范围。读者可以查看撰写精良的 Flash Attention 论文以获取更多细节。

这里的主要要点是：

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

从公式上看，人们直观地会认为Flash Attention的速度要比默认的自注意力公式慢得多，因为需要进行更多的计算。事实上，与普通的注意力相比，Flash Attention需要更多的浮点运算操作，因为softmax归一化统计量必须不断重新计算（如果感兴趣，请参阅论文中的更多细节）。

然而，在推理中，与默认的注意力相比，Flash Attention要快得多，这是因为它能够大大降低对GPU（VRAM）上较慢的高带宽内存的需求，而是更侧重于更快的片上内存（SRAM）。

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

实际上，目前没有任何理由不使用Flash Attention（如果有的话）。该算法在数学上给出相同的输出，而且速度更快，内存使用效率更高。

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

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

system_prompt = """以下是各种人与AI技术助手之间的对话。
助手试图提供帮助、礼貌、诚实、精致、情感意识和谦卑但知识渊博。
助手乐于协助解答代码问题，并尽力理解确切的需求。
它还试图避免提供虚假或误导性的信息，并在不确定正确答案时加以限定。
也就是说，助手在实践中真正做到了最好，并且不会让谨慎过多地妨碍其有用性。

Starcoder模型是一系列基于80多种编程语言的15.5B参数模型，来自The Stack（v1.2）（不包括退出请求）。
该模型使用多查询注意力，使用填充中间目标进行训练，并使用8192个标记的上下文窗口进行逐字节重复数据的训练。

-----

问题：编写一个函数，该函数接受两个列表，并返回一个交替包含来自每个输入列表的元素的列表。

答案：当然。这是一个能够实现该功能的函数。

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

问题：你能为这个函数编写一些测试用例吗？

答案：当然，这是一些测试用例。

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([], []) == []

问题：修改该函数，使其在列表长度不相等时返回所有的输入元素。较长列表的元素应位于末尾。

答案：这是修改后的函数。

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"生成时间：{time.time() - start_time} 秒。")
result

输出：

生成时间为10.96854019165039秒。
可以的。这是一个可以实现此功能的函数。\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\n答案：可以的。这是一个可以实现此功能的函数。\n\ndef

我们得到了与之前相同的输出，但是这次模型将答案重复多次直到达到60个标记的限制。这并不奇怪，因为我们为了演示目的重复了系统提示十次，从而提示模型重复自己。

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

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

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出：

37.668193340301514

我们可以看到，峰值GPU内存需求现在比一开始显著增加，这主要是由于输入序列变长。此外，生成时间现在稍微超过一分钟。

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

flush()

为了比较，让我们运行相同的函数，但启用Flash Attention。为此，我们将模型转换为BetterTransformers，并通过此方式启用了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"生成时间为{time.time() - start_time}秒。")
result

输出：

生成时间为3.0211617946624756秒。
可以的。这是一个可以实现此功能的函数。\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\n答案：可以的。这是一个可以实现此功能的函数。\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的位置嵌入

自注意力（Self-attention）将每个标记与其他标记关联起来。例如，文本输入序列“Hello”，“I”，“love”，“you”的Softmax(QKT)矩阵可以如下所示：

每个词标记被赋予一个概率质量，表示它与所有其他词标记的关联程度。例如，词语“love”与“Hello”的关联度为0.05%，与“I”的关联度为0.3%，与自身的关联度为0.65%。

一个基于自注意力的LLM，如果没有位置嵌入，将很难理解文本输入之间的位置关系。这是因为由QKT计算出的概率分数将每个词标记与其他词标记关联起来，而不考虑它们之间的相对位置距离。因此，对于没有位置嵌入的LLM来说，每个标记看起来与所有其他标记的距离相同，例如，“Hello I love you”和“You love I hello”之间的区别将非常困难。

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

Attention Is All You Need论文的作者引入了正弦波位置嵌入P=p1,…,pN。其中每个向量pi是其位置i的正弦函数计算得到的。然后，位置编码简单地添加到输入序列向量X^=x^1,…,x^N中，即x1+p1,…,xN+pN，从而提示模型更好地学习句子顺序。

与使用固定位置嵌入不同，其他人（如Devlin等人）使用了学习到的位置编码，即在训练过程中学习位置嵌入P。

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

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

最近，相对位置嵌入变得越来越受欢迎，最值得注意的是：

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

RoPE和ALiBi都认为，最好直接在自注意力算法中向LLM提示句子顺序，因为正是在那里将单词标记与彼此关联起来。更具体地说，应通过修改QKT计算来提示句子顺序。

不详述细节，RoPE指出，位置信息可以编码到查询-键对中，例如通过将每个向量旋转一个角度θ∗i和θ∗j，其中i，j描述了每个向量的句子位置：

q^iTx^j=qiTRθ,i−jxj. q^​iT​x^j​=qiT​Rθ,i−j​xj​. Rθ,i−j因此表示一个旋转矩阵。θ在训练过程中不会学习，而是根据训练期间最大输入序列长度设定为预定义值。

通过这样做，qi和qj之间的概率得分只会受到影响，如果i≠j并且仅取决于相对距离i−j，而不考虑每个向量的具体位置i和j。

RoPE在今天最重要的多个LLM中使用，例如：

• Falcon
• Llama
• PaLM

作为一种替代方案，ALiBi提出了一种更简单的相对位置编码方案。输入标记之间的相对距离以预定义值m的负整数形式添加到softmax计算之前的每个查询-键条目的QKT矩阵中。

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

ALiBi在今天最重要的多个LLM中使用，例如：

• MPT
• BLOOM

RoPE和ALiBi的位置编码都可以推广到训练期间未见过的输入长度，然而，已经证明对于ALiBi来说，推广效果比RoPE更好。对于ALiBi，只需增加下三角位置矩阵的值，使其与输入序列的长度相匹配。对于RoPE，在传递比训练期间见过的文本输入要长得多的文本输入时，保持相同的θ会导致结果不佳，参见Press等人的研究。然而，社区已经找到了一些有效的技巧来调整θ，从而使RoPE位置嵌入在推广的文本输入序列中表现良好（请参阅此处）。

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

• 关于文本输入的位置线索应直接提供给自注意力层的QKT矩阵
• 应鼓励LLM学习每个位置编码之间的恒定相对距离
• 文本输入标记之间的距离越远，它们的查询-值概率越低。RoPE通过增加查询-键向量之间的角度来减小它们的向量乘积。ALiBi通过向向量乘积添加大的负数来实现。

总之，用于处理大型文本输入的任务的LLM最好使用相对位置嵌入，例如RoPE和ALiBi。还要注意，即使只对长度为N1=2048的固定长度进行了RoPE和ALiBi的训练，仍然可以通过推广位置嵌入的方式在实践中使用比N1大得多的文本输入，如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("input_ids的形状", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

输出:

input_ids的形状 torch.Size([1, 21])
input_ids的形状 torch.Size([1, 22])
input_ids的形状 torch.Size([1, 23])
input_ids的形状 torch.Size([1, 24])
input_ids的形状 torch.Size([1, 25])
[' 这是一个Python函数']

我们可以看到，每次我们通过刚刚采样的标记增加文本输入标记。

除了很少的例外情况外，LLMs都是使用因果语言建模目标进行训练的，因此屏蔽了注意力得分的上三角矩阵 – 这就是为什么上面的两个图表中的注意力得分是空白的（即概率为0）。有关因果语言建模的快速回顾，您可以参考Illustrated Self Attention博客。

因此，标记永远不依赖于先前的标记，更具体地说，如果j>i j > i j>i，那么qi \mathbf{q}_i qi​向量将从未与任何键、值向量kj,vj \mathbf{k}_j, \mathbf{v}_j kj​,vj​相关联。相反，qi \mathbf{q}_i qi​只与先前的键值向量km<i,vm<i , for m∈{0,…i−1} \mathbf{k}_{m < i}, \mathbf{v}_{m < i} \text{ , for } m \in \{0, \ldots i – 1\} km<i​,vm<i​ , for m∈{0,…i−1} 进行关注。为了减少不必要的计算，因此可以为每个层的键值向量缓存所有先前的时间步。

在下面的示例中，我们将要求LLM使用键值缓存，通过将use_cache标志传递给forward调用并将其与当前标记一起传递。

past_key_values = None # past_key_values是键值缓存
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("input_ids的形状", input_ids.shape)
  print("键值缓存的长度", len(past_key_values[0][0]))  # past_key_values的形状为[num_layers, 0表示k, 1表示v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

输出:

input_ids的形状 torch.Size([1, 20])
键值缓存的长度 20
input_ids的形状 torch.Size([1, 20])
键值缓存的长度 21
input_ids的形状 torch.Size([1, 20])
键值缓存的长度 22
input_ids的形状 torch.Size([1, 20])
键值缓存的长度 23
input_ids的形状 torch.Size([1, 20])
键值缓存的长度 24
[' 这', '是', '一个', 'Python', '函数']

正如我们所看到的，使用键值缓存时，文本输入令牌的长度不会增加，而是保持为单个输入向量。而键值缓存的长度在每个解码步骤中都会增加一个。

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

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

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

应始终使用键值缓存，因为它可以产生相同的结果，并且对于较长的输入序列可以显著加速。当使用文本流水线或 generate 方法时，Transformers 默认启用键值缓存。

请注意，键值缓存对于聊天等需要多次自回归解码的应用特别有用。让我们看一个例子。

User: 法国有多少人口？
Assistant: 大约有7500万人口住在法国
User: 德国有多少人口？
Assistant: 德国有大约8100万居民

在这个聊天中，LLM 运行两次自回归解码：

• 1. 第一次，键值缓存为空，输入提示为"User: 法国有多少人口？"，模型自回归生成文本"大约有7500万人口住在法国"，同时在每个解码步骤中增加键值缓存。
• 1. 第二次，输入提示为"User: 法国有多少人口？ \n Assistant: 大约有7500万人口住在法国 \n User: 德国有多少人口？"。由于缓存的存在，前两个句子的所有键值向量已经计算出来。因此，输入提示只包含"User: 德国有多少人口？"。在处理缩短的输入提示时，它的计算的键值向量会与第一次解码的键值缓存进行连接。第二个助手的回答"德国有大约8100万居民"是在键值缓存中包含编码的键值向量"User: 法国有多少人口？ \n Assistant: 大约有7500万人口住在法国 \n User: 德国有多少人口？"的情况下自回归生成的。

这里需要注意两件事：

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

然而，有一个问题。虽然 QKT \mathbf{QK}^T QKT 矩阵所需的峰值内存大大减少，但将键值缓存保留在内存中可能会在长输入序列或多轮对话中导致内存占用过高。请记住，键值缓存需要存储所有先前输入向量 xi, 对于 i∈{1,…,c−1} \mathbf{x}_i \text{, for } i \in \{1, \ldots, c – 1\} xi​, 对于 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亿个浮点数值需要大约15GB的内存，这大约是模型权重本身的一半！研究人员提出了两种方法，可以显著减少存储键值缓存的内存成本：

• 1. 多查询注意力（MQA）

多查询注意力是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，我们可以将所需的内存消耗从15GB减少到小于400MB，输入序列长度为16000。

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

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

MQA已经被社区广泛采用，现在许多最受欢迎的LLM都在使用：

• Falcon
• PaLM
• MPT
• BLOOM

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

• 1. 分组查询注意力（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等将变得更快，可提供更多的内存，但仍应始终确保使用最佳的可用算法和架构以获得最好的性能 🤗