大模型推理优化方法整理

优化技术 优化原理
硬件升级 硬件的性能和特性对于模型推理速度和效率有着重要影响，通过硬件升级，大语言模型可以实现更高的推理速度。
子图融合技术 子图融合通过将计算图中的一系列操作合并为更大的操作，减少计算和内存开销，以提高推理性能。它通过优化计算流程来加速模型推理。
模型压缩技术 模型压缩技术包括稀疏化、参数剪枝、量化、知识蒸馏等方法，通过减小模型的尺寸和计算量，提高模型在推理阶段的速度和资源效率。
并行化技术 并行化技术包括数据并行化、模型并行化、管道并行等方法，利用多核 CPU、多个 GPU 或其他硬件资源，同时处理多个数据样本或模型部分，以提高推理效率。
Transformer 结构优化 Transformer 结构优化通过对 Transformer 模型的各个组件（如注意力机制、前馈网络等）进行优化，减少计算量、参数数量和内存开销，以加速推理。
动态批处理 动态批处理技术根据输入数据的长度和特性自适应地调整批处理大小，以提高计算资源的利用率，适应不同输入情况，从而提高推理性能。

硬件升级

NVIDIA H100 PCIe

相比A100，H100在训练和推理上都有了很大的提升，具体对比如下图所示：

H100.png

子图融合（subgraph fusion）

Subgraph Fusion（子图融合）是一种大语言模型推理加速技术。在深度学习中，模型通常以计算图的形式表示，其中每个节点表示一个操作（比如加法、乘法、激活函数等），边表示数据流向。在推理阶段，为了执行模型，计算图中的各个节点会逐一执行，导致计算和内存开销较大。因此，子图融合技术主要针对深度学习模型的计算图进行优化，其目标是将计算图中的一系列连续操作（子图）合并为更大的操作，从而减少多余的计算和数据传输。通过将多个操作融合为一个更大的操作，可以减少操作之间的开销，同时也有助于提高并行性。

FasterTransformer

DeepSpeed

MLC LLM

模型压缩（Model Compression）

在大型语言模型的推理阶段，模型压缩技术是一种关键的优化手段，旨在减小模型的尺寸和计算量，从而提高推理速度和降低资源消耗。这类方法可以在精度损失很小的情况下实现模型小型化，主要包括3类手段：稀疏、量化、蒸馏。

稀疏(Sparsity)

量化(Quantization)

蒸馏(Distillation)

并行化（Parallelism）

在提及推理的并行化技术时，往往会提到3D Parallelism，它是指在三个维度上实现并行化的策略。这三个维度分别是：数据并行，张量并行和流水线并行。这种并行化方法旨在充分利用硬件资源，加速模型的推理过程，从而提高系统的性能和效率。

3Dparallelism.png

数据并行 (Data Parallelism, DP)

张量并行(Tensor Parallelism, TP)

流水线并行(Pipeline Parallelism, PP)

Transformer 结构优化

该类方法主要通过优化 Transformer 的结构以实现推理性能的提升。

FlashAttention

