LLM推理优化
大型语言模型(LLMs)通过生成显示出高度理解和流畅性的文本,将文本生成应用(如聊天和代码补全模型)推向了新的高度。但使LLMs如此强大的因素——即它们的规模——也给推理带来了挑战。
基本推理速度较慢,因为需要反复调用LLMs来生成下一个标记。随着生成的进行,输入序列不断增加,LLMs处理所需的时间也越来越长。LLMs还拥有数十亿个参数,这使得在内存中存储和处理所有这些权重成为一个挑战。
本指南将向您展示如何使用Transformers中可用的优化技术来加速LLM推理。
Hugging Face 还提供了 Text Generation Inference (TGI),这是一个专门用于部署和服务高度优化的LLM进行推理的库。它包含了一些Transformers中没有的面向部署的优化功能,例如用于提高吞吐量的连续批处理和多GPU推理的张量并行。
静态kv缓存和torch.compile
在解码过程中,LLM会为每个输入标记计算键值(kv)值,由于它是自回归的,每次都会计算相同的kv值,因为生成的输出现在成为输入的一部分。这并不高效,因为你每次都在重新计算相同的kv值。
为了优化这一点,你可以使用kv-cache来存储过去的键和值,而不是每次都重新计算它们。然而,由于kv-cache随着每个生成步骤的增长并且是动态的,它阻止你利用torch.compile
,这是一个强大的优化工具,可以将PyTorch代码融合到快速和优化的内核中。我们有一个完整的指南专门介绍kv-caches 这里。
静态kv缓存通过预先分配kv缓存的大小到最大值来解决这个问题,这允许你将其与torch.compile
结合使用,以实现高达4倍的加速。你的加速效果可能会因模型大小(较大的模型加速效果较小)和硬件而异。
根据任务的复杂性,静态kv-cache的使用有三种方式:
- 基本用法:只需在
generation_config
中设置一个标志(推荐); - 高级用法:处理用于多轮生成的缓存对象或自定义生成循环;
- 高级用法:将整个
generate
函数编译成单个图,如果单个图对您有相关性。
请选择下面的正确标签以获取每种风格的进一步说明。
无论使用哪种策略与torch.compile
,如果你将LLM输入左填充到一组有限的值,你可以避免与形状相关的重新编译。pad_to_multiple_of
tokenizer flag是你的好帮手!
在这个例子中,我们使用Gemma模型。我们需要做的就是:
- 访问模型的
generation_config
属性并将cache_implementation
设置为“static”; - 在模型上调用
torch.compile
以使用静态 kv-cache 编译前向传递。
就这样!
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # To prevent long warnings :)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", torch_dtype="auto", device_map="auto")
model.generation_config.cache_implementation = "static"
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device.type)
outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The theory of special relativity states 1. The speed of light is constant in all inertial reference']
在底层,generate
将尝试重用相同的缓存对象,从而避免每次调用时重新编译。避免重新编译对于充分利用 torch.compile
至关重要,您应该注意以下几点:
- 如果批次大小发生变化或最大输出长度在调用之间增加,缓存将不得不重新初始化,从而触发新的编译;
- 编译后的函数的前几次调用较慢,因为函数正在被编译。
对于更高级的静态缓存使用,例如多轮对话,我们建议在generate()之外实例化和操作缓存对象。请参阅高级用法选项卡。
推测解码
如需更深入的解释,请查看辅助生成:低延迟文本生成的新方向博客文章!
自回归的另一个问题是,对于每个输入标记,你需要在每次前向传播时加载模型权重。这对于拥有数十亿参数的大型语言模型(LLMs)来说既慢又麻烦。推测性解码通过使用第二个更小更快的辅助模型来生成候选标记,这些候选标记由更大的LLM在单次前向传播中验证,从而缓解了这种减速。如果验证的标记是正确的,LLM基本上可以“免费”获得它们,而不必自己生成它们。由于验证前向传播确保了与LLM自己生成相同的输出,因此准确性不会降低。
为了获得最大的加速效果,辅助模型应该比LLM小得多,以便能够快速生成标记。辅助模型和LLM模型还必须共享相同的标记器,以避免重新编码和解码标记。
推测解码仅支持贪婪搜索和采样解码策略,并且也不支持批量输入。
通过加载一个辅助模型并将其传递给generate()方法,启用推测性解码。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from accelerate.test_utils.testing import get_backend
device, _, _ = get_backend() # automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("Einstein's theory of relativity states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype="auto").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
["Einstein's theory of relativity states that the speed of light is constant. "]
提示查找解码
提示查找解码是一种推测性解码的变体,它也兼容贪婪搜索和采样。提示查找特别适用于输入基础任务——例如摘要——其中提示和输出之间经常有重叠的单词。这些重叠的n-gram被用作LLM候选标记。
要启用提示查找解码,请在prompt_lookup_num_tokens
参数中指定应重叠的令牌数量。然后,您可以将此参数传递给generate()方法。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from accelerate.test_utils.testing import get_backend
device, _, _ = get_backend() # automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("The second law of thermodynamics states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype="auto").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, prompt_lookup_num_tokens=3)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The second law of thermodynamics states that entropy increases with temperature. ']
注意力优化
Transformer模型的一个已知问题是,自注意力机制的计算和内存需求随着输入标记数量的增加呈二次方增长。这一限制在LLMs中更加明显,因为它们处理更长的序列。为了解决这个问题,可以尝试使用FlashAttention2或PyTorch的缩放点积注意力(SDPA),这些是更高效的内存注意力实现,可以加速推理。
FlashAttention-2
FlashAttention 和 FlashAttention-2 将注意力计算分解为更小的块,并减少了对GPU内存的中间读写操作次数,从而加速推理。FlashAttention-2 在原始 FlashAttention 算法的基础上进行了改进,通过在序列长度维度上进行并行化,并更好地在硬件上分配工作,以减少同步和通信开销。
要使用 FlashAttention-2,请在 from_pretrained() 方法中设置 attn_implementation="flash_attention_2"
。
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2b",
quantization_config=quant_config,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
使用 torch.compile 和无填充数据整理的微调
除了优化推理外,您还可以通过在微调过程中利用torch.compile和使用无填充数据整理器来提高大型语言模型的训练效率。这种方法可以显著加快训练速度并减少计算开销。
以下是您如何使用TRL库中的SFTTrainer微调Llama模型的方法,启用torch_compile并使用无填充数据整理器:
#################### IMPORTS ###################
import math
import datasets
import dataclasses
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments
)
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM
#################### MODEL LOADING WITH FLASH ATTENTION ###################
model_name = "meta-llama/Llama-3.2-1B"
model = AutoModelForCausalLM.from_pretrained(
model_name,
attn_implementation="flash_attention_2" # Enables FlashAttention-2
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
#################### DATA PREPROCESSING (PADDING-FREE) ###################
response_template = "\n### Label:"
response_template_ids = tokenizer.encode(
response_template, add_special_tokens=False
)[2:] # Exclude special tokens
data_collator = DataCollatorForCompletionOnlyLM(
response_template_ids=response_template_ids,
tokenizer=tokenizer,
ignore_index=-100,
padding_free=True # Enables padding-free collation
)
def format_dataset(example):
return {
"output": example["output"] + tokenizer.eos_token
}
data_files = {"train": "path/to/dataset"} # Replace with your dataset path
json_dataset = datasets.load_dataset("json", data_files=data_files)
formatted_train_dataset = json_dataset["train"].map(format_dataset)
################# TRAINING CONFIGURATION ############################
train_args = TrainingArguments(
num_train_epochs=5,
per_device_train_batch_size=4,
per_device_eval_batch_size=4,
gradient_accumulation_steps=4,
learning_rate=1e-5,
weight_decay=0.0,
warmup_ratio=0.03,
lr_scheduler_type="cosine",
logging_steps=1,
include_tokens_per_second=True,
save_strategy="epoch",
output_dir="output",
torch_compile=True, # Enables torch.compile
torch_compile_backend="inductor",
torch_compile_mode="default"
)
# Convert TrainingArguments to SFTConfig
transformer_train_arg_fields = [x.name for x in dataclasses.fields(SFTConfig)]
transformer_kwargs = {
k: v
for k, v in train_args.to_dict().items()
if k in transformer_train_arg_fields
}
training_args = SFTConfig(**transformer_kwargs)
####################### FINE-TUNING #####################
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=formatted_train_dataset,
data_collator=data_collator,
dataset_text_field="output",
args=training_args,
)
trainer.train()
PyTorch 缩放点积注意力
在 PyTorch 2.0 中,Scaled dot product attention (SDPA) 自动启用,并支持 FlashAttention、xFormers 和 PyTorch 的 C++ 实现。如果你使用的是 CUDA 后端,SDPA 会选择性能最佳的注意力算法。对于其他后端,SDPA 默认使用 PyTorch 的 C++ 实现。
SDPA支持FlashAttention-2,只要您安装了最新版本的PyTorch。
使用 torch.backends.cuda.sdp_kernel 上下文管理器来显式启用或禁用三种注意力算法中的任何一种。例如,设置 enable_flash=True
以启用 FlashAttention。
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2b",
torch_dtype=torch.bfloat16,
)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
outputs = model.generate(**inputs)
量化
量化通过以较低精度存储LLM权重来减小其大小。这意味着内存使用量减少,并且如果你受限于GPU内存,加载LLM进行推理将更加容易。如果你不受GPU限制,你不一定需要量化你的模型,因为这可能会带来一些小的延迟成本(除了AWQ和融合AWQ模块),因为需要额外的步骤来量化和反量化权重。
有许多量化库可供使用(详情请参阅量化指南),例如Quanto、AQLM、AWQ和AutoGPTQ。欢迎尝试它们,看看哪个最适合您的用例。我们还建议阅读🤗 Transformers中本地支持的量化方案概述博客文章,该文章比较了AutoGPTQ和bitsandbytes。
使用下面的模型内存计算器来估算和比较加载模型所需的内存。例如,尝试估算加载Mistral-7B-v0.1所需的内存。
要以半精度加载Mistral-7B-v0.1,请在from_pretrained()方法中将torch_dtype
参数设置为torch.bfloat16
。这需要13.74GB的内存。
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-7B-v0.1", torch_dtype=torch.bfloat16, device_map="auto",
)
要加载量化模型(8位或4位)进行推理,请尝试使用bitsandbytes并将load_in_4bit
或load_in_8bit
参数设置为True
。以8位加载模型仅需要6.87 GB的内存。
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-7B-v0.1", quantization_config=quant_config, device_map="auto"
)