Paper Source: A Survey on Efficient Inference for Large Language Models

PRELIMINARIES

Transformer-Styles LLMs

主流LLM是基于Transformer架构设计的，典型的Transformer架构模型由数个堆叠的（stacked）Transformer Block 组成。

Transformer Block：

  • Attention-Linear(generate matrix Q, K, V)
  • Multi-Head Self-Attention（MHSA）
  • Feed Forward Network（FFN）
  • Layer Norm

每个 Transformer Block 接收前一个 Transformer Block 的输出特征，并将其作为输入，并将特征串行送进每个子模块中，最后输出。

特别的，在第一个 Transformer Block 前，需要用一个 tokenizer 将输入语句（prompts）转化为 token 序列，紧接 word embedding 及 positional embedding 层将 token 序列转化为输入特征。

Transformer架构的核心是自注意力机制，这一机制在多头自注意力模块（MHSA）中使用。自注意力机制可以让模型在不考虑距离的情况下识别不同输入部分的重要性,也因此可以得到输入语句的长距离依赖及句子内部的复杂关系。

首先，MHSA模块对输入进行线性变换得到 Q, K, V向量，如(1):

Qi=XWQi,Ki=XWKi,Vi=XWVi(1)Q_i = XW^{Q_i}, K_i = XW^{K_i}, V_i = XW^{V_i} \tag{1}

其中X=[x1,x2,…,xn]X = [x_1,x_2,\dots,x_n]为输入特征, WQi,WKi,WViW^{Q_i},W^{K_i},W^{V_i}为第 ii​ 个 attention head 的变换矩阵

接着, 每个 (Qi,Ki,Vi)(Q_i,K_i,V_i) tuple 都被用于 attention 操作, 得到第 ii 个 attention head 的特征 ZiZ_i, 如(2):

Zi=Attention(Qi,Ki,Vi)=Softmax(QiKiTdk)Vi(2)Z_i = Attention(Q_i,K_i,V_i) = Softmax(\frac{Q_iK_i^T}{\sqrt{d_k}})V_i \tag{2}

其中, dkd_k 为queries(keys)的维度

自注意力计算包含矩阵乘法，计算复杂度为输入长度的二次方

最后, 将所有 attention head 特征进行拼接, 并对它们进行映射矩阵变换, 如(3):

Z=Concat(Z1,Z2,…,Zh)WOZ = Concat(Z_1,Z_2,\dots,Z_h)W^O

其中, WOW^O 为映射矩阵

FFN 为 Transformer Block 的另一个重要模块，其被设置在 MHSA 模块之后，其中包含两个线性变换和一个ReLU激活函数（非线性）。

FFN(X)=W2σ(W1X)FFN(X) = W_2\sigma(W_1X)

Inference Process of LLMs

目前主流 LLM 采用 decoder-only 架构，使用自回归方式生成输出语句，即逐 token 进行输出。在每一次 generation step， 使用过去全部 token 序列（包含原始输入 token 及刚刚生成的 token）作为输入，来生成下一个 token。随着序列长度的增加，生产过程的时间成本会显著增加，为了解决这一问题，一个关键技术， Key-Value（KV） Cache 被提出用于加速生成。

KV Cahce 包含在MHSA模块内，存储和复用前面 token 对应的 key 向量（K）和 value 向量（V）。基于这一项技术，LLM推理过程可以划分为两个阶段：

  1. Pre-fill Phase: 计算并存储原始输入 token 的 KV Cache 并生成第一个输出 token。
  2. Decoding Phase：利用 KV Cache 逐个输出 token，并利用新生成 token 的 KV对 KV Cache进行更新。

Demonstration of the prefilling stage (a) and decoding stage (b)

Illustration of the memory variation through time (latency) during one generation process. Note that we ignore the activation size in this figure for a simplification

上图展示了提升推理效率的关键指标。

横轴 Latency：

  • 在 pre-fill 阶段，将 First Token Latency 记作生成第一个 token 的时间；
  • 在 decoding 阶段，将 Per-output Token Latency 记作生成一个 token 的平均时间；
  • Generation Latency 表示输出整个 token 序列的时间。

纵轴 Memory：

  • Model Size 表示存储模型权重所需要的内存大小；
  • KV Cache Size 代表存储存储 KV Cache 的内存大小；
  • Peak Memory 代表在生成工程中需要占用的最大内存，其大约为 Model Size 与 KV Cache Size 之和，即模型权重和 KV Cache的内存和。

除去延迟和内存中，吞吐量（throughput）也是大模型推理服务系统中的一个广泛使用的指标。token throughput表示每秒生成的token数量，request throughput表示每秒完成的请求数。

Efficiency Analysis: 大模型推理效率瓶颈分析