FlashAttention的作者们发现，这些Efficient Transformer虽然能够有效降低模型的FLOPS，但它们的计算速度并没有显著降低。导致该现象的根本原因是模型的计算速度除了与FLOPS有很大关系，同时也与MAC（Memory Access Cost，存储访问开销）有关。尤其是当计算本身已经很高效的情况下，MAC的开销更加不能忽略。MAC的开销主要来自两方面。一是从存储中读取数据；二是向存储中写数据。与CPU的情况类似，在GPU中，当需要计算时，需将数据从显存中读取并由计算单元进行计算操作。在计算完毕后，再写回到显存中。(这与ShuffleNet v2 十分类似。当时很多轻量化模型涌现，也是只关注FLOPS。ShuffleNet v2中却强调了MAC，这与FlashAttention背后的研究动机异曲同工。）

我们可以根据计算的密集程度，将操作分为两类：计算密集型和存储访问密集型。前者耗时主要在于计算本身，对显存的读写几乎可以忽略。例如大矩阵乘法、大channel size的卷积操作等。对于这类操作，它们的FLOPS决定了计算的时耗。后者耗时主要集中在存储的访问上，计算本身耗时较低。例如逐元素操作（ReLU，Dropout等）、以及Reduce操作（求和、softmax、BatchNorm等）。对于这类操作，它们的MAC决定了计算的时耗。

对于一个transformer而言，主要有两个部分需要关注：Attention以及FFN，后者是一个计算密集型模块，在这篇工作中没有关注，前者中有一些操作是存储访问密集型的，因此这篇工作主要关注Attention模块。

标准Attention

首先我们回顾一下标准 Attention 的操作:
$$
\text{Attention}(Q,K,V)=\text{softmax}\biggl(\frac{QK^\top}{\sqrt{d_k}}\biggr)V
$$
其中$Q,K,V\in\mathbb{R}^{N\times d}(N\text{表示序列长度，}d\text{表示维度})$，上述公式可拆解为：
$$
\mathbf{S}=\mathbf{Q}\mathbf{K}^\top\in\mathbb{R}^{N\times N},\quad\mathbf{P}=\mathrm{softmax}(\mathbf{S})\in\mathbb{R}^{N\times N},\quad\mathbf{O}=\mathbf{P}\mathbf{V}\in\mathbb{R}^{N\times d}
$$
其中$\mathbf{S},\mathbf{P}$（对于decoder来说还有 mask）的空间复杂度都是$O(N^2)$，另外还有几个存储访问密集型操作：对$\mathbf{S}$的 scale, mask 和 softmax 操作，对$\mathbf{P}$的 dropout 操作。下图算法展示了 HBM 与 SRAM 之间的数据传输过程:

FlashAttention.png

以 A100 (40GB HBM) 为例，内存层次结构分为三层：GPU SRAM（19 TB/s，20 MB），GPU HBM（1.5 TB/s，40 GB），CPU DRAM（12.8 GB/s，>1 TB）。SRAM内存分布在108个流式多处理器(SMs)上，每个处理器192KB。片上SRAM比HBM快得多，但比HBM小得多，在计算方面，使用Tensor Core的BFLOAT16 的理论峰值吞吐量为 312 TFLOPS。GPU 的典型操作方式是使用大量的线程来执行一个操作，这个操作被称为内核。输入从HBM加载到寄存器和SRAM，并在计算后写回HBM。

不难发现，大矩阵对HBM进行重复读写是一个主要瓶颈。为此 FlashAttention 提出了两种方法来分别解决上述问题：tiling 和 recomputation：

  • tiling：注意力计算被重新构造，将输入分割成块，并通过在输入块上进行多次传递来递增地执行softmax操作。
  • recomputation：存储来自前向的 softmax 归一化因子，以便在反向中快速重新计算芯片上的 attention，这比从HBM读取中间矩阵的标准注意力方法更快。

tiling技术

之前提到了FlashAttention的核心是对Self-Attention的计算进行分块计算。对于矩阵乘法而言，可以直接通过分块来达到分块计算的目的。但Self-Attention中有 softmax操作，而 softmax 的分母包含与所有元素相关的求和项，所以对Self-Attention进行分块计算的真正难点在于对 softmax 的分块计算。

假设有维度为n的向量$x\in \mathbb{R}^n$我们知道softmax的计算公式如下：
$$
\text{softmax}(x_{i})={\frac{e^{x_{i}}}{\sum_{j=1}^{n}e^{x_{j}}}}
$$
但是由于计算公式中包含指数项，在x较大的时候有溢出的可能，因此大多数框架采用稳定版的softmax，如下所示：
$$
\begin{aligned}
m(x)&=\text{max}(x_1,x_2,\ldots,x_n)\\
f(x)&=[e^{x_1-m(x)},\ldots,e^{x_n-m(x)}]\\
l(x)&=\sum_{i=1}^{n} f(x_i)\\
\text{softmax}(x_i)&=\frac{f(x_i)}{l(x)}
\end{aligned}
$$
值得注意的是上面这个公式的分子是向量，分母是标量，这是一种逐元素的除法，是存储访问密集型操作。接下来我们介绍如何对softmax进行分块计算。

对于一个维度为 2n 的向量，将其一分为二分块，即：$x=[x^{(1)},x^{(2)}]$。我们先计算前半部分的局部softmax，当有：
$$
\begin{aligned}
m(x^{(1)})&=\text{max}(x^{(1)}_1,x^{(1)}_2,\ldots,x^{(1)}_n)\\
f(x^{(1)})&=[e^{x^{(1)}_1-m(x^{(1)})},\ldots,e^{x^{(1)}_n-m(x^{(1)})}]\\
l(x^{(1)})&=\sum_{i=1}^{n} f(x^{(1)}_i)\\
\text{softmax}(x^{(1)}_i)&=\frac{f(x^{(1)}_i)}{l(x^{(1)})}
\end{aligned}
$$

显然，这个局部softmax与最终的softmax不等，因此我们需要保存一些额外变量，在处理完后半部分之后更新这个局部的softmax值。我们需要保存的变量有：$m(x^{(1)})，l(x^{(1)})，m_{max}，l_{all}$，前两者如公式所写，这两个标量的保存要比保存整个向量的开销小得多，后两者分别表示当前最大值和当前全局指数求和项。

接下来，用同样的方式处理向量$x^{(2)}$：
$$
\begin{aligned}
m(x^{(2)})&=\text{max}(x^{(2)}_1,x^{(2)}_2,\ldots,x^{(2)}_n)\\
f(x^{(2)})&=[e^{x^{(2)}_1-m(x^{(2)})},\ldots,e^{x^{(2)}_n-m(x^{(2)})}]\\
l(x^{(2)})&=\sum_{i=1}^{n} f(x^{(2)}_i)\\
\text{softmax}(x^{(2)}_i)&=\frac{f(x^{(2)}_i)}{l(x^{(2)})}
\end{aligned}
$$
然后用新得到的变量更新之前的变量：
$$
\begin{aligned}
m_{max}^{new}&=max(m_{max},m(x^{(2)}))\\
l_{all}^{new}&=e^{m_{max}-m_{max}^{new}}l_{all}+e^{m_{x^{(2)}}-m_{max}^{new}}l(x^{(2)})
\end{aligned}
$$
这两个公式也很简单，第一个无需赘述，第二个变量中之所以是局部而非全局只是因为除以的最大值是局部最大值，因此将其乘以局部最大值的指数再除以新的全局最大值的指数即可。这时，如果想将局部softmax变为全局softmax，当有：
$$
\begin{aligned}
f^{new}(x^{(2)})& =f(x^{(2)})\cdot e^{m(x^{(2)})-m_{max}^{new}} \\
&=[e^{x_{1}^{(2)}-m(x^{(2)})},\ldots,e^{x_{B}^{(2)}-m(x^{(2)})}]\cdot e^{m(x^{(2)})-m_{max}^{new}} \\
&=[e^{x_{1}^{(2)}-m_{max}^{new}},\ldots,e^{x_{B}^{(2)}-m_{max}^{new}}]\\
l^{new}(x^{(2)})&=l_{all}^{new}
\end{aligned}
$$

基于此二式便能得到两个新的softmax：
$$
\begin{aligned}
\mathrm{softmax}^{(new)}(x^{(1)})&=\frac{\mathrm{softmax}(x^{(1)})\cdot l(x^{(1)})\cdot e^{m(x^{(1)})-m_{max}^{new}}}{l_{all}^{mew}}\\
\mathrm{softmax}^{(new)}(x^{(2)})&=\frac{\mathrm{softmax}(x^{(2)})\cdot l(x^{(2)})\cdot e^{m(x^{(2)})-m_{max}^{new}}}{l_{all}^{mew}}
\end{aligned}
$$

前向传播

在了解了tiling技术后，我们开始了解完整的前向传播过程，如下图所示：

feedforward.png

其中第2行是要根据SRAM的大小M以及向量的维度d选择合适的分块大小$B_c,B_r$。

第4、5行是要将矩阵$\mathrm {Q，O}$沿着行方向分为$T_r$块，每一分块的大小为$B_r \times d$；将向量$l$ 和向量$m$分为$T_r$块，每一个子向量大小为$B_r$。将矩阵$\mathrm {K，V}$沿着行方向分为$T_c$块，每一块的大小为$B_c \times d$。如下图所示：

divide.png

接下来是一个两层的循环，先遍历$\mathrm {K，V}$，再遍历$\mathrm {Q，O，l，m}$，然后按照tiling策略操作即可获得最终结果，至此，前向传播过程结束。

recomputation和反向传播

在标准Attention中，有如下公式：
$$
\mathbf{S}=\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}},\mathbf{P}=\text{Softmax(}\mathbf{S}),\mathbf{O}=\mathbf{P}\mathbf{V}
$$
根据矩阵求导法则，当有：
$$
\text{d}\mathbf{V}=\mathbf{P}^\top\text{d}\mathbf{O}
$$
其余的变量同理。

