LLM工程师手册——推理优化

596 阅读35分钟

部署LLMs具有挑战性,因为它们对计算和内存的需求非常高。高效运行这些模型需要使用专用加速器,如GPU或TPU,以并行化操作并提高吞吐量。尽管一些任务(如文档生成)可以在夜间批量处理,但其他任务(如代码补全)需要低延迟和快速生成。因此,优化推理过程——即这些模型根据输入数据进行预测的过程——对于许多实际应用至关重要。优化包括减少生成第一个token的时间(延迟)、增加每秒生成的token数量(吞吐量)以及最小化LLM的内存占用。

确实,简单的部署方法会导致硬件利用率低下,吞吐量和延迟表现不佳。幸运的是,各种优化技术已经出现,可显著加速推理。本章将探讨关键方法,如推测解码、模型并行和权重量化,展示如何通过合理的实现实现2-4倍或更高的加速效果。我们还将介绍三个流行的推理引擎(Text Generation Inference、vLLM和TensorRT-LLM),并比较它们在推理优化方面的功能。

本章将涵盖以下主题:

  • 模型优化策略
  • 模型并行
  • 模型量化

完成本章后,您将理解LLM推理中的核心挑战,并熟悉包括模型并行和权重量化在内的最新优化技术。

本章中的所有代码示例都可以在GitHub上找到:github.com/PacktPublis…

模型优化策略

当前使用的大多数LLMs,如GPT或Llama,基于仅解码器的Transformer架构。仅解码器架构专为文本生成任务而设计,根据前文预测序列中的下一个词,使其在生成上下文相关的文本续写方面非常有效。

相比之下,BERT等仅编码器架构则侧重于通过详细的嵌入来理解和表示输入文本,擅长需要全面上下文理解的任务,如文本分类和命名实体识别。最后,T5等编码器-解码器架构结合了两者的功能:编码器处理输入文本生成丰富上下文的表示,解码器则利用该表示生成输出文本。这种双重结构在翻译和摘要等序列到序列任务中特别强大,因为它既需要理解输入上下文,又需要生成相关输出。

本书只关注在LLM领域占主导地位的仅解码器架构。

image.png

如图8.1所示,仅解码器模型的基本推理过程包括以下步骤:

  1. 对输入提示进行分词,并通过嵌入层和位置编码。
  2. 使用多头注意力机制计算每个输入token的键和值对。
  3. 使用计算出的键和值,逐步生成输出token,每次生成一个。

尽管步骤1和2计算量大,但它们由高度可并行的矩阵乘法组成,可以在GPU和TPU等加速器上实现高硬件利用率。

真正的挑战在于第3步的token生成过程本质上是顺序的——要生成下一个token,必须已经生成所有先前的token。这导致输出序列逐个token地迭代增长,难以利用硬件的并行计算能力。解决这一瓶颈是推理优化的核心目标之一。

在本节中,我们将详细介绍多种常用的优化策略,以加快推理速度并减少视频随机存取存储器(VRAM)使用量,包括实施(静态)KV缓存、持续批处理、推测解码和优化注意力机制。

KV缓存

我们已经看到,LLMs逐个token生成文本,这很慢,因为每个新预测都依赖于之前的整个上下文。例如,为预测序列中的第100个token,模型需要tokens 1到99的上下文。预测第101个token时,它又需要tokens 1到99的信息以及第100个token。这种重复计算尤其低效。

键-值(KV)缓存解决了此问题,它存储由自注意力层生成的键-值对。模型无需为每个新token重新计算这些对,而是从缓存中获取它们,从而显著加快生成速度。

该技术的示例可参见图8.2。

image.png

当生成一个新token时,仅需要计算该token的键和值并将其添加到缓存中。KV缓存是一种立即见效的优化,每个流行的工具和库中都实现了它。一些实现会为模型的每一层单独维护KV缓存。

KV缓存的大小随着token数量(TTT)和若干模型维度(如层数LLL、注意力头数HHH、头部维度DDD以及参数精度(以字节计)PPP)的增加而增长:

image.png

对于使用16位精度的典型70亿参数模型,若序列长度较长(超过2048个token),其KV缓存大小会超过2GB。更大规模、更多层数和更高嵌入维度的模型会要求更大的内存。

由于KV缓存随着每一步生成动态增长,因此它无法利用torch.compile这一强大的优化工具,该工具能够将PyTorch代码融合成快速优化的内核。静态KV缓存通过将KV缓存大小预分配为一个最大值,解决了这一问题,从而可以将其与torch.compile结合,在前向传递中实现高达4倍的加速。