大语言模型在实际部署应用中，我们通常关注其时延、吞吐、功耗和存储（latency、throughout、energy and power consumption、storage），而在大语言模型推理过程中，有三个重要因素会直接影响上述效率指标，分别是计算开销（Computational Cost）、访存开销（Memory Access Cost）和存储开销（Memory Cost）。进一步地，本综述深入分析探究，并总结归纳除出影响上述指标的三个根本因素，分别为：

  • Model Scale（模型规模）:
    主流大语言模型庞大的模型规模会导致巨大的计算量、访存量和存储量。例如，LLaMA-70B模型包括700亿参数，而GPT-3为1750亿参数。在推理过程中，模型大小对计算成本、内存访问成本和内存使用产生了显著影响。
  • Attention Operation（注意力计算）: Pre-fill 阶段中，自注意操作的计算复杂度为输入长度的2次方
    ，因此随着输入长度的增加，计算成本、内存访问成本和内存使用都会显著增加。
  • Decoding Approach（解码方式）: 主流的自回归解码方式导致极低的计算-访存比和硬件利用率，同时动态增长的KV缓存会导致碎片化的内存使用，对访存开销和存储开销带来增长。

LLM效率优化技术分类法

针对上述影响LLM推理性能的根本因素，文章将LLM效率优化技术划分为三个层次，分别为：

  • Data-level Optimization: 指通过优化输入提示词（Input Compression）或规划模型输出内容（Output Organization）优化推理效率。这类优化技术通常不需要修改模型本身，因此避免了大量的模型训练或微调开销；
  • Model-level Optimization: 指通过设计高效的模型结构（Efficient Structure Design）或模型压缩（Model Compression）技术优化推理效率。前者通常需要预训练或微调来保留或者恢复模型性能，而后者通常会给模型性能带来损失；
  • System-level Optimization: 指通过优化推理引擎（Inference Engine）或服务系统（Serving System）优化推理效率。推理引擎的优化不需要进行模型训练，服务系统的优化对于模型性能而言是无损的。

TAXONOMY

Data-level Optimization 数据层面优化

Data-level Optimization 可以划分为两大类：输入压缩(Input Compression)和输出规划（Output Organization）。

输入压缩技术直接缩短了模型的输入长度来减少推理损失；

输出规划技术通过组织输出内容的结构来实现批量(并行)推理，此方法可以提升硬件利用率和降低模型的生成延迟。

Input Compression

在实际利用大语言模型做回答时，通常会在输入提示词中加入一些辅助内容来增强模型的回答质量，例如，上下文学习技术（In-Context Learning，ICL）提出在输入中加入多个相关的问答例子来教模型如何作答。然而，这些技术不可避免地会增长输入词提示的长度，导致模型推理的开销增大。为了解决该问题，
输入压缩技术通过直接减小输入的长度来优化模型的推理效率。

文章将该类技术进一步划分为四个小类：

  • Prompt Pruning（提示词剪枝）：通常根据设计好的重要度评估指标删除输入提示词中不重要的词块、句子或文段，对被压缩的输入提示词执行压缩。
  • Prompt Summary（提示词总结）：通过对输入提示词做文本总结任务，在保证其语义信息相同地前提下缩短输入的长度。
  • Soft Prompt-based Compression（软提示词压缩）：通过微调训练的方式得到一个长度较短的软提示词，代替原先的输入提示词或其中固定的一部分内容。其中，软提示词指连续的、可学习的词块序列，可以通过训练的方式学习得到。
  • Retrieval-Augmented Generation（RAG，检索增强生成）：通过检索和输入相关的辅助内容，并只将这些相关的内容加入到输入提示词中，来降低原本的输入长度（相比于加入所有辅助内容）。

Input Compression

Output Organization

传统的解码方式是完全串行的，输出规划技术通过组织输出内容的结构，实现（部分的）并行生成来降低端到端的推理延时。