但是FlashAttention将矩阵分成了若干个块，因此可以使用tiling策略。此外，由于$\mathbf{P}$的大小是$O(N^2)$量级，为了减小访存代价，FlashAttention选择在反向传播时重新计算，而非在前向传播中保存此矩阵（结果确实比取出来更快）。

MAC分析

接下来我们来分析两种方法的访存次数，以下图为例：

FlashAttention.png

在不使用FlashAttention时，对于上图的1、2、3三行，当有：

  • 第一行，读 $Q,K$ 的MAC次数为 $2Nd$ ，写 $S$ 的MAC次数为 $N^2$ 。
  • 第二行，读 $S$ 的MAC次数为 $N^2$ ，写 $P$ 的MAC次数为 $N^2$ 。
  • 第三行，读 $P$ 的MAC次数为 $N^2$ ，读 $V$ 的MAC次数为 $Nd$ ，写 $O$ 的MAC次数为 $Nd$ 。

共计$O(Nd+N^2)$。

在使用FlashAttention时，对于上图当有：

  • 对于内循环，需要读取完整的$Q$，开销为$Nd$。
  • 对于外循环，需要执行$T_c=\lceil \frac{4dN}{M} \rceil$次。

共计$O(N^2d^2M^{-1})$。由于M通常远大于d，因此FlashAttention会更快。

结论

FlashAttention基本可以达到2-4倍的计算加速，10-20倍的内存节省，已被广泛使用。

FlashAttentionV2

尽管 FlashAttention 的速度已经达到baseline的 2-4 倍，但它仍然有相当大的改进空间。FlashAttention 仍然不如优化过的矩阵乘法 (GEMM) 运算快，仅达到理论最大 FLOPs/s 的 25-40%。

因此研究者们又推出了 FlashAttentionV2 。FlashAttentionV2 的速度是 FlashAttention 的 2 倍，在 A100 GPU 上达到 230 TFLOPs/s。在端到端训练 GPT 类语言模型时，FlashAttentionV2 可让训练速度高达 225 TFLOPs/s（模型 FLOP 利用率为 72%）。

本文主要贡献和创新点为：

  1. 减少了非矩阵乘法 FLOPs的数量（消除了原先频繁rescale）。虽然非矩阵乘法 FLOPs仅占总FLOPs的一小部分，但它们的执行时间较长，这是因为GPU有专用的矩阵乘法计算单元，其吞吐量高达非矩阵乘法吞吐量的16倍。因此，减少非矩阵乘法 FLOPs并尽可能多地执行matmul FLOPs非常重要。
  2. 提出了在序列长度维度上并行化。该方法在输入序列很长（此时batch size通常很小）的情况下增加了GPU利用率。即使对于单个head，也在不同的thread block之间进行并行计算。
  3. 在一个attention计算块内，将工作分配在一个thread block的不同warp上，以减少通信和共享内存读/写。

GPU基础知识

由于FlashAttentionV2的优化偏底层一些，因此我们需要先了解一些关于GPU的基础知识。GPU在软件和硬件两个方面由低到高分别分为如下层次：

GPU.png

  1. 从硬件角度来看：

    • Streaming Processor（SP）：是最基本的处理单元，也就是 CUDA core。
    • Streaming MultiProcessor（SM）：一个SM由多个CUDA core（SP）组成，每个SM在不同GPU架构上有不同数量的CUDA core，例如Pascal架构中一个SM有128个CUDA core。

    SM还包括特殊运算单元，共享内存，寄存器文件和调度器等。其中寄存器和共享内存是稀缺资源，这些有限的资源就使得每个SM中active warps有非常严格的限制，也就限制了并行能力。

  2. 从软件角度来看：

    • thread：thread是最小的执行单元，一个CUDA并行程序由多个thread来执行。
    • warp（线程束）：一个warp是在一个时钟周期内一起执行的一组32个线程。每个warp中的thread可以同时执行相同的指令，从而实现SIMT（单指令多线程）并行。warp是SM中最小的调度单位，一个SM可以同时处理多个warp。
    • thread block：一个thread block可以包含多个warp，同一个block中的thread可以同步，也可以通过shared memory进行通信。一个warp中的threads必然在同一个block中，如果block所含thread数量不是warp大小的整数倍，那么多出的那个warp中会剩余一些inactive的thread。也就是说，即使warp的thread数量不足，硬件也会为warp凑足thread，只不过这些thread是inactive状态，但也会消耗SM资源。
    • grid：在GPU编程中，grid是一个由多个thread block组成的二维或三维数组。grid的大小取决于计算任务的规模和thread block的大小，通常根据计算任务的特点和GPU性能来进行调整。在执行过程中，网格中的线程块会被分配到不同的SM上，以实现并行计算。

Hardware和Software的联系：

在执行任务时，grid中的thread block被分配到SM上，大量的thread可能被分到不同的SM上，但是一个线程块的thread只能在一个SM上调度，SM一般可以调度多个block。每个thread拥有自己的程序计数器和状态寄存器，并且可以使用不同的数据来执行指令，从而实现并行计算，这就是所谓的Single Instruction Multiple Thread。

一个CUDA core可以执行一个thread，一个SM中的CUDA core会被分成几个warp，由warp scheduler负责调度。GPU规定warp中所有thread在同一周期执行相同的指令，尽管这些thread执行同一程序地址，但可能产生不同的行为，比如分支结构。一个SM同时并发的warp是有限的，由于资源限制，SM要为每个block分配共享内存，也要为每个warp中的thread分配独立的寄存器，所以SM的配置会影响其所支持的block和warp并发数量。

GPU执行模型小结：

GPU有大量的threads用于执行操作（an operation，也称为a kernel）。这些thread组成了thread block，接着这些blocks被调度在SMs上运行。在每个thread block中，threads被组成了warps（32个threads为一组）。一个warp内的threads可以通过快速shuffle指令进行通信或者合作执行矩阵乘法。在每个thread block内部，warps可以通过读取/写入共享内存进行通信。每个kernel从HBM加载数据到寄存器和SRAM中，进行计算，最后将结果写回HBM中。

算法——减少了非矩阵乘法运算

因为现代 GPU 具有专门的计算单元（例如 Nvidia GPU 上的张量核心），使得矩阵乘法速度更快，可以达到非矩阵乘法的16倍速度。为了保持高吞吐量，研究者希望在矩阵乘法 FLOP 上花费尽可能多的时间。因此他们调整了 FlashAttention 中使用的在线 softmax 技巧，以减少重新缩放操作、边界检查和因果掩码操作的数量，而无需更改输出。