使用transformers库配置模型以启用静态KV缓存,请按照以下步骤操作:

  1. 我们导入希望优化的分词器和模型:

    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    model_id = "google/gemma-2b-it"
    tokenizer = AutoTokenizer.from_pretrained(model_id) 
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
    
  2. 实现静态缓存,我们将模型生成配置中的缓存实现更改为静态:

    model.generation_config.cache_implementation = "static"
    
  3. 现在KV缓存已经设置为静态,我们可以使用torch.compile编译模型:

    compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
    
  4. 我们对输入问题“2+2等于多少?”进行分词并将其存储在GPU上(如果没有GPU,则存储在CPU上):

    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer("What is 2+2?", return_tensors="pt").to(device)
    
  5. 使用generate()方法获取模型输出,并通过batch_decode()解码并打印答案:

    outputs = model.generate(**inputs, do_sample=True, temperature=0.7, max_new_tokens=64)
    print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    

    输出结果为包含输入和输出的列表,正确回答了我们的问题:

    ['What is 2+2?\n\nThe answer is 4. 2+2 = 4.']
    

注意,静态缓存并不适用于所有架构。有关支持的架构详情,请查阅transformers文档。

高效管理KV缓存至关重要,因为它会快速耗尽可用的GPU内存并限制可处理的批量大小。这促使人们开发了内存高效的注意力机制和其他技术,我们将在最后一节中详细介绍这些内容。

连续批处理

批处理(同时处理多个推理请求)是一种标准的高吞吐量方法。较大的批处理量可以分摊模型权重的内存成本,并一次将更多数据传输到GPU,更好地饱和其并行计算能力。

然而,仅解码器模型由于输入提示长度和期望输出长度的高度可变性,带来了特定的挑战。一些请求的提示较短,仅需一个词的回答,而其他请求则可能输入长篇上下文,并期望多段式回应。

在传统批处理中,需要等待一个批次中最长请求完成后才能开始新的批次,这会导致加速器部分闲置,等待较慢的请求完成,造成资源利用不足。连续批处理(也称为实时批处理)旨在通过在一个请求完成后立即将新请求加入批次来防止空闲时间。

批处理过程的开始仍是填充初始请求,但一旦某个请求完成生成,便从批处理中移除,并加入一个新请求。这样一来,加速器始终处理满批次,实现最高效的硬件利用率。另一个需考虑的因素是需要定期暂停生成过程,以进行预填充,即等待请求的嵌入和编码。找到生成和预填充之间的最佳平衡需要调整等待服务比这一超参数。

连续批处理在大多数推理框架中已原生实现,如Hugging Face的Text Generation Inference(TGI)、vLLM和NVIDIA TensorRT-LLM。

推测解码

另一种强大的优化技术是推测解码,也称为辅助生成。其关键在于,即使有了连续批处理,逐个token的生成过程仍无法充分利用加速器的并行处理能力。推测解码旨在利用剩余的计算容量,通过使用较小的代理模型同时预测多个token(参见图8.3)。

image.png

一般方法如下:

  1. 使用一个更小的模型(例如主模型的蒸馏或剪枝版本)并行预测多个token完成结果,例如在一个步骤中预测5–10个token。
  2. 将这些推测完成结果输入完整模型,以验证哪些预测符合大模型的生成结果。
  3. 保留推测完成结果中最长的匹配前缀,丢弃任何不正确的token。

如果小模型能较好地近似大模型,则可以在一个步骤中生成多个token,从而避免多次运行昂贵的大模型。加速程度取决于小模型的预测准确性——90%的匹配率可能带来3–4倍的加速。

确保两个模型使用相同的分词器非常关键,否则草稿模型生成的token将与大模型生成的不匹配,导致不兼容。以下是使用transformers库实现推测解码的步骤。我们将使用来自阿里云的两个Qwen1.5模型:一个1.8B版本作为主模型,一个0.5B版本作为草稿模型。注意,如果有足够的VRAM,可以使用更大的主模型(如14B、32B、72B或110B)。在这里,由于Google Colab中的T4 GPU VRAM限制,辅助模型应比大模型小得多,以获得最大加速效果。

实现推测解码的步骤:

  1. 加载分词器和两个模型:

    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    model_id = "Qwen/Qwen1.5-1.8B-Chat"
    tokenizer = AutoTokenizer.from_pretrained(model_id) 
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
    draft_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B-Chat", device_map="auto")
    
  2. 对同一输入进行分词并存储在加速器中(如果可用):

    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer("What is 2+2?", return_tensors="pt").to(device)
    
  3. 使用model.generate()并传入assistant_model参数以启用推测解码:

    outputs = model.generate(**inputs, do_sample=True, assistant_model=draft_model, temperature=0.7, max_new_tokens=64)
    print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    

输出结果:

['What is 2+2? 2 + 2 equals 4!']

在此小示例中,加速效果不明显,但在更大模型上会更为显著。

