清华大学汪玉教授团队力作《高效深度学习》出版:有效降低大模型应用成本
共 14417字,需浏览 29分钟
·
2024-07-03 18:08
--文末赠书--
以 ChatGPT、CoPilot 等应用为代表,大型语言模型的应用已经在各个领域展示出极大的潜力,甚至有望引领人机交互范式的巨大改变。然而,随着大语言模型规模的增长,模型的计算和存储开销与实际可用的硬件资源之间的矛盾日益加剧。在这个背景下,高效深度学习领域显得尤为关键,得到了广泛关注。
本文节选自《高效深度学习:模型压缩与设计(全彩)》一书。
作者是清华大学电子工程系汪玉教授团队,致力于提供通用人工智能算力解决方案。
近期,汪玉团队在“高效大模型”推理方向做了一系列新工作,包括“Skeleton-of-Thought:Large Language Models Can Do Parallel Decoding”的高效算法设计工作、“FlashDecoding++:Faster Large Language Model Inference on GPUs”的 GPU 软件算子设计工作,以及“FlightLLM:Efficient Large Language Model Inference witha Complete Mapping Flow on FPGAs”的 FPGA 硬件加速器设计工作等。
汪玉教授团队在高效深度学习方法论上的认知总结为本书,希望通过本书的出版帮助对这个领域感兴趣的读者,推动领域的发展。
除了 Encoder-Decoder结构,还有两种常用的 Transformer 模型结构,分别是 Encoder-only 结构和 Decoder-only结构。顾名思义,这两种结构分别仅包含编码器和解码器。Encoder-only 结构的代表模型包括BERT、RoBERTa等。2021 年 OpenAI 发布 GPT 系列模型后,Decoder-only 结构逐渐成为主流,其代表模型包括 GPT、OPT、LLaMA 等。下面以常用的 Decoder-only 结构为例,介绍基于 Transformer 的大语言模型的基本组成和推理过程。
Decoder-only 的 Transformer 模型由若干个 Transformer 模块重复堆叠而成,其中每个模块包含多头自注意力(Multi-Head self-Attention,MHA)模块、前馈神经网络(Feed ForwardNetwork,FFN)模块和层归一化操作,如图 10.1 所示。
Transformer 模型的完整推理流程通常包括预处理、模型推理和后处理三个步骤。其中,预处理将用户提供的文本切分成词块序列用于模型推理,后处理则将模型推理得到的词块序列转换成文本反馈给用户。通常,使用预训练的分词器(Tokenizer)完成上述词块序列和文本之间的转换,常见的分词器包括 BPE、SentencePiece等。模型推理时,大语言模型需要进行多次模型前向推理,每次前向推理均基于已有的词块序列生成一个新词块。若每次前向推理时,均将新生成词块拼接在已有词块序列之后,并对整个词块序列进行前向计算,就会产生大量的重复计算,增加模型端到端的推理延时。
为了解决该问题,Pope 等人提出了 Key-Value 缓存技术。该技术的核心思想是,每次进行前向推理时,无须重新计算多头注意力机制中的 Key 和 Value 矩阵,而是可以复用先前的计算结果进行完整的多头自注意力计算。因此,可以将每次前向推理得到的 Key 和 Value 矩阵存储起来,以便在后续的计算中复用。Key-Value 缓存技术用额外存储量换取了计算量的减少,被广泛应用于主流大语言模型推理引擎中。采用 Key-Value 缓存技术之后,可以将模型推理分为以下两个阶段。
(1)预填充阶段(Prefill Stage): 在该阶段中,模型接收初始输入文本并计算得到初始的Key-Value 缓存用以之后的词块生成。
(2)解码阶段(Decoding Stage): 在该阶段中,模型逐一生成词块。产生每个词块时,模型会将初始的输入和已经产生的词块拼接在一起作为输入,并根据模型输出的概率分布采样当前词块。
效率瓶颈
笔者将大语言模型推理时的效率瓶颈归因于以下三个方面。
(1)模型规模大。大语言模型拥有数十亿至数千亿的庞大参数量,例如 Meta 发布的开源大语言模型 LLaMA系列中最大的模型拥有 700 亿个参数。庞大的参数量会带来巨大的存储开销、访存开销及计算开销。
(2)注意力模块开销大。注意力模块是 Transformer 模型的核心组件之一。以 LLaMA-7B 模型为例,当输入序列长度为 8 000 个词块时,预填充阶段的注意力模块的延时占总延时的 71.8%。
(3)解码方式开销大。大语言模型推理过程的解码阶段会依次、逐个地生成新词块。每生成一个新词块都需要进行一次模型前向推理,这意味着模型的所有参数需要完成一次从 GPU显存到 GPU 计算核心上的加载过程。这样的负载计算强度低,在 GPU 平台上运行时访存受限,不能充分利用 GPU 的算力,在批大小(Batch Size)较小时尤为显著。随着解码的进行,词块序列的长度不断增长,导致所需的 Key-Value 缓存空间持续增大。若无合适的存储管理机制,可能导致内存碎片和非连续访存,增加存储开销和访存开销。
优化路径
1. 模型层次优化
模型层次优化通过调整模型实现延时、吞吐率、存储等目标,包括模型结构设计和模型压缩两方面。其中,模型结构设计方法通常需要较大的训练开销,主要包括以下三点。
(1)混合专家动态推理:对于给定的输入,仅激活部分专家模块完成推理,让大参数量的模型能拥有更高的推理效率。
(2)低复杂度注意力机制:降低注意力算子的平方复杂度。
(3)多查询注意力(Multi-Query Attention,MQA)机制:减小注意力算子的存储量和访存量。
2. 系统层次优化
系统层次优化通过设计软硬件系统无损地加速模型推理。针对大语言模型推理的访存受限特性,许多工作通过优化计算图、实现融合算子来降低访存开销。代表性工作为 FlashAt-tention,它通过算子融合和分块计算技术避免了注意力图在显存和高速缓存之间的搬运,大幅降低了推理过程中预填充阶段或训练过程的访存和存储开销。
2023 年,笔者团队提出的FlashDecoding++ 对大语言模型的两类主要计算都进行了进一步优化。对于注意力计算,FlashDecoding++ 基于注意力输入的数值分布特性,引入预设好的最大值,去除了 FlashAt-tention 中的额外更新操作。对于线性层计算,FlashDecoding++ 实现了解码阶段矩阵乘算子的针对性优化,并通过预先建立的映射表,动态选择大语言模型中所有线性层的最优底层实现,实现了端到端推理性能的提升。
服务系统软件(Serving Systems Software)旨在提高系统吞吐率。批次化技术是提升服务器系统软件效率的必备技术,它将多个用户的查询打包成一个批次,以提升硬件利用效率,从而提高整体吞吐率。随后的研究对批次化技术进行了各种改进和变体。模型并行(Model Parallelism)和计算/存储卸载(Offloading)等技术也常被用于提升大语言模型训练或推理的吞吐率。基于给定的计算图和硬件配置,这些技术决定输入计算图上计算和存储的切分、硬件资源的分配,以及时序调度。
此外,高效内存管理技术也是服务系统中重要的组成部分。在硬件加速器设计方面,笔者团队提出了一种面向现场可编程门阵列(Field-Programmable Gate Array,FPGA)的大语言模型轻量化部署流程 FlightLLM,首次在单块 Xilinx U280FPGA 上实现了 LLaMA2-7B 的高效推理。
3. 解码方式优化
在大语言模型的推理过程中,自回归式的解码阶段占据了主要的端到端延时。由于词块之间存在前后依赖关系,多个词块对应的计算无法并行进行,导致 GPU 计算单元的利用率极低。针对这一问题,解决思路可分为以下三类。
(1)猜测解码技术(Speculative Decoding)。这类技术的核心思想是先使用更廉价的方式猜测接下来的候选词块序列,再使用大语言模型并行地验证并接收其中匹配大语言模型验证结果的词块。
(2)非自回归式解码技术(Non-Autoregressive Generation,NAG)。这些工作设计专用解码方式并行采样相邻的若干个词块,且大多需要额外的模型设计或训练。为了保证生成质量,许多非自回归式解码方法会迭代多次修正生成结果。
(3)笔者团队提出了思维骨架(Skeleton-of-Thought,SoT),利用大语言模型的规划和指令跟随能力实现并行生成,从而优化计算利用率和端到端延时。具体地,思维骨架通过提示词引导大语言模型组织答案的提纲或框架,再并行解码提纲中的每个要点。
上述三类工作分别实现了大语言模型的连续并行验证、连续并行采样、分段并行生成,能提升计算利用率,减少端到端生成延时。
模型量化可以从优化问题的视角来看,即在满足硬件约束的前提下,通过优化量化格式、量化参数和量化值,最小化模型的算法性能损失。按工作流程来分类,模型量化可分为训练后量化和量化感知训练两种。
在研究量化之前,首先需要明确两个不同的过程。
1. 离线的模型量化过程
运行量化算法将高位宽、高精度的权重转换为低位宽的权重。若需要对激活值进行量化,并且不是在线统计激活值量化参数,则还需要确定激活值的量化参数。
2. 在线量化推理过程
这是从离线过程得到量化模型,再用量化模型做推理的过程。根据是否使用低精度计算,将在线推理阶段分为以下两类。
(1)使用低精度计算:计算单元采用低精度。一个典型的流程如下:执行两个 INT8 低精度数的低精度乘法后进行累加,累加器一般需要更大位宽(如 INT32),最后还需要对位宽放大的累加输出做再量化,将输出量化回 INT8。
(2)使用高精度计算: 将低位宽表示的权重或激活值去量化为高位宽高精度表示,然后使用高精度计算单元进行计算。最后,可能需要将计算结果量化回低位宽表示,如用于降低训练时激活值的存储和通信开销等。
相比于量化传统小模型,量化大语言模型存在以下新挑战和新需求。
(1)算法性能上:大语言模型权重包含分布范围更广的离群值,且保留这些离群值对模型精度很关键。对离群权重的保留会导致小权重的舍入误差较大。
(2)推理效率上:在不同场景下,大语言模型推理的效率瓶颈不同,需要不同的量化方法。大语言模型在不同阶段、不同批大小、不同输入文本长度的效率瓶颈分析如图 10.2 所示。
在预填充阶段,当输入文本较长或批大小较大时,模型中线性层的计算访存比 I 将大于 I0(I0 表示对应 Roofline 拐点的计算强度),推理瓶颈为计算瓶颈。反之,当输入文本较短或批大小较小时,推理瓶颈为权重访存瓶颈。在解码阶段,仅有批大小会影响效率瓶颈,与预填充阶段类似,当批大小较大时,模型中线性层的实际计算访存比 I 将大于 I0,推理瓶颈为计算瓶颈。反之,当批大小较小时,推理瓶颈为权重访存瓶颈。对于瓶颈为权重访存瓶颈的场景,常用“仅权重量化方法”(Weight-only Quantization)降低权重的访存量,实现推理加速。对于瓶颈为计算瓶颈的场景,由于“仅权重量化方法”不使用低精度计算单元,甚至还会引入额外的去量化计算开销,因而并不能实现推理加速。此时,需依赖使用低精度计算的“权重激活量化方法”(Weight-activation Quantization)加速推理,在这些方法里,权重和激活值都被量化为低位宽和低精度。
研究者针对上述特点设计了适合大语言模型的量化算法,表 10.1 总结了若干个大语言模型量化方法在量化格式、量化参数、量化值等多个维度的策略选取以及量化效果。
从主要目的来看,LLM.int8()、SmoothQuant和 AWQ的目的是保持算法性能。这些工作均针对大语言模型中的权重离群值进行特殊处理,更好地平衡截断和舍入误差。GPTQ的目的是加速离线量化过程,其基于已有的 OBQ方法做改进。LUT-GEMM的目的是加速在线量化推理。FlexGen的目的是降低在线量化推理过程的内存开销。从量化方式来看,GPTQ、LUT-GEMM和 AWQ仅量化权重,计算时会将低精度权重去量化为高精度表示,再和高精度的激活值进行计算;而 LLM.int8()、SmoothQuant同时量化了权重和激活值,使用了高效的低精度计算;除了量化权重,FlexGen还量化了 Key-Value 缓存这一类激活值。下面列出这几个代表工作所使用的方法。
(1)GPTQ:传统量化算法 OBQ基于逐层重建误差最小化原则,对矩阵的每一行采用逐权重量化。具体来说,OBQ 对每一行选取其最优的权重量化顺序,选取方法依赖重建误差相对于当前未被量化权重的二阶梯度信息(Hessian 矩阵)。OBQ 的计算复杂度非常高,原因是在每个权重被量化后,Hessian 矩阵都需要更新。若不同行所选取的量化顺序不同,则每一行量化过程中都要进行大量的 Hessian 矩阵计算,计算开销大。针对该问题,GPTQ 提出了所有行都使用相同的权重量化顺序,即按照统一、从左往右的顺序量化权重矩阵的所有行。此时,不同行在量化过程中所使用的多个 Hessian 矩阵是完全相同的,因此仅需要在一行的量化过程中计算 Hessian 矩阵,即可复用在其他行的量化过程中,实现量化过程的加速。
(2)LUT-GEMM:当前的仅权重量化方法一般会将权重去量化操作和高精度矩阵乘操作融合在一个算子里,以降低访存开销,但是去量化操作仍然引入了额外的计算开销。为降低去量化操作的计算开销,LUT-GEMM 设计了基于查找表(Look-Up-Table,LUT)的去量化操作。这种去量化方法可简单地支持均匀量化和非均匀量化格式的去量化。在数据格式层面,作者使用了一种基于可学习量化间隔的非均匀量化格式。
(3)AWQ:在大语言模型中,不同通道的权重并非同等重要。作者观察到,激活值中存在离群数据的输入通道对应的权重通道更为重要。为了更有效地保存重要通道的权重,作者采用了重参数化技术,并通过网格搜索的方法选择重参数化系数,以最小化重建误差。
(4)LLM.int8():在大语言模型中,激活值的数据范围显著大于权重的数据范围,这使得激活值更难以被量化为低精度。作者发现激活值中的离群数据仅分布在少数通道中。基于这一观察,为了降低激活值的量化误差,作者按照输入通道将激活值和权重拆分为两个部分,包含离群数据的激活值和权重通道以 FP16 格式存储,其他激活值和权重通道量化为 INT8 格式存储。这两部分矩阵分别进行 FP16 和 INT8 的矩阵乘法,将 INT8 数据的计算结果去量化,再将两部分的计算结果相加,得到最终结果。
(5)SmoothQuant:如上所述,大语言模型中的激活值比权重更难量化,并且激活值中离群数据仅分布于少数通道。为了缩小激活值中离群数据通道的数据范围,作者使用重参数化技术扩大权重中对应通道的数据范围,使权重和激活值均易于量化。
(6)FlexGen:在大批次(批大小 ≫ 1)推理时,存储和访存开销的主要来源是模型权重量和 Key-Value 缓存量。为了实现更大的吞吐率,作者通过将权重和 Key-Value 缓存量化为 4-bit 来减小存储和访存开销。
大语言模型是一类能够解决多种任务的通用模型,因此充分了解不同量化方案对任务性能的影响,能够切实地指导实际应用场景中量化方案的选取。
限时优惠,快快抢购吧!
互动有奖
按以下方式与博文菌互动,即有机会获赠图书!
活动方式:在评论区留言参与“你如何看待大模型的加速和压缩”等话题互动,届时会在参与的小伙伴中抽取1名幸运鹅赠送图书盲盒一份!
说明:留言区收到回复“恭喜中奖”者将免费获赠本图书,中奖者请在收到通知的24小时内将您的“姓名+电话+快递地址”留言至原评论下方处即可,隐私信息不会被放出,未在规定时间内回复视作自动放弃兑奖资格。
活动时间:截至7月8日开奖。
快快拉上你的小伙伴参与进来吧~~
温馨提示:可以将“博文视点”设为星标,以免错过赠书活动哦!
发布:刘恩惠
审核:陈歆懿
如果喜欢本文 欢迎 在看丨留言丨分享至朋友圈 三连 < PAST · 往期回顾 >
书单 | 7月新书速递!