并行——优化thread block

FlashAttention在batch和heads两个维度上进行了并行化：使用一个thread block来处理一个attention head，总共需要thread block的数量等于batch size × number of heads。每个block被调到到一个SM上运行，例如A100 GPU上有108个SMs。当block数量很大时（例如≥80），这种调度方式是高效的，因为几乎可以有效利用GPU上所有计算资源。但是在处理长序列输入时，由于内存限制，通常会减小batch size和head数量，这样并行化程度就降低了。因此，FlashAttention-2还在序列长度这一维度上进行并行化，显著提升了计算速度并提高了GPU的占用率。

分区——优化warp

对于每个块，FlashAttention 将 K 和 V 分割到 4 个 warp 上，同时保持 Q 可被所有 warp 访问。这种方案是低效的，原因在于所有 warp 都需要将它们的中间结果写入共享内存，并同步，然后将中间结果相加。这些共享内存读写会减慢 FlashAttention 中的前向传递速度。在 FlashAttention-2 中，研究者将 Q 分割在 4 个 warp 上，同时保持 K 和 V 可被所有的 warp 访问。每个 warp 执行矩阵乘法以获得 Q K^T 的切片，然后只需与 V 的共享切片相乘就能获得相应的输出切片。这种新的方法减少了不同 warp 之间的同步和通信量，也减少了共享内存的读写。

flashattention2.png

新特性

FlashAttention 仅支持最高 128 的头维数，这对一些模型而言并不够，因此，FlashAttention-2 支持了高达 256 的头维数，这意味着 GPT-J、CodeGen 和 CodeGen2、StableDiffusion 1.x 等模型可以使用 FlashAttention-2 来获得加速和节省内存。此外，FlashAttention-2 还支持了多查询注意力（multi-query attention, MQA）以及分组查询注意力（grouped-query attention, GQA）。它们是注意力的变体，其中多个查询头关注相同的键和值头，以减少推理过程中 KV Cache的大小，并可以显著提高推理吞吐量。

PagedAttention

vLLM 是在加州大学伯克利分校开发，配备了PagedAttention的vLLM重新定义了 LLM 服务的最新技术水平：它的吞吐量比 HuggingFace Transformers 高出 24 倍，且无需更改任何模型架构。

现有问题

众所周知的，当前注意力算法在自回归解码过程中，需要将所有输入令牌的注意力键和值张量存储在GPU内存中，以生成下一个令牌，这也就是我们所说的 KV Cache，这个被存储在HBM中的缓存有以下三种缺陷：

  1. 显存占用大：在 LLaMA-13B 中，缓存单个序列最多需要 1.7GB 显存。
  2. 动态变化：KV 缓存的大小取决于序列长度，这是高度可变和不可预测的。因此，这对有效管理 KV cache 挑战较大。该研究发现，由于碎片化和过度保留，现有系统浪费了 60% - 80% 的显存。
  3. 读写低效：根据前面对于GPU存储层次的介绍我们知道，HBM虽然具有比较大的存储容量，但是读写速度只有1.5T/S，跟SRAM有着数量级的差距，因此频繁地对其进行读写势必会拖慢计算速度。

因此，有效管理KV缓存是一个重大挑战。对此，研究团队发现现有系统由于碎片化和过度保留而浪费了60%至80%的内存。

解决方案

为了上述问题，PagedAttention应运而生，这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在非连续的内存空间中存储连续的 key 和 value 。具体来说，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。下面是一个KV Cache在物理空间的存储示例：

annimation0.gif

因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理 key 和 value ，就像在操作系统的虚拟内存中一样：可以将块视为页面，将 token 视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块中。而物理块则在生成新 token 时按需分配，下面是一个为新token分配新块并更新块表的示例：

annimation1.gif

在 PagedAttention 中，内存浪费只会发生在序列的最后一个块中。这使得在实践中可以实现接近最佳的内存使用，仅浪费不到 4%。这种内存效率的提升被证明非常有用，允许系统将更多序列进行批处理，提高 GPU 使用率，显著提升吞吐量。

PagedAttention 还有另一个关键优势 —— 高效的内存共享。例如在并行采样中，多个输出序列是由同一个 prompt 生成的。在这种情况下，prompt 的计算和内存可以在输出序列中共享。下面是一个并行采样的示例：

annimation2.gif

PagedAttention 自然地通过其块表来启动内存共享。与进程共享物理页面的方式类似，PagedAttention 中的不同序列可以通过将它们的逻辑块映射到同一个物理块的方式来共享块。为了确保安全共享，PagedAttention 会对物理块的引用计数进行跟踪，并实现写时复制（Copy-on-Write）机制。下面是一个从同一输入并行生成多个输出的示例：

annimation3.gif

PageAttention 的内存共享大大减少了复杂采样算法的内存开销，例如并行采样和集束搜索的内存使用量降低了 55%。这可以转化为高达 2.2 倍的吞吐量提升。在实验中，基于pagedAttention的vLLM实现了huggingface 24 倍的吞吐量，TGI 2.5 倍的吞吐量。

vllm.png

补充：关于KV Cache

生成式模型的推理过程很有特点，我们给定一个输入文本，模型会输出一个回答（长度为N），其实该过程中执行了N次推理过程。即一次推理只输出一个token，输出token会与输入tokens 拼接在一起，然后作为下一次推理的输入，这样不断反复直到遇到终止符。如下代码及图所示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
from transformers import LlamaForCausalLM, LlamaTokenizer
tokenizer = LlamaTokenizer.from_pretrained("/data1/wt/Pretrained_Model/llama2_hf")
model = LlamaForCausalLM.from_pretrained("/data1/wt/Pretrained_Model/llama2_hf").eval()
in_text = "Lionel Messi is a"
in_tokens = torch.tensor(tokenizer.encode(in_text))
token_eos = torch.tensor([198])
out_token = None
i = 0
with torch.no_grad():
    while out_token != token_eos:
        in_tokens = in_tokens.unsqueeze(0)
        res = model(in_tokens)
        out_token = torch.argmax(res["logits"][:, -1, :])
        in_tokens = torch.cat((in_tokens, out_token.unsqueeze(0).unsqueeze(0)), -1)
        in_tokens = in_tokens.squeeze()
        text = tokenizer.decode(in_tokens)
        print(f"step {i} input: {text}", flush=True)
        i += 1
out_text = tokenizer.decode(in_tokens)
print(f" Input: {in_text}")
print(f"Output: {out_text}")

