Transformers 文档

优化LLMs的速度和内存

优化LLMs的速度和内存

大型语言模型(LLMs)如GPT3/4、FalconLlama在处理以人为中心的任务方面的能力正在迅速提升,它们已成为现代知识型产业中不可或缺的工具。然而,将这些模型部署到实际任务中仍然具有挑战性:

  • 为了展示接近人类的文本理解和生成能力,目前的大型语言模型(LLMs)需要由数十亿个参数组成(参见Kaplan et al, Wei et. al)。这因此增加了推理时的内存需求。
  • 在许多现实世界的任务中,LLMs需要被提供大量的上下文信息。这就要求模型在推理过程中具备管理非常长的输入序列的能力。

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

在本指南中,我们将介绍高效LLM部署的有效技术:

  1. 低精度:研究表明,以降低的数值精度操作,即8位和4位,可以在不显著降低模型性能的情况下实现计算优势。

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

  3. 架构创新:考虑到LLMs在推理过程中总是以相同的方式部署,即使用长输入上下文进行自回归文本生成,已经提出了专门的模型架构,以实现更高效的推理。模型架构中最重要的进展是AlibiRotary embeddingsMulti-Query Attention (MQA)Grouped-Query-Attention (GQA)

在本指南中,我们将从张量的角度提供自回归生成的分析。我们深入探讨采用低精度的优缺点,全面探索最新的注意力算法,并讨论改进的LLM架构。在此过程中,我们运行实际示例,展示每个功能的改进。

1. 低精度

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

在撰写本指南时,LLMs 至少由数十亿个参数组成。每个参数由一个十进制数组成,例如 4.5689,通常存储在 float32bfloat16float16 格式中。这使得我们可以轻松计算将 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 的显存
  • 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和H100,提供80GB的显存。之前列出的大多数模型仅加载就需要超过80GB,因此必须使用张量并行和/或流水线并行

🤗 Transformers 默认不支持张量并行,因为这需要模型架构以特定的方式编写。如果您有兴趣以支持张量并行的方式编写模型,请随时查看 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的显存。

如今几乎所有模型都以bfloat16进行训练,如果你的GPU支持bfloat16,就没有理由以完整的float32精度运行模型。Float32不会比用于训练模型的精度提供更好的推理结果。

如果您不确定模型权重在Hub上以何种格式存储,您可以随时查看检查点配置中的"torch_dtype"例如 这里。建议在加载时使用from_pretrained(..., torch_dtype=...)将模型设置为与配置中相同的精度类型,除非原始类型是float32,在这种情况下可以使用float16bfloat16进行推理。

让我们定义一个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的显存怎么办?研究发现,模型权重可以量化为8位或4位,而不会显著影响性能(参见Dettmers等人)。 模型甚至可以量化为3位或2位,性能损失在可接受范围内,如最近的GPTQ论文所示 🤯。

在不深入太多细节的情况下,量化方案旨在减少权重的精度,同时尽量保持模型的推理结果尽可能准确(尽可能接近bfloat16)。 请注意,量化在文本生成中特别有效,因为我们只关心选择最可能的下一个标记集,而不太关心下一个标记logit分布的确切值。 重要的是下一个标记logit分布大致保持不变,以便argmaxtopk操作给出相同的结果。

有多种量化技术,我们在这里不会详细讨论,但一般来说,所有量化技术的工作方式如下:

    1. 将所有权重量化为目标精度
    1. 加载量化权重,并以bfloat16精度传递输入向量序列
    1. 动态将权重反量化为bfloat16,以便以bfloat16精度与输入向量进行计算

简而言之,这意味着输入权重矩阵的乘法,其中X X 输入W W 是权重矩阵,Y Y 是输出: Y=XW Y = X * W

被更改为 Y=Xdequantize(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, 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位量化使用了更激进的量化方法,导致quantize \text{quantize} dequantize \text{dequantize} 在推理过程中耗时更长。

del model
del pipe
flush()

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

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

有关量化的更多信息,并了解如何量化模型以需要比4位更少的GPU VRAM内存,我们建议查看AutoGPTQ实现。

作为结论,重要的是要记住,模型量化在提高内存效率的同时,可能会牺牲准确性和在某些情况下的推理时间。

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

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

2. Flash Attention

当今表现最佳的LLMs或多或少共享相同的基本架构,该架构包括前馈层、激活层、层归一化层,以及最关键的自注意力层。

自注意力层是大型语言模型(LLMs)的核心,它们使模型能够理解输入标记之间的上下文关系。 然而,自注意力层的峰值GPU内存消耗在计算和内存复杂度上随着输入标记数量(也称为序列长度)的增加而呈二次方增长,我们在下文中用N N 表示。 对于较短的输入序列(最多1000个输入标记),这一点并不明显,但对于较长的输入序列(大约16000个输入标记),这就成为一个严重的问题。

让我们更仔细地看一下。计算自注意力层输出的公式,对于长度为N N 的输入X \mathbf{X} 是: 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} X=(x1,...xN) \mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N}) 因此是注意力层的输入序列。投影Q \mathbf{Q} K \mathbf{K} 将各自由N N 个向量组成,导致QKT \mathbf{QK}^T 的大小为N2 N^2

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