提示查找解码是推测解码的一种变体,适用于摘要等以输入为基础的任务,通常提示与输出之间有重复内容。共享的n元组(n-grams)作为LLM候选token。可以在model.generate()中使用prompt_lookup_num_tokens参数启用提示查找解码:

outputs = model.generate(**inputs, prompt_lookup_num_tokens=4)

结合静态KV缓存与torch.compile、连续批处理以及推测解码技术,LLMs可以在不损失质量的情况下实现2–4倍甚至更高的推理加速。

另一种创建小型代理模型的方法是同时微调小模型与大模型,以获得最大近似性。Medusa技术即是一种代表性方法,它在主模型中插入专用的推测头。Medusa-1方法微调这些推测头,同时冻结大模型;Medusa-2方法则同时微调推测头和大模型。Medusa方法展示了令人印象深刻的效果,使70M参数模型能够在一系列任务上接近7B参数模型的表现。推测解码在TGI中原生支持。

优化的注意力机制

Transformer架构基于注意力机制,计算复杂度随着输入token数量(或序列长度)成平方增长。这对于较长序列来说尤其低效,KV缓存的大小会迅速膨胀。

由Kwon、Li等人(2023)提出的PagedAttention受操作系统中的虚拟内存和分页技术启发,解决了这些内存挑战。PagedAttention将KV缓存划分为块,消除了连续内存分配的需求。每个块包含固定数量的token的键和值。在计算注意力时,PagedAttention内核可以有效地获取这些块,而不论其物理内存位置。

这种分块方式允许实现接近最佳的内存利用率,有助于批量处理更多序列,从而增加吞吐量和GPU利用率。此外,PagedAttention的基于块的方式天然支持从同一提示生成的多个输出序列之间的内存共享。这对于并行采样和束搜索特别有利,因为它们使用相同的提示生成多个输出。共享的内存块减少了冗余计算和内存使用,作者指出这可以减少高达55%的内存开销,并提高多达2.2倍的吞吐量。vLLM库最早实现了PagedAttention,随后TGI和TensorRT-LLM也实现了PagedAttention。

另一种流行的选择是FlashAttention-2,由Tri Dao(2023)开发。它通过将输入和输出矩阵分割为更小的块,确保这些块可以存入GPU的片上SRAM(比高带宽内存更快),从而显著减少了GPU主内存与处理单元之间的数据传输频率。

这种方法结合了在线softmax计算,分别对注意力得分矩阵的每个块计算softmax,而不是一次性对整个矩阵计算。通过保持一个运行中的最大值和指数和,FlashAttention-2可以计算注意力概率,而无需存储大型中间矩阵。

此外,FlashAttention-2的在线softmax计算支持块式处理,既保证了准确性,又显著降低了内存需求。对于训练来说,这尤其重要,因为在反向传播中重新计算中间值(而非存储它们),将内存需求从序列长度的平方降低到线性。

与PagedAttention不同,FlashAttention-2可以通过transformers库中的attn_implementation参数轻松使用:

首先,通过--no-build-isolation安装flash-attn库,以避免安装依赖项:

pip install flash-attn --no-build-isolation

要在推理中使用FlashAttention-2,可以在加载模型时将attn_implementation参数指定为flash_attention_2。例如,以下是如何使用FlashAttention-2加载Mistral-7B-Instruct-v0.3:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
    attn_implementation="flash_attention_2",
)

本节介绍的技术侧重于提高模型处理token的效率。在下一节中,我们将讨论如何将模型和计算分布到多个GPU上。

模型并行

模型并行可以将LLM的内存和计算需求分布在多个GPU上,从而使超大模型的训练和推理成为可能,同时提升吞吐量(每秒生成的token数)等性能。

模型并行的主要方法有三种,每种方法都以不同方式拆分模型的权重和计算:数据并行、流水线并行和张量并行。尽管这些方法最初是为训练开发的,我们可以通过仅关注前向传播来将它们复用于推理。

数据并行

数据并行(DP)是最简单的模型并行类型。它通过在不同GPU上复制模型的副本(见图8.4),使每个GPU同时处理数据的一个子集。在训练过程中,每个GPU上计算出的梯度会被平均并用于更新模型参数,以确保每个副本保持同步。当批量大小过大,无法在单台机器上容纳,或希望加速训练过程时,该方法尤为有用。

image.png

在推理过程中,数据并行(DP)有助于处理并发请求。通过将工作负载分布在多个GPU上,DP方法可减少延迟,因为多个请求可以同时被处理。这种并发处理还提高了吞吐量,使得可以同时处理更多请求。

然而,DP的有效性受到模型大小和GPU间通信开销的限制。确实,在每个GPU上复制模型的参数会降低效率。这意味着该技术仅适用于模型足够小、能够适配单个GPU的情况,导致输入数据所占空间减少,从而限制了批量大小。对于更大的模型或在内存受限时,这可能成为显著缺点。