clm.png

在上面的推理过程中，每 step 内，输入一个 token序列，经过Embedding层将输入token序列变为一个三维张量[b, s, h]，经过一通计算，最后经logits层将计算结果映射至词表空间，输出张量维度为[b, s, vocab_size]。当前轮输出token与输入tokens拼接，并作为下一轮的输入tokens，反复多次。可以看出第 i+1 轮输入数据只比第 i 轮输入数据新增了一个token，其他全部相同！因此第 i+1 轮推理时必然包含了第 i 轮的部分计算。KV Cache的出发点就在这里，缓存当前轮可重复利用的计算结果，下一轮计算时直接读取缓存结果，就是这么简单，不存在操作系统中需要考虑的 Cache miss 问题。

为了实现 KV Cache，我们对代码做如下改动：

  1. 在推理时新增了 past_key_values 参数，该参数就会以追加方式保存每一轮的K V值。KV Cache变量内容为((k,v), (k,v), …, (k,v))，即有 layers 个 k,v 组成的一个元组。
  2. 推理输出的token直接作为下一轮的输入，不再拼接，因为上文信息已经在 kvcache 中。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import torch
from transformers import LlamaForCausalLM, LlamaTokenizer
tokenizer = LlamaTokenizer.from_pretrained("/data1/wt/Pretrained_Model/llama2_hf")
model = LlamaForCausalLM.from_pretrained("/data1/wt/Pretrained_Model/llama2_hf").eval()
in_text = "Lionel Messi is a"
in_tokens = torch.tensor(tokenizer.encode(in_text))
token_eos = torch.tensor([2])
out_token = None
kvcache = None
out_text = in_text
i = 0
with torch.no_grad():
    while out_token != token_eos:
        in_tokens = in_tokens.unsqueeze(0)
        res = model(in_tokens, past_key_values=kvcache)
        kvcache = res["past_key_values"]
        out_token = torch.argmax(res["logits"][:, -1, :])
        in_tokens = out_token
        in_tokens = in_tokens.squeeze()
        text = tokenizer.decode(in_tokens)
        in_tokens = in_tokens.unsqueeze(0)
        print(f"step {i} input: {text}", flush=True)
        i += 1
        out_text += text
out_text = tokenizer.decode(in_tokens)
print(f" Input: {in_text}")
print(f"Output: {out_text}")

clm2.png

当时候huggingface的Transformer时，KV Cache 配置开启后，推理过程可以分为2个阶段：

  1. 预填充阶段：发生在计算第一个输出token过程中，这时Cache是空的，计算时需要为每个 transformer layer 计算并保存 key cache 和 value cache ，在输出token时Cache完成填充；FLOPs同KV Cache关闭一致，存在大量 gemm 操作，推理速度慢。
  2. 使用KV Cache阶段：发生在计算第二个输出token至最后一个token过程中，这时Cache是有值的，每轮推理只需读取Cache，同时将当前轮计算出的新的Key、Value追加写入至Cache；FLOPs降低，gemm变为gemv操作，推理速度相对第一阶段变快，这时属于Memory-bound类型计算。

借助KV Cache的推理流程如下图所示（关于这部分的理解参考FlexGen论文的第三部分更好）：

KVCache.png

动态批处理（Dynamic Batch, Continuous batch）

在大语言模型的推理优化技术中，动态批处理技术是一种针对不同输入样本长度进行自适应调整批处理大小的策略。大语言模型通常会对输入文本进行分词，生成词嵌入或子词嵌入。由于不同的文本长度可能会导致不同大小的嵌入矩阵，传统的固定批处理大小可能会导致在处理不同长度的输入时产生性能瓶颈。动态批处理技术通过在推理过程中根据当前输入的长度自适应地调整批处理大小，以充分利用计算资源并提高推理效率。

ORCA

FastServe

vLLM

Text Generation Inference

总结

Inference优化目标

Inference服务关注两个指标：Latency和Throughput。Latency关注服务体验，也就是返回结果要快，用户体验好。Throughput则是关注系统成本，高Throughput则系统单位时间处理的量就大，系统利用率高，但是会影响latency。这两个指标一般情况下需要trade-off。

  • Latency：延时，主要从用户的视角来看，也就是用户提交一个prompt，然后得到response的时间。特殊情况batch_size=1只给一个用户进行服务，Latency是最低的。计算方法为生成一个token所需要的单位时间数，如16 ms/token。
  • Throughput：吞吐率，主要是从系统的角度来看，单位时间内能处理的tokens数，如16 tokens/sec。扩大Throughput的方法一般就是提升Batch_size，也就是将一个一个用户的请求由之前的串行改为并行。

高并发时，把用户的prompt合在扩大batch_size能提升Throughput，但会一定程度上损害每个用户的latency，因为以前只计算一个请求，现在合并计算多个请求，每个用户等待的时间就长了。从我们实际的测试结果可以看到，Throuput随着batch_size的增大而增大，但是latency是随着减小的，当然Latency在可接受范围内就是ok的。因此指标需要trade-off。

简单计算，对于一次请求来说：latency=batch_size * output sequence length / Throughput. 提升batch_size会提升Throughput，但Throughput与batch_size并不是同比例增大的，因此导致Latency随着batch_size增大而增大。

Inference过程

有两个阶段Prefill Phase和Decoding Phase（ FlexGen中讲的比较清楚）。

  1. Prefill Phase：称为预处理/Encoding。计算并缓存每一层的key和value，其他的不需要缓存。每一个请求的prompt需要经过这个阶段，它只计算一次，是并行计算的。这个缓存称为KV Cache，KV Cache是整个解码过程中最为核心关键的一块。
  2. Decoding Phase：生成新token阶段，它是串行的，也就是decode one by one。它用上一步生成的token，称为当前token放到input中，然后生成下一个token。具体包括两步，一是Lookup KV Cache计算并输出当前token最终embedding用来预测下一个token，二是缓存计算过程中得到的当前token在每一层的key和value，update到第一阶段Prefill Phase中的KV Cache中。

这样划分的好处是，每次解码，不需要对整个网络全部进行计算，相当于计算复杂度从O(n^3)变为O(n)了。Prefill Phase只需要计算一次，核心就是GPT-style的transformer是单向的，而不是双向的。每一个token只与它前面的tokens有关系，其key，value也只与它前面的tokens有关系，因此可以只计算一次且可以并行。后面decoding phase就很快了，避免了重复计算。整个的开销就只有key-value cache的update以及look-up的时间。