以该领域最早的工作 Skeleton-of-Thought (SoT)（ICLR`2023）为例，SoT技术的核心思想是让大语言模型自行规划输出的并行结构
，并基于该结构进行并行生成，提升硬件利用率，减少端到端生成延时。

具体来说，如下图所示，SoT将大语言模型的生成分为两个阶段：

  • 在第一阶段（The Skeleton Stage），SoT通过设计的提示词让大语言模型输出答案的大纲；
  • 在第二阶段（The Point-expanding Stage），SoT让大语言模型基于大纲中的每一个分点并行做扩展，最后将所有分点扩展的答案合并起来。

SoT

在SoT技术发布后，一些研究工作通过微调大语言模型、前后端协同优化等方式优化输出规划技术，达到了更好的加速比和回答质量之间的权衡。

Model-level Optimization

大模型高效推理的模型级别优化主要集中在模型结构或数据表示的优化上。模型结构优化包括直接设计有效的模型结构、修改原模型和调整推理时间结构。在数据表示优化方面，通常采用模型量化技术。

模型层优化技术可以划分为两大类：高效结构设计（Efficient Structure Design）和模型压缩（Model
Compression）。前者通常需要将新设计的模型从头进行预训练，而后者通常需要微调恢复精度。

Model-level Optimization

Efficient Structure Design

目前主流的大语言模型大多采用Transformer架构，从结构上看，参数量占比最多的前馈神经网络（FFN）和平方复杂度的注意力运算是导致模型推理效率低下的主要原因。

基于此，本文将 Efficient Structure Design 的技术进一步划分为三类：

  • Efficient FFN Design: Mixture-of-Experts (MoE, 混合专家)。通过减少激活FFN的数量来降低内存占用。
    在基于MoE的Transformers中，多个并行FFN（Expert），与可训练的路由模块一起。在推理过程中，面对不同的输入 token 时，路由模块为每个（组）token 激活特定的专家。

    该领域的研究工作主要关注三个方向：

    1. 更高效地获取专家FFN的权重或构建更轻量化的专家FFN；
    2. 优化路由模型使其分配更加平衡，避免因分配不平衡导致的精度和效率下降；
    3. 优化MoE模型的训练方式，使得训练更加稳定。

  • Efficient Attention: 降低 attention 计算的复杂度

    • Across attention heads：

      • Multi/Group-query attention: 在不同 attention head 之间共享 KV Cache来降低访存开销及内存占用

    • Within single attention head：

      • Kernel-based attention：使用核函数，以点乘来近似Softmax运算来降低复杂度。

        Softmax(QKT)V≈ϕ(Q)(ϕ(K)TV)Softmax(QK^T)V \approx \phi(Q)(\phi(K)^TV)

      • Low-rank attention: 在执行 attnetion 计算之前，将 K 和 V 矩阵的 token 维度（如 n)压缩到较小的固定长度（如 k)。该方法基于对 n×nn \times n 矩阵通常表现出低秩特性的观察，使得在 token维度上压缩是可行的。不需要计算完整的 n×nn \times n​ 矩阵，以此降低计算开销。

        X∈Rn×d→X‘∈Rk×dX \in \R^{n \times d} \rightarrow X^` \in \R^{k \times d}

  • Transformer Alternates:

    • State Space Model(SSM)
    • Non-SSM：
      • Long Convolution
      • attention-like Linear Operation

Model Compression

满足一些硬性限制（如时延）的同时最大化降低对模型精度的影响

  • Model Quantization: 模型量化技术通过将模型权重和激活层从 high bit-width 转换为 low bit-width 表示，从而降低LLM显存和计算开销，在 LLM 推理优化领域应用广泛。

    • Post-Training Quantization（PTQ，训练后量化）：指对预训练后的模型直接做量化，不需要重新训练量化后的模型。这类方法对量化后模型的精度损失较大，但因其不需要额外的模型训练开销，因此在大语言模型领域应用最为广泛。
    • Quantization-Aware Training（QAT，量化感知训练）：在模型的训练中加入模型量化过程，并通过训练减小量化带来的精度损失。相比于训练后量化，训练感知量化方法通常具有更高的精度，但其需要大量的数据和计算量来做模型训练。因此，目前该子领域的研究主要关注于如何在数据层面和计算量层面降低模型训练带来的开销。

  • Model Sparsification：模型稀疏分为 Weight Pruning（权重剪枝）和 Sparse attention（注意力稀疏）。

    • Weight Pruning：通过将模型中不重要的权重和对应结构移除，降低模型的计算和存储开销。

      Weight Pruning

    • Sparse Attention：通过减少冗余的注意力计算，来降低预填充阶段的计算开销和解码阶段中KV cache带来存储和访存开销。

  • Structure Optimization: 结构优化，通过修改模型的架构或结构来达到更好的精度-效率之间的权衡。在该领域有两类代表性的技术：

    • Neural Architecture Search：
      神经网络架构搜索，自动化地搜索出最优的模型架构。然而，这类方法目前只在中等规模的语言模型上得到应用，在大语言模型上还未获得实际的优化效果，原因是该技术通常需要在搜索过程中对采样到的架构进行训练评估，对大模型来说需要花费巨大的训练开销。
    • Low Rank Factorization：低秩分解，将一个大的权重矩阵近似分解成两个低质小矩阵的乘积，通过该技术，可以降低大语言模型权重的存储开销和访存开销。

  • Knowledge Distillation: 知识蒸馏，指用一个大模型（教师模型）来辅助训练一个小模型（学生模型），从而将大模型的知识传递给小模型，通过小模型更小的开销达到相似的精度效果。

    • White-box: 可以获得教师模型的架构和权重，因此可以利用更多的信息（例如特征、输出概率等）训练学生模型
    • Black-box: 多针对基于API接口访问的大模型，这类模型的架构和权重无法获取，因此仅能通过构造数据来训练学生模型。

  • Dynamic Inference：动态推理，在推理过程中，基于输入数据的不同，动态决定最合适的模型子结构。涉及到 early exiting 技术。