通常情况下,DP主要用于训练,而流水线并行和张量并行更适合推理。

流水线并行

流水线并行(PP)由Huang等人在2019年的GPipe论文中提出,是一种将大型神经网络的计算负载分布到多个GPU上的策略。与将整个模型复制到每个GPU上的传统数据并行不同,流水线并行将模型的层划分到不同的GPU上。这样,每个GPU只需处理模型的特定部分,从而减轻了单个GPU上的内存负担。

image.png

如图8.5所示,在典型的四路流水线并行拆分中,模型被分为四个部分,每个部分分配给不同的GPU。例如,模型前25%的层可能由GPU 1处理,接下来的25%由GPU 2处理,以此类推。在前向传播过程中,计算激活值并将其传递到下一个GPU。对于训练而言,反向传播则按照相反的顺序执行,将梯度依次传回各个GPU。通常,GPU的数量被称为并行度。

流水线并行的主要优势在于它可以显著减少每个GPU的内存需求。然而,这种方法也引入了新的挑战,尤其是与流水线的顺序特性相关的问题。其中一个主要问题是“流水线气泡”的出现。这些气泡产生于部分GPU空闲,等待来自前一层的激活数据时,这种空闲时间会降低整体效率。

为缓解流水线气泡的影响,引入了微批处理。通过将输入批次拆分为更小的子批次,微批处理确保GPU保持更高的工作率,因为下一个子批次可以在前一个子批次尚未完全完成时开始处理。

image.png

图8.6展示了带有微批处理的流水线并行示例。在此例中,流水线有四个阶段(F0、F1、F2、F3),输入批次被分为四个微批次。GPU 0将依次处理前向路径F0,0、F0,1、F0,2和F0,3。一旦F0,0完成,GPU 1可以立即开始处理F1,0,以此类推。在完成这些前向传播后,GPU 0将等待其他GPU完成各自的前向计算,然后再开始反向路径(B0,3、B0,2、B0,1和B0,0)。

流水线并行在诸如Megatron-LM、DeepSpeed(ZeRO)以及PyTorch的专用流水线并行库(PiPPy)等分布式训练框架中得到实现。截至撰写本文时,只有TensorRT-LLM等少数推理框架支持流水线并行。

张量并行

张量并行(TP)由Shoeby、Patwary、Puri等人在2019年的Megatron-LM论文中提出,是另一种将LLM层的计算分布到多个设备上的流行技术。与流水线并行不同,TP将各层中的权重矩阵拆分,使得计算可以同时进行,从而显著减少内存瓶颈并提高处理速度。

在TP中,大型矩阵(例如MLP中的权重矩阵或自注意力层中的注意力头)被划分到多个GPU上。每个GPU持有这些矩阵的一部分,并在其对应的部分上执行计算。

image.png

例如,在MLP层中,权重矩阵被划分,使每个GPU仅处理权重的一部分(见图8.7)。输入被广播到所有GPU,每个GPU独立计算其对应的输出。然后,通过一次全归约操作将这些部分结果聚合,形成最终输出。

在自注意力层中,由于注意力头的固有并行性,张量并行(TP)尤其高效。每个GPU可以独立计算部分注意力头,从而让模型更有效地处理大型序列。这使得TP比流水线并行更高效,因为流水线并行需要等待前一层的完成。

尽管TP具有优势,但并不适用于神经网络的所有层。像LayerNorm和Dropout这类层涉及整个输入的依赖关系,无法高效分割,通常在设备间复制这些操作。不过,这些操作可以在输入的序列维度上进行分割(即序列并行)。不同的GPU可以在输入序列的不同切片上计算这些层,避免权重复制。此技术仅适用于少数特定层,但对非常长的输入序列可提供额外的内存节省。

此外,TP需要设备之间有高速互连,以最小化通信开销,使其在互连带宽不足的节点间难以实现。

TP在Megatron-LM、DeepSpeed(ZeRO)和PyTorch(FSDP)等分布式训练框架中得以实现,并在大多数推理框架中可用,如TGI、vLLM和TensorRT-LLM。

组合方法

数据并行、张量并行和流水线并行是相互独立的技术,可以结合使用。图8.8展示了如何根据每种方法对一个模型进行拆分:

image.png

结合这些技术可以缓解各自的缺点。流水线并行在内存减少方面效果最佳,但由于流水线气泡牺牲了效率。如果主要限制是将模型适配到GPU内存中,那么流水线并行可能是理想选择。相反,如果低延迟至关重要,那么优先考虑张量并行并接受较大的内存占用可能是更好的折衷方案。在实际应用中,模型可以在深度上划分为多个流水线阶段,并在每个阶段内使用张量并行。

在部署LLMs时,平衡这些折衷并将模型架构映射到可用的硬件加速器上是一个关键挑战。

模型量化