有文章里面提出LLM inference is memory-IO bound, not compute bound。从下面的量化分析来看确实如此。

Inference量化分析

Inference的核心是KV Cache，以FP16为例，其对显存的占用量化分析如下。

其对显存占用分为两块，一是Weights、二是KV Cache。

Weights占用大约为layer_num * ( 8 * hidden_size * hidden_size + 4 * hidden_size * MLP_hidden_size )。

KV Cache的占用为4 * batch_size * layer_num * hidden_size * ( input_length + output_length )。

以OPT-175B 为例（layer_num = 96, hidden_size = 12288, MLP_hidden_size = 49152，batch_size=512，input_length=512, output length=32)。

Weights差不多占用 325G, KV cache 差不多占用 1.2T。对内存消耗是非常惊人。里面唯一可以调节的参数是batch_size，显存占用基本与batch_size呈正比关系。显存的大小限制了batch_size，而batch_size的大小就限制了Throughput。因此就有很多加速工作就是想办法节省显存，进而扩大batch_size。

加速优化方法

总结来看，目前的加速算法有两个优化方向，一是针对Latency的优化；二是针对Throughput的优化。

针对Latency的优化，主要还是底层的OP算子、矩阵优化、并行、更高效的C++解码等，如FasterTrnasformer以及DeepSpeed。针对Latency的优化可以提升Throughput，但没有直接用batch_size提升的更显著。

针对Throughput优化，主要是KV Cache存取优化，本质是降低显存开销，从而可以提升batch size。

这方面工作相对多一些，如offloading技术，就是如何高效利用第三方存储CPU/DRAM/Disk，使得GPU显存能空出来进而增大batch_size。

再如vLLM中的PagedAttention技术就是借鉴OS中的分页以及虚拟存储思想实现显存动态分配，也能节省很多显存空间。

还有如continuous batching ，变传统的static batch为动态可复用的batch分配，同样也能尽可能扩大batch_size，进而提升Throughput。

一些主流加速框架

FasterTransformer

  • atency-oriented
  • 方法：90%的时间消耗在12层Transformer的前向计算上，总结优化点如下：https://zhuanlan.zhihu.com/p/79528308
  • 为了减少kernel调用次数，将除了矩阵乘法的kernel都尽可能合并（这个可能是主要的）
  • 针对大batch单独进行了kernel优化
  • 支持选择最优的矩阵乘法
  • 在使用FP16时使用half2类型，达到half两倍的访存带宽和计算吞吐
  • 优化gelu、softmax、layernorm的实现以及选用rsqrt等

DeepSpeed Zero-Inference

  • 微软

  • 同时优化latency和Throughput

  • 优化Latency：a multi-GPU inference solution.主要是下面三个技术

  • parallelism： Tensor parallelism、Pipeline parallelism、Expert Parallelism（MoE）。对多机多卡之间的通信带宽要求较高。

  • communication optimization

  • optimized sparse kernels

  • 优化Throughput：Zero-Inference也是用到了offloading技术，如何结合GPU显存以及其他外部存储设备如DRAM、NVMe等加载大模型，问题变为How to apportion GPU memory among model weights, inference inputs and intermediate results.然后可以接受大的batch size，进而提升Throughput。

LLaMA.cpp

  • 最火热的开源社区产品

  • 面向消费级CPU/GPU的Inference框架，主打易用性，CPU支持

  • 方法：offloading、高效C++解码（没有用任何复杂的语句）

vLLM

  • UC Berkeley，文档里面说是Vicuna/Alpaca等的后台服务框架

  • Throughput-Oriented，目前看到claim提升最大的。

  • 方法：提出paged attention，动态分配K-V Cache，提升Batch_size

FlexGen

目前我们的努力方向都是减小延迟，因为我们的目标是构造一个实时系统，因此用户体验极为重要，那么延迟和准确率就是最重要的两个指标。但是大语言模型不应该只有这一种应用场景，其他场景比如大规模数据预处理、批量推荐系统、自动化报告生成、大规模文档生成等等离线批量计算场景，它们更在意吞吐量而不是延迟，因此我关注到了FlexGen这项工作。

与前面讲过的FlashAttention、PagedAttention等方法通过优化KV Cache存储、减小SRAM与HBM之间的读写次数的手段来减小LLMs内存浪费，加快推理速度不同，FlexGen这篇工作主要目标是利用有限资源（如单个GPU）进行高吞吐量的LLMs推断。FlexGen可以通过聚合GPU、CPU和磁盘的内存和计算，在各种硬件资源的约束下进行灵活的配置，通过解决线性规划问题，寻找出有效的模式来存储和访问张量。此外，FlexGen进一步将权重和注意力缓存压缩为4位，但却几乎没有精度损失。当在单个16GB GPU上运行OPT-175B时，FlexGen实现了有效批大小为144时每秒生成1个token的吞吐量。

这项工作采用的最主要并行策略是卸载策略（Offloading Strategy）。

LLM推理

典型的大语言模型生成推理任务包括两个阶段：i) 预填充阶段，该阶段接受一个提示序列以生成大语言模型每个transformer层的KV Cache； ii) 解码阶段，该阶段利用并更新 KV Cache逐步生成token，当前token的生成依赖于先前生成的token。具体公式以及吞吐量的定义见论文第三部分。

问题的形式化定义

下图是两个很简单的例子，图里面的彩色方块就表示在GPU上计算一个batch的一个layer，相同颜色块是共享权重的，这里假设模型有4个layer，并且每个prompt生成3个token。第一个调度就是很简单的把一个batch的从头算到尾一次算完，其实就是面向延迟的操作了，第二个是在batch维度加了一个block的粒度，在block内做zigzag，block间还是从头到尾，这才有面向吞吐量的意思。这些块也是有约束的：

  1. 每一个块都要确保它左边的所有块都被计算过，它才能被计算
  2. 一个块如果要在某个设备上进行计算，那么它所有的输入都要被加载到同一个设备上
  3. 一个块有两个输出：激活值和KV缓存。激活值要保存到右边的相邻块算完才能释放；KV缓存则要保存到同一行右边所有块都算完才能释放
  4. 存储的tensor的大小不能超过设备的内存大小

目标就是要找出一条有效的路径最小化总运行时间（计算开销+IO开销）。

schedule.png