简而言之,对于大型输入上下文,默认的自注意力算法很快就会变得内存消耗过大。

随着LLMs在文本理解和生成方面的进步,它们被应用于越来越复杂的任务。虽然模型曾经只能处理几句话的翻译或摘要,但现在它们能够管理整个页面,这要求它们具备处理大量输入长度的能力。

我们如何摆脱大输入长度所需的高昂内存需求?我们需要一种新的方法来计算自注意力机制,以摆脱QKT QK^T 矩阵。Tri Dao et al. 开发了这样一种新算法,并将其称为 Flash Attention

简而言之,Flash Attention 将 V×Softmax(QKT\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) 计算分解,并通过多次 softmax 计算步骤迭代计算输出的较小块: OisijaOi+sijbVj×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}

使用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的内存节省效果。 我们附加了原始文本提示 "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个标记的截断。这并不令人惊讶,因为我们为了演示目的重复了系统提示十次,从而提示模型重复自己。

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

让我们测量一下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 改进LLMs的位置嵌入

自注意力机制将每个标记与其他标记相关联。例如,文本输入序列“Hello”, “I”, “love”, “you”Softmax(QKT) \text{Softmax}(\mathbf{QK}^T) 矩阵可能如下所示:

每个单词标记都被赋予了一个概率质量,它关注所有其他单词标记,因此与所有其他单词标记建立了关系。例如,单词“love”以5%的概率关注单词“Hello”,以30%的概率关注“I”,并以65%的概率关注自身。

一个基于自注意力机制但没有位置嵌入的LLM在理解文本输入之间的位置关系时会遇到很大的困难。这是因为由QKT \mathbf{QK}^T 计算出的概率分数将每个词标记与O(1) O(1) 计算中的其他词标记相关联,而不管它们之间的相对位置距离如何。因此,对于没有位置嵌入的LLM来说,每个标记似乎与所有其他标记的距离相同,例如区分“Hello I love you”“You love I hello”将非常具有挑战性。

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

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

与使用固定的位置嵌入不同,其他人(如Devlin等人)使用了学习的位置编码,其中位置嵌入P \mathbf{P} 在训练过程中学习。

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

  1. 正弦和学习的嵌入都是绝对位置嵌入,为每个位置ID编码一个唯一的嵌入:。正如Huang et al.Su et al.所示,绝对位置嵌入会导致长文本输入的LLM性能较差。对于长文本输入,如果模型学习输入标记之间的相对位置距离而不是它们的绝对位置,则是有利的。
  2. 当使用学习到的位置嵌入时,LLM 必须在固定的输入长度上进行训练,这使得它难以外推到比训练时更长的输入长度。

最近,能够解决上述问题的相对位置嵌入变得更加流行,最显著的是:

无论是RoPE还是ALiBi都认为,最好在自注意力算法中直接提示LLM关于句子顺序的信息,因为正是在那里,词标记被相互关联起来。更具体地说,句子顺序应该通过修改QKT \mathbf{QK}^T 计算来提示。

在不深入太多细节的情况下,RoPE 指出位置信息可以编码到查询-键对中,例如qi \mathbf{q}_i xj \mathbf{x}_j 通过将每个向量旋转一个角度θi \theta * i θj \theta * j 分别与i,j i, j 描述每个向量的句子位置: q^iTx^j=qiTRθ,ijxj. \mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j. Rθ,ij \mathbf{R}_{\theta, i - j} 因此表示一个旋转矩阵。θ \theta 在训练期间学习,而是设置为一个预定义的值,该值取决于训练期间的最大输入序列长度。

通过这样做,qi \mathbf{q}_i qj \mathbf{q}_j 之间的概率得分仅在ij i \ne j 时受到影响,并且仅取决于相对距离ij i - j ,而不受每个向量的具体位置i i j j 的影响。

RoPE 被用于当今多个最重要的LLM中,例如:

作为替代方案,ALiBi 提出了一种更简单的相对位置编码方案。输入标记之间的相对距离被添加为负整数,并通过预定义的值 m 进行缩放,然后添加到QKT \mathbf{QK}^T 矩阵的每个查询-键条目中,在 softmax 计算之前。

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

ALiBi 被用于当今多个最重要的LLM中,例如:

无论是RoPE还是ALiBi位置编码,都可以外推到训练期间未见过的输入长度,而已经证明,与RoPE相比,ALiBi的外推效果更好。 对于ALiBi,只需增加下三角位置矩阵的值以匹配输入序列的长度。 对于RoPE,保持训练期间使用的相同θ \theta 会导致在传递比训练期间见过的文本输入长得多时效果不佳,参见 Press et al.。然而,社区已经找到了一些有效的技巧来调整θ \theta ,从而使RoPE位置嵌入能够很好地处理外推的文本输入序列(参见这里)。

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

  • 关于文本输入的位置提示应直接提供给 自注意力层的矩阵
  • LLM 应该被激励去学习一个恒定的相对距离位置编码彼此之间的关系
  • 文本输入标记之间的距离越远,它们的查询值概率就越低。RoPE和ALiBi都降低了彼此远离的标记的查询键概率。RoPE通过增加查询键向量之间的角度来减少它们的向量积。ALiBi通过在向量积上添加大的负数来实现这一点。

总之,对于需要处理大量文本输入的任务,使用相对位置嵌入(如RoPE和ALiBi)训练的LLM效果更好。还需注意的是,即使一个使用RoPE和ALiBi的LLM仅在固定长度(例如N1=2048 N_1 = 2048 )上进行了训练,它仍然可以在实践中用于处理比N1 N_1 大得多的文本输入,例如N2=8192>N1 N_2 = 8192 > N_1 ,通过外推位置嵌入来实现。

3.2 键值缓存

使用LLMs进行自回归文本生成的工作原理是:迭代地输入一个序列,采样下一个标记,将该标记附加到输入序列中,并继续这样做,直到LLM生成一个表示生成完成的标记。

请查看Transformer’s Generate Text Tutorial以获取更直观的自回归生成工作原理的解释。

让我们运行一个快速的代码片段来展示自回归在实践中是如何工作的。我们将简单地通过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 博客

因此,标记从不依赖于先前的标记,更具体地说,qi \mathbf{q}_i 向量从不与任何键、值向量kj,vj \mathbf{k}_j, \mathbf{v}_j 如果j>i j > i 。相反qi \mathbf{q}_i 只关注先前的键值向量km<i,vm<i , 对于 m{0,i1} \mathbf{k}_{m < i}, \mathbf{v}_{m < i} \text{ , 对于 } 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']

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

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

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

  • 计算效率显著提高,因为与计算完整的矩阵相比,执行的计算更少。这导致推理速度的提高
  • 所需的最大内存不会随着生成的令牌数量呈二次方增加,而是仅线性增加。

人们应该始终利用键值缓存,因为它能带来相同的结果,并且对于较长的输入序列能显著加速。当使用文本管道或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. 第一次,键值缓存为空,输入提示为"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才能理解对话的所有先前上下文。例如,对于上面的例子,当用户询问"And how many are in Germany"时,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 矩阵所需的峰值内存显著减少,但对于长输入序列或多轮聊天,将键值缓存保存在内存中可能会变得非常耗费内存。请记住,键值缓存需要存储所有先前输入向量的键值向量xi, 对于 i{1,,c1} \mathbf{x}_i \text{, 对于 } 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的RAM,这大约是模型权重本身的一半! 研究人员提出了两种方法,可以显著减少存储键值缓存的内存成本,这些方法将在接下来的小节中探讨。

3.2.2 多查询注意力 (MQA)

Multi-Query-Attention 是在 Noam Shazeer 的论文《快速Transformer解码:一个写头就是你所需要的》中提出的。正如标题所说,Noam 发现,与其使用 n_head 键值投影权重,可以使用一个在所有注意力头之间共享的单一头值投影权重对,而不会显著降低模型的性能。

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

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

除了节省内存外,MQA 还提高了计算效率,如下所述。 在自回归解码中,需要重新加载大的键值向量,将其与当前的键值向量对连接起来,然后输入到qcKT \mathbf{q}_c\mathbf{K}^T 计算中。对于自回归解码,持续重新加载所需的内存带宽可能成为一个严重的时间瓶颈。通过减小键值向量的大小,需要访问的内存减少,从而减轻了内存带宽瓶颈。更多详情,请查看 Noam 的论文

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

MQA 已被社区广泛采用,现在许多最流行的 LLM 都在使用它:

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

3.2.3 分组查询注意力 (GQA)

Grouped-Query-Attention,由Google的Ainslie等人提出,发现与使用传统的多头键值投影相比,使用MQA通常会导致质量下降。论文认为,通过不那么剧烈地减少查询头投影权重的数量,可以保留更多的模型性能。应该使用n < n_head个键值投影权重,而不是仅使用单个键值投影权重。通过选择n为比n_head小得多的值,例如2、4或8,几乎可以保留MQA的所有内存和速度优势,同时牺牲较少的模型容量,从而可能减少性能损失。

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

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

作为结论,强烈建议在LLM部署自回归解码并需要处理大输入序列(例如聊天)时使用GQA或MQA。

结论

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

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

< > Update on GitHub