量化指的是使用低精度数据类型表示神经网络的权重和激活。在LLMs的背景下,量化主要集中在降低模型权重和激活的精度。

默认情况下,权重通常以16位或32位浮点格式(FP16或FP32)存储,这种高精度格式带来了较高的内存占用和计算复杂度。量化可以减少内存占用并加速LLMs的推理。

除了这些好处外,将超过300亿参数的大型模型量化到2位或3位精度时,其性能在质量上可能优于较小的模型(如7B–13B的LLMs),也就是说,它们可以在保持可比内存占用的同时获得更好的性能。

在本节中,我们将介绍量化的概念、使用llama.cpp的GGUF、GPTQ和EXL2,并概述其他一些技术。除了本节提供的代码外,您还可以参考AutoQuant(bit.ly/autoquant)使用Google Colab notebook量化您的模型。

量化简介

权重量化有两种主要方法:训练后量化(Post-Training Quantization,PTQ)和量化感知训练(Quantization-Aware Training,QAT)。PTQ是一种简单的技术,即直接将预训练模型的权重转换为低精度格式,而无需重新训练。尽管PTQ易于实现,但可能会导致性能有所下降。相反,QAT在训练或微调阶段进行量化,使模型能够适应低精度权重。与PTQ相比,QAT通常能获得更好的性能,但需要额外的计算资源和代表性训练数据。

数据类型的选择在量化中起着至关重要的作用。浮点数如FP32、FP16(半精度)和BF16(脑浮点)在深度学习中常被使用。这些格式分配固定的位数来表示数值的符号、指数和有效数(尾数)。

image.png

符号位为0表示正数,而1表示负数。指数位控制表示的数值范围(大或小),而有效数位控制数值的精度(小数位数)。将这些表示转换为实数的公式如下:

image.png

图7.7中的数据类型展示了不同的权衡,通过不同的表示方式来说明。例如,FP32使用32位,提供高精度,但需要更多内存。相反,FP16和BF16使用16位,降低了内存占用,但也降低了精度。总体而言,神经网络更倾向于更大的范围而非更高的精度,这也是当硬件支持时BF16成为最受欢迎数据类型的原因。例如,NVIDIA的Ampere架构(如A100、A30等)支持BF16,但前几代架构(如Turing系列的T4、T40等)不支持。

然而,我们并不局限于这三种数据类型。可以使用如INT8(8位整数)等更低精度的数据类型进行量化,从而进一步减少内存占用。简单的量化技术如绝对最大值(absmax)量化和零点量化可以用于将FP32、FP16或BF16权重转换为INT8,如图8.10所示:

image.png

Absmax量化通过将原始权重映射到范围[−127,127][-127, 127][−127,127],具体做法是将权重除以其绝对最大值absmax(w)\text{absmax}(w)absmax(w)并进行缩放:

Xquant=absmax(X)X×127Xquant​=absmax(X)X​×127

例如,如果绝对最大值为3.2(见图8.8),那么权重0.1将被量化为:

Xquant=3.20.1×1273.974Xquant​=3.20.1​×127≈3.97≈4

要进行反量化,我们执行逆运算:

Xdequant=Xquant×127absmax(X)Xdequant​=Xquant​×127absmax(X)​

这意味着在反量化权重时,我们得到:

Xdequant=4×1273.20.1008Xdequant​=4×1273.2​≈0.1008

在此示例中,我们可以看到四舍五入误差为0.00080.00080.0008。在Python中,我们可以使用PyTorch库实现如下代码:

import torch

def absmax_quantize(X):
    # 计算缩放比例
    scale = 127 / torch.max(torch.abs(X))
    # 量化
    X_quant = (scale * X).round()
    return X_quant.to(torch.int8)

另一方面,零点量化考虑了非对称输入分布,通过引入一个零点偏移将权重映射到[−128,127][-128, 127][−128,127]的范围:

image.png

其中,scale\text{scale}scale 和 zeropoint\text{zeropoint}zeropoint 分别表示缩放因子和零点偏移。

如果我们以权重0.1为例,得到的缩放因子为 scale\text{scale}scale 和零点值为 zeropoint\text{zeropoint}zeropoint。权重0.1会被量化为与Absmax方法不同的值,而不是由Absmax提供的值。

我们可以通过应用逆运算轻松地进行反量化:

X_dequant = (X_quant - zeropoint) / scale

在Python中,零点量化可以按以下方式实现:

def zeropoint_quantize(X):
    # 计算值的范围(分母)
    x_range = torch.max(X) - torch.min(X)
    x_range = 1 if x_range == 0 else x_range
    # 计算缩放因子
    scale = 255 / x_range
    # 计算零点偏移
    zeropoint = (-scale * torch.min(X) - 128).round()
    # 按零点偏移缩放并四舍五入
    X_quant = torch.clip((X * scale + zeropoint).round(), -128, 127)
   
    return X_quant.to(torch.int8)