卸载策略搜索空间

  • 计算调度：i ) 计算顺序：按照上面的图，逐行算延迟最低，是完成一个批次生成的最快方法，而且可以在处理完一行后立即释放KV Cache，但是由于每两个相邻的块不共享权重，这个方法不得不反复加载权重，从而产生巨大的I/O成本。逐列计算时，相邻计算的块共享权重，因此我们可以让权重保留在GPU上以便重用，只加载/卸载激活值和KV Cache，但是这样会有很大的存储压力。所以batch维度上block分块，可以平衡一下，但这种还不是最优的，不过作者为了调度的简单起见就只设计了这种，但他也证明了这个最多就比最优的差两倍。ii ) overlapping：第二个优化点在于重叠，可以在计算当前batch的时候，同时加载下一层的权重、保存上一batch的激活值和KV Cache、加载下一batch的激活值和KV Cache（也就是下图中最内层循环的六个函数可以并行进行）。这其实就是操作系统中pipeline的思想了，把不同的操作叠在一起，见缝插针。这里有两个参数：GPU的batch大小和block内的batch数量，这两的乘积也就是一个block的大小了。

    algrithm.png

  • Tensor位置：Tensor有三个位置可以放：显存、内存、硬盘。所有Tensor在这三种设备上各放多少是要考虑的，同时存放粒度也需要考虑：一个模型放多少层到GPU上(模型粒度)？一个层放多少tensor到GPU上(层粒度)？一个tensor放多少元素到GPU上(tensor粒度)？较粗粒度的划分会具有较低的运行时开销，但灵活性较低。作者权衡了一下，在权重上使用层粒度，在激活和 KV 缓存上使用张量粒度。

  • 计算设备：虽然用GPU算会很快，但当KV缓存很大的时候，是把一大堆KV Cache从内存挪到显存用GPU算呢？还是只把激活从显存挪到内存，用CPU算呢？作者发现当序列很长的时候，KV Cache如果是在内存的话，直接用CPU算比挪到显存用GPU算要好。

代价模型和策略搜索

作者将所有行为（比如从CPU读到GPU、从CPU写到硬盘等等）都用9个变量进行了表示，得到了代价模型。接着借助线性规划方法得到了满意的卸载策略。

Hugging Face pipeline Accelerate

Text Generation Inference

  • HuggingFace 官方支持的推理部署工具

  • 和 vllm 类似的 continuous batching

  • 支持了 flash-attention 和 Paged Attention。

  • 支持了 Safetensors 权重加载。

  • TGI 支持部署 GPTQ 模型服务，这使得我们可以在单卡上部署拥有 continous batching 功能的，更大的模型。

  • 支持采用 Tensor Parallelism 部署多 GPU 服务，模型水印等其他功能

Parameter Efficient Fine Tuning（PEFT）回顾

近来的PEFT方法大致可以分为三个种类：Adapter、Prompt、LoRA。相应论文的发表时间排序如下：Adapter(2019)、Prefix-Tuning(2021.1)、P-Tuning(2021.3)、Prompt-Tuning(2021.4)、LoRA(2021.6)、P-Tuning v2(2021.10)。

Adapter

paper：Parameter-Efficient Transfer Learning for NLP. Neil Houlsby, etc. (Google Research) PMLR, 2019.

2019年谷歌的研究人员首次在论文《Parameter-Efficient Transfer Learning for NLP》提出针对 BERT 的 PEFT微调方式，拉开了 PEFT 研究的序幕。他们指出，在面对特定的下游任务时，如果进行 Full-Fintuning（即预训练模型中的所有参数都进行微调），太过低效；而如果采用固定预训练模型的某些层，只微调接近下游任务的那几层参数，又难以达到较好的效果。

于是他们设计了如下图所示的 Adapter 结构，将其嵌入 Transformer 的结构里面，在训练时，固定住原来预训练模型的参数不变，只对新增的 Adapter 结构进行微调。同时为了保证训练的高效性（也就是尽可能少的引入更多参数），他们将 Adapter 设计为这样的结构：

  • 首先是一个 down-project 层将高维度特征映射到低维特征
  • 然后过一个非线形层之后，再用一个 up-project 结构将低维特征映射回原来的高维特征
  • 同时也设计了 skip-connection 结构，确保了在最差的情况下能够退化为identity（类似残差结构）。

Adapter.png

从实验结果来看，该方法能够在只额外对增加的 3.6% 参数规模（相比原来预训练模型的参数量）的情况下取得和Full-Finetuning 接近的效果（GLUE指标在0.4%以内）。

Prompt

paper：Prefix-Tuning: Optimizing Continuous Prompts for Generation. Xiang Lisa Li, Percy Liang. (Stanford University) ACL,2021.

2021年斯坦福的研究人员在论文《Prefix-Tuning: Optimizing Continuous Prompts for Generation》中提出了 Prefix Tuning 方法。与Full-finetuning 更新所有参数的方式不同，该方法是在输入 token 之前构造一段任务相关的 virtual tokens 作为 Prefix，然后训练的时候只更新 Prefix 部分的参数，而 Transformer 中的其他部分参数固定。该方法其实和构造 Prompt 类似，只是 Prompt 是人为构造的“显式”的提示，并且无法更新参数，而Prefix 则是可以学习的“隐式”的提示。

PrefixTuning.png

同时，为了防止直接更新 Prefix 的参数导致训练不稳定的情况，他们在 Prefix 层前面加了 MLP 结构(相当于将Prefix 分解为更小维度的 Input 与 MLP 的组合后输出的结果)，训练完成后，只保留 Prefix 的参数。

paper：P-tuning: GPT Understands, Too. Xiao Liu, Yanan Zheng, etc. (Tsinghua University) 2021.

PTuning.png

P-Tuning 提出将 Prompt 转换为可以学习的 Embedding 层，只是考虑到直接对 Embedding 参数进行优化会存在这样两个挑战：

  • Discretenes： 对输入正常语料的 Embedding 层已经经过预训练，而如果直接对输入的 prompt embedding进行随机初始化训练，容易陷入局部最优。
  • Association：没法捕捉到 prompt embedding 之间的相关关系。

P-tuning 依然是固定 LLM 参数，利用多层感知机和 LSTM 对 Prompt 进行编码，编码之后与其他向量进行拼接之后正常输入 LLM。注意，训练之后只保留 Prompt 编码之后的向量即可，无需保留编码器。

P-Tuning 和 Prefix-Tuning 差不多同时提出，做法其实也有一些相似之处，主要区别在：

  • Prefix Tuning 是将额外的 embedding 加在开头，看起来更像是模仿 Instruction 指令；而 P-Tuning 的位置则不固定。
  • Prefix Tuning 通过在每个 Attention 层都加入 Prefix Embedding 来增加额外的参数，通过 MLP 来初始化；而 P-Tuning 只是在输入的时候加入 Embedding，并通过 LSTM+MLP 来初始化。

paper：Prompt-tuning: The Power of Scale for Parameter-Effificient Prompt Tuning. Brian Lester, etc. (Google Research) EMNLP, 2021.

Prompt Tuning 是2021年谷歌在论文《The Power of Scale for Parameter-Efficient Prompt Tuning》中提出的微调方法。

该方法可以看作是 Prefix Tuning 的简化版本，只在输入层加入 prompt tokens，并不需要加入 MLP 进行调整来解决难训练的问题，主要在 T5 预训练模型上做实验。似乎只要预训练模型足够强大，其他的一切都不是问题。作者也做实验说明随着预训练模型参数量的增加，Prompt Tuning的方法会逼近 Fine-tune 的结果。

固定预训练参数，为每一个任务额外添加一个或多个 embedding，之后拼接 query 正常输入 LLM，并只训练这些 embedding。左图为单任务全参数微调，右图为 Prompt tuning。

PromptTuning.png

作者做了一系列对比实验，都在说明：随着预训练模型参数的增加，一切的问题都不是问题，最简单的设置也能达到极好的效果。

  • Prompt 长度影响：模型参数达到一定量级时，Prompt 长度为1也能达到不错的效果，Prompt 长度为20就能达到极好效果。
  • Prompt初始化方式影响：Random Uniform 方式明显弱于其他两种，但是当模型参数达到一定量级，这种差异也不复存在。
  • 预训练的方式：LM Adaptation 的方式效果好，但是当模型达到一定规模，差异又几乎没有了。
  • 微调步数影响：模型参数较小时，步数越多，效果越好。同样随着模型参数达到一定规模，zero shot 也能取得不错效果。
  • 当参数达到100亿规模与全参数微调方式效果无异。

paper：P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks. Xiao Liu, etc. (Tsinghua University) 2021.

PTuning.png

P-Tuning 的问题是在小参数量模型上表现差。于是就有了v2版本：《P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks》。

从标题就可以看出，P-Tuning v2 的目标就是要让 Prompt Tuning 能够在不同参数规模的预训练模型、针对不同下游任务的结果上都达到匹敌 Fine-tuning 的结果。

那也就是说当前 Prompt Tuning 方法在这两个方面都存在局限性。

  • 不同模型规模：Prompt Tuning 和 P-tuning 这两种方法都是在预训练模型参数规模够足够大时，才能达到和Fine-tuning 类似的效果，而参数规模较小时效果则很差。
  • 不同任务类型：Prompt Tuning 和 P-tuning 这两种方法在 sequence tagging 任务上表现都很差。

相比 Prompt Tuning 和 P-tuning 的方法， P-tuning v2 方法在多层加入了 Prompts tokens 作为输入，带来两个方面的好处：

  1. 带来更多可学习的参数（从 P-tuning 和 Prompt Tuning 的0.1%增加到0.1%-3%），同时也足够 parameter-efficient。
  2. 加入到更深层结构中的 Prompt 能给模型预测带来更直接的影响。

v1 到 v2 的可视化：蓝色部分为参数冻结，橙色部分为可训练部分。

LORA

paper：LoRA: Low-Rank Adaptation of Large Language Models. Edward Hu, etc. (Microsoft Corporation) ICLR, 2022.

自然语言处理目前存在一个重要范式：一般领域数据的大规模预训练，对特定任务或领域的适应微调（finetune）。但是随着预训练语言模型越来越大，这个范式存在以下问题：

  • 当我们finetune大模型时，由于训练成本太高，不太可能重新训练所有模型参数。

  • 以前的方法（论文发表于2021年）都或多或少有其它性能问题，如adapter增加了模型层数，引入了额外的推理延迟；prefix-tuning比较难训练，效果不如直接finetune。

基于上述背景，论文作者得益于前人的一些关于内在维度（intrinsic dimension）的发现：模型是过参数化的，它们有更小的内在维度，模型主要依赖于这个低的内在维度（low intrinsic dimension）去做任务适配。假设模型在任务适配过程中权重的改变量是低秩（low rank）的，由此提出低秩自适应（LoRA）方法，LoRA允许我们通过优化适应过程中密集层变化的秩分解矩阵来间接训练神经网络中的一些密集层，同时保持预先训练的权重不变。

方法

LoRA的实现思想很简单，如下图所示，就是冻结一个预训练模型的矩阵参数，并选择用A和B矩阵来替代，在下游任务时只更新A和B。

lora.png

结合图片来看，LoRA的实现流程如下：

  • 在原始预训练语言模型（PLM）旁边增加一个旁路，做一个降维再升维的操作，来模拟所谓的内在秩。
  • 训练的时候固定PLM的参数，只训练降维矩阵A与升维矩阵B。
  • 模型的输入输出维度不变，输出时将BA与PLM的参数叠加。
  • 用随机高斯分布初始化A，用0矩阵初始化B，保证训练的开始此旁路矩阵依然是0矩阵。

实现

接下来我们从公式上解释LoRA的实现。假设要在下游任务微调一个预训练语言模型（如GPT3），则需要更新预训练模型参数，公式表示如下：
$$
h=W_0x+\Delta Wx=W_0x+BAx
$$
W0是预训练模型初始化的参数，ΔW就是需要更新的参数。如果是全参数微调，则$\Delta W$参数量=W0参数量（如果是GPT3，则ΔW≈175B）。从这可以看出要全参数微调大语言模型，小家小户是不可能的。由于前人的工作发现预训练的语言模型具有较低的“内部维度（intrinsic dimension）”，在任务适配过程中，即使随机投影到较小的子空间，仍然可以有效地学习。因此，LoRA做的就是增加小参数模块去学习改变量ΔW。在训练过程中，W0是固定不变的，只有A和B包含训练参数，是变化的。而在推理的过程中，只需要把改变量放回原模型，就不会有任何延迟。如果想切换任务，只需要切换任务的过程中，减去BA，然后换上用其它任务训练好的新BA就可以了。

总结

总的来说，基于大模型的内在低秩特性，增加旁路矩阵来模拟full finetuning，LoRA是一个能达成lightweight finetuning的简单有效的方案。目前该技术已经广泛应用于大模型的微调，如Alpaca，stable diffusion+LoRA，而且能和其它参数高效微调方法有效结合，例如 State-of-the-art Parameter-Efficient Fine-Tuning (PEFT)。