然而,简单的量化方法存在局限性,特别是在处理LLMs中的离群特征时。离群特征是极端的权重值(大约占总值的0.1%),这些值会显著影响量化过程,导致其他值的精度降低。

丢弃这些离群值是不可行的,因为这会降低模型的性能。您可以在图8.11中看到离群值的示例:

image.png

为了解决离群值问题,提出了更先进的量化技术。一个显著的例子是LLM.int8(),由Dettmers等人(2022年)提出。LLM.int8()采用混合精度量化方案,其中离群特征使用FP16处理,而其余值则量化为INT8。这种方法有效地将LLMs的内存占用减少了近2倍,同时最小化了性能下降。

LLM.int8()通过三步执行矩阵乘法。首先,它使用自定义阈值从输入的隐藏状态中提取包含离群特征的列。其次,它对离群值(使用FP16)和非离群值(使用INT8)分别进行矩阵乘法,采用按向量量化的方法。最后,它对非离群结果进行反量化,并将其与离群结果结合,得到最终的FP16输出。

LLM.int8()的有效性已通过实验证明,与原始的FP32模型相比,性能下降几乎可以忽略(<1%)。然而,它确实引入了额外的计算开销,导致大型模型的推理速度约慢20%。使用transformers库,模型可以直接以8位精度加载,如下所示:

from transformers import AutoModelForCausalLM
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)

由Dettmers等人(2023年)提出的NF4是一种4位精度格式,专为QLoRA设计(在第5章讨论)。它也已集成到transformers库中,但需要bitsandbytes库作为依赖项。要以NF4(4位精度)加载模型,可以使用load_in_4bit参数,如下所示:

from transformers import AutoModelForCausalLM
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_4bit=True)

使用GGUF和llama.cpp进行量化

llama.cpp项目是由Georgi Gerganov创建的一个开源C++软件库,旨在执行各种LLM的推理。它是最流行的量化技术,Hugging Face Hub上有许多量化模型可用。

与依赖CUDA等硬件特定的闭源库的其他库相比,llama.cpp可以运行在更广泛的硬件上,尤其在没有专用硬件的用户中,因其可以在CPU和Android设备上运行,受到了极大的欢迎。此外,llama.cpp还可以将层卸载到GPU,从而加速推理速度。它与不同的推理优化技术兼容,如FlashAttention-2和推测解码。

该项目具有自己独特的量化格式,GGUF,旨在简化并加速模型加载。GGUF文件存储张量和元数据,支持从1位到8位精度的各种格式。它遵循基于位数和特定变体的命名约定,如下所示:

  • IQ1_S 和 IQ1_M:1位精度 – 极低质量
  • IQ2_XXS/XS/S/M 和 Q2_K:2位精度 – 通常低质量,但IQ2适用于大模型
  • IQ3_XXS/XS/S/M 和 Q3_K_S/M/L:3位精度 – 低质量,但适用于大模型
  • IQ4_XS/NL 和 Q4_K_S/M,Q4_0/1:4位精度 – 良好质量,适用于大多数模型
  • Q5_K_S/M 和 Q5_0/1:5位精度 – 高质量
  • Q6_K:6位精度 – 非常高质量
  • Q8_0:8位精度 – 最高质量

为了简要概述GGUF量化,llama.cpp将值分为块并将其四舍五入到较低精度。例如,传统的Q4_0格式处理每个块32个值,根据块中的最大权重值进行缩放和量化()。在Q4_1中,还会加入块中的最小值()。在Q4_K中,权重被划分为超级块,每个超级块包含8个块,每个块有32个值。块的缩放和最小值也会使用更高精度(6位)进行量化。最后,像IQ4_XS这样的i-量化方法受到另一种量化技术QuIP#的启发。该方法确保在八个值的组中有相同数量的正(或负)量化符号,并使用格架存储它们的幅度。

以下是如何将模型量化为GGUF格式的实际示例。以下步骤可以在Google Colab的免费T4 GPU上执行:

  1. 安装llama.cpp和所需的库:

    !git clone https://github.com/ggerganov/llama.cpp
    !cd llama.cpp && git pull && make clean && LLAMA_CUBLAS=1 make
    !pip install -r llama.cpp/requirements.txt
    
  2. 下载要转换的模型。我们将提供Hugging Face Hub上的模型ID,例如mistralai/Mistral-7B-Instruct-v0.2:

    MODEL_ID = "mlabonne/EvolCodeLlama-7b"
    MODEL_NAME = MODEL_ID.split('/')[-1]
    !git lfs install
    !git clone https://huggingface.co/{MODEL_ID}
    
  3. 首先,将模型转换为FP16。这是一个中间文件,将用于每种GGUF量化类型。请注意,llama.cpp中有不同的转换脚本,适用于不同的模型:

    fp16 = f"{MODEL_NAME}/{MODEL_NAME.lower()}.fp16.bin"
    !python llama.cpp/convert.py {MODEL_NAME} --outtype f16 --outfile {fp16}
    
  4. 选择一个格式(例如Q4_K_M)并开始量化。此过程在T4 GPU上可能需要一小时:

    METHOD = "q4_k_m"
    qtype = f"{MODEL_NAME}/{MODEL_NAME.lower()}.{method.upper()}.gguf"
    !./llama.cpp/quantize {fp16} {qtype} {METHOD}
    
  5. 完成后,您的量化模型已准备好。您可以将其下载到本地,或使用以下代码将其上传到Hugging Face Hub:

    from huggingface_hub import create_repo, HfApi
    hf_token = "" # 指定您的令牌
    username = "" # 指定您的用户名
    api = HfApi()
    # 创建空仓库
    create_repo(
        repo_id = f"{username}/{MODEL_NAME}-GGUF",
        repo_type="model",
        exist_ok=True,
        token=hf_token
    )
    # 上传gguf文件
    api.upload_folder(
        folder_path=MODEL_NAME,
        repo_id=f"{username}/{MODEL_NAME}-GGUF",
        allow_patterns=f"*.gguf",
        token=hf_token
    )
    

GGUF模型可以与如llama-cpp-python等后端及LangChain等框架一起使用。如果您想将量化模型集成到更广泛的系统中,这是非常有用的。您还可以通过前端直接与模型进行对话,如llama.cpp的轻量级服务器、LM Studio和Text Generation Web UI。这些工具使得与GGUF模型的交互变得非常容易,提供了类似于ChatGPT的体验。

使用GPTQ和EXL2进行量化

虽然GGUF和llama.cpp提供了CPU推理并支持GPU卸载,但GPTQ和EXL2是专门为GPU设计的量化格式,因此它们在推理时比llama.cpp更快。特别是,EXL2通过其专用库ExLlamaV2提供了最高的吞吐量。

GPTQ和EXL2量化基于GPTQ算法,该算法由Frantar等人(2023年)提出。GPTQ通过改进最佳大脑量化(OBQ)方法来优化LLM的权重量化,旨在高效处理大型矩阵。它从Hessian逆矩阵的Cholesky分解开始,确保数值稳定性。GPTQ不是按照严格顺序量化权重,而是将权重分批处理,迭代更新列和相关的块。这种方法利用懒批量更新,减少了计算冗余和内存瓶颈。

虽然GPTQ仅支持4位精度,但EXL2提供了更多灵活性,具有高度可定制的精度,可以混合不同的量化级别。这允许在2到8位之间的精确比特率,例如2.3、3.5或6.0。它还可以对每个线性层应用多个量化级别,通过对重要权重使用更高位的量化来进行优化。参数是自动选择的,通过多次量化每个矩阵并选择最小化量化误差的组合,同时满足目标比特率。实际上,这使得70B模型能够在单个24 GB GPU上以2.55位精度运行。

推理本身由ExLlamaV2库处理,该库支持GPTQ和EXL2模型。

在以下示例中,我们将使用ExLlamaV2将模型量化为EXL2格式。这些步骤可以在Google Colab的免费T4 GPU上执行:

  1. 从源代码安装ExLlamaV2库:

    !git clone https://github.com/turboderp/exllamav2
    !pip install -e exllamav2
    
  2. 通过从Hugging Face Hub克隆模型的仓库来下载要量化的模型:

    MODEL_ID = "meta-llama/Llama-2-7b-chat-hf"
    MODEL_NAME = MODEL_ID.split('/')[-1]
    !git lfs install
    !git clone https://huggingface.co/{MODEL_ID}
    
  3. 下载用于测量量化误差的校准数据集。在本例中,我们将使用WikiText-103,这是一个标准的校准数据集,包含来自Wikipedia的高质量文章:

    !wget https://huggingface.co/datasets/wikitext/resolve/9a9e482b5987f9d25b3a9b2883fc6cc9fd8071b3/wikitext-103-v1/wikitext-test.parquet
    
  4. 在给定精度下(例如4.5)对模型进行量化:

    !mkdir quant
    !python exllamav2/convert.py \
        -i {MODEL_NAME} \
        -o quant \
        -c wikitext-test.parquet \
        -b 4.5
    

量化后的模型可以像之前一样上传到Hugging Face Hub。

GPTQ和EXL2量化的支持不如GGUF广泛。例如,像LM Studio这样的前端当前不集成它们。您可以使用其他工具,如oobabooga的Text Generation Web UI。它也直接集成在transformers库中,并且得到TGI的支持。GPTQ模型也在TensorRT-LLM中得到支持。

尽管不如GGUF流行,但在Hugging Face Hub上您可以找到许多GPTQ和EXL2模型。

其他量化技术

除了GGUF、GPTQ和EXL2外,还有多种量化技术。本小节将简要介绍激活感知权重量化(AWQ)以及极端量化技术,如QuIP#(量化与不一致性处理)和HQQ(半二次量化)。

由Lin等人(2023年)提出的AWQ是另一种流行的量化算法。它识别并保护最重要的权重,这些权重是根据激活值的大小而非权重大小来确定的。该方法为这些显著权重应用最优的每通道缩放,而无需依赖反向传播或重构,从而确保LLM不会过拟合校准集。尽管它依赖于不同的范式,AWQ与GPTQ和EXL2版本非常相似,尽管稍慢一些。它们在推理引擎中得到了良好的支持,并集成到TGI、vLLM和TensorRT-LLM中。

一个有趣的趋势是将模型量化为1位或2位精度。虽然一些格式,如EXL2,允许极端量化,但模型的质量通常会显著下降。然而,像QuIP#和HQQ这样的最新算法已经针对这一领域进行了优化,提供了能够更好地保持原始模型性能的量化方法。这对于大型模型(超过30B参数)尤其有效,这些模型可能比7B或13B参数的模型占用更少的空间,同时提供更高质量的输出。

预计这一趋势将继续,进一步优化这些量化方法。

本章总结了我们在前几节中介绍的三大推理引擎的特点,见下表:

技术TGIvLLMTensorRT-LLM
连续批处理
推测解码
FlashAttention2
PagedAttention
流水线并行
张量并行
GPTQ
EXL2
AWQ

表8.1 – TGI、vLLM和TensorRT-LLM的功能总结

总结

总之,推理优化是有效部署LLM的关键方面。本章探讨了各种优化技术,包括优化的生成方法、模型并行性和权重量化。通过利用推测解码技术并行预测多个令牌,或使用FlashAttention-2优化的注意力机制,可以显著加速推理。此外,我们讨论了如何通过数据、流水线和张量并行等模型并行方法将计算负载分配到多个GPU,从而提高吞吐量并减少延迟。权重量化,如GGUF和EXL2格式,进一步减少了内存占用并加速推理,尽管可能会有输出质量的折衷。

理解并应用这些优化策略对于在LLM的实际应用中实现高性能至关重要,如聊天机器人和代码补全。技术和工具的选择取决于具体要求,包括可用的硬件、期望的延迟和吞吐量。通过结合不同的方法,如连续批处理、推测解码、先进的注意力机制和模型并行,用户可以定制其部署策略,以最大限度地提高效率。

回顾第四章时,我们仅关注了实现数据摄取管道,这是标准RAG应用的一个组成部分。在下一章中,我们将通过实现检索和生成组件并将它们集成到推理管道中,来完成RAG系统。

参考文献

  • Hugging Face, Text Generation Inference, github.com/huggingface…, 2022.
  • W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C.H. Yu, J.E. Gonzalez, H. Zhang, I. Stoica, Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023.
  • Nvidia, TensorRT-LLM, github.com/NVIDIA/Tens…, 2023.
  • Y. Leviathan, M. Kalman, Y. Matias, Fast Inference from Transformers via Speculative Decoding, 2023.
  • T. Cai, Y. Li, Z. Geng, H. Peng, J.D. Lee, D. Chen, T. Dao, Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, 2024.
  • W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C.H. Yu, J.E. Gonzalez, H. Zhang, I. Stoica, Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023.
  • R.Y. Aminabadi, S. Rajbhandari, M. Zhang, A.A. Awan, C. Li, D. Li, E. Zheng, J. Rasley, S. Smith, O. Ruwase, Y. He, DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022.
  • Y. Huang, Y. Cheng, A. Bapna, O. Firat, M.X. Chen, D. Chen, H. Lee, J. Ngiam, Q.V. Le, Y. Wu, Z. Chen, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019.
  • K. James Reed, PiPPy: Pipeline Parallelism for PyTorch, github.com/pytorch/PiP…, 2022.
  • M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, B. Catanzaro, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020.
  • Verma and Vaidya, Mastering LLM Techniques: Inference Optimization, NVIDIA Developer Technical Blog, developer.nvidia.com/blog/master…, 2023.
  • T. Dettmers, M. Lewis, Y. Belkada, L. Zettlemoyer, LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, 2022.
  • G. Gerganov, llama.cpp, github.com/ggerganov/l…, 2023.
  • E. Frantar, S. Ashkboos, T. Hoefler, D. Alistarh, GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers, 2023.
  • Tuboderp, exllamav2, github.com/turboderp/e…, 2023.
  • J. Lin, J. Tang, H. Tang, S. Yang, W.-M. Chen, W.-C. Wang, G. Xiao, X. Dang, C. Gan, S. Han, AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration, 2024.