大模型推理优化技术

Tw_xxxx
757 阅读10分钟

推理过程

推理会分成 prefill 和 decoding 两个阶段。每一个请求发起后产生的推理过程都会先经历一个 Prefill 过程，prefill 过程会计算用户所有的输入，并生成对应的 KV 缓存，再经历若干个 decoding 过程，每一个 decoding 过程服务器都会生成一个字符，并将其放入到 KV 缓存当中，推理出来的预测结果又放入输入中，如此循环往复，直到推理出最终结果。新的请求进来在进行完 prefill 之后会不断迭代进行 decoding，每一个 decoding 阶段结束之后都会将结果当场返回给客户。这样的生成过程称为流式传输。

202408021943.png

推理性能的评价指标

Throughput（吞吐量）

  • 吞吐量是指当系统的负载达到最大的时候，在单位时间内能够执行多少个 decoding，即生成多少个字符。
  • 测试吞吐量的方法是，假设所有用户都会在同一时刻到来，并且这些用户问的都是一样的问题，这些用户可以同时启动和结束，且他们生成的文本的长度和输入的文本长度都是一样的。通过使用完全相同的输入，组成一个完整的 batch。在这种情况下，系统的吞吐量会达到最高。
  • 如果用户输入的长度和生成的长度很长，那么系统吞吐量不会很高。

First Token Latency（首字延迟）

  • 当一批用户进入到推理系统之后用户完成 Prefill 阶段的过程需要花多长时间。这也是系统生成第一个字符所需的响应时间。
  • 希望用户在系统上输入问题后得到回答的时间小于 2~3 秒。
  • 与首字延迟最相关的就是用户的输入长度，用户输入的长度越长，首字延迟也会越高

Latency（延迟）

  • 每一个 decoding 所需要的时长。它反映的是大模型每生成一个字符的间隔是多长时间，也就是生成的过程有多么流畅。
  • 大部分情况下希望生成的延迟小于 50 毫秒，也就是一秒钟生成 20 个字符。这样生成是比较流畅的。
  • 主要受到 batch size 的影响，batch size 越大推理延迟也会越大

QPS（每秒请求数）

  • 一秒钟能够处理多少个用户的请求。表示系统中每秒可以处理多少个请求。
  • 有些用户会提前生成完，而有些用户要生成很多长度之后才会结束。所以有很多地方的 GPU 会空闲。因此QPS 并不能够发挥完全的吞吐量优势。吞吐量可能很大，但实际的处理能力可能会很差。

模型压缩

模型压缩的基本动机在于当前的模型是冗余的，可以在精度损失很小的情况下实现模型小型化

稀疏(Sparsity)

实现稀疏(Sparsity)的一个重要方法是剪枝(Pruning)。剪枝是在保留模型容量的情况下，通过修剪不重要的模型权重或连接来减小模型大小。 它可能需要也可能不需要重新培训。修剪可以是非结构化的或结构化的。

  • 非结构化剪枝允许删除任何权重或连接，因此它不保留原始网络架构。 非结构化剪枝通常不适用于现代硬件，并且不会带来实际的推理加速。
  • 结构化剪枝旨在维持某些元素为零的密集矩阵乘法形式。 他们可能需要遵循某些模式限制才能使用硬件内核支持的内容。 当前的主流方法关注结构化剪枝，以实现 Transformer 模型的高稀疏性。

量化(Quantization)

  • 训练后量化（Post-Training Quantization，PTQ）：模型首先经过训练以达到收敛，然后我们将其权重转换为较低的精度，而无需进行更多训练。 与训练相比，实施起来通常相当便宜。
  • 量化感知训练（Quantization-Aware Training，QAT）：在预训练或进一步微调期间应用量化。 QAT 能够获得更好的性能，但需要额外的计算资源和对代表性训练数据的访问。
  • 实际上，由于 GPU 内核缺乏对某些类型的矩阵乘法（例如 INT4 x FP16）的支持，理论最优量化策略与硬件内核支持之间的差距，并非以下所有方法都能加速实际推理。

蒸馏(Distillation)

  • 知识蒸馏是一种构建更小、更便宜的模型的直接方法，通过从预先训练的昂贵模型中转移技能来加速推理融入 student。 除了与 teacher 匹配的输出空间以构建适当的学习目标之外，对于如何构建 student 架构没有太多限制。
  • 在大模型时代，蒸馏可以与量化、剪枝或稀疏化技术相结合，其中 teacher 模型是原始的全精度密集模型，而 student 模型则经过量化、剪枝或修剪以具有更高的稀疏级别，以实现模型的小型化。

202402192344.png

Transformer 结构优化

  • 由于 Transformer 固有的 O(N^2) 计算复杂度和内存限制的键值缓存，在推理过程中表现出低效率。特别对于长序列来说，低效率无法接受。

202403012201.png

  • 可以通过减少头的数量，减少 kv cache的 size，达到减小带宽的压力的目的，那么推理速度势必更快。

202401301832.png

202401301711.png

202401301808.png

Multi Head Attention

标准的多头注意力机制，h 个 Query、Key 和 Value 矩阵。

Multi Query Attention

MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。

Group Query Attention

GQA将查询头分成 G 组，每个组共享一个 Key 和 Value 矩阵。GQA-G 是指具有 G 组的 grouped-query attention。GQA-1 具有单个组，因此具有单个 Key 和 Value，等效于 MQA。而 GQA-H 具有与头数相等的组，等效于MHA。

KVCache

  • 在decoding phase中需要计算当前token和之前所有已生成token的attention，因此需要计算所有token的k和v向量，但是前面的token的kv值在每轮decoding中都被重复计算了，因此我们可以把它们存下来，存成两个[seq_len-1, inner_dim]的Tensor，在每轮计算中只需要计算当前token的kv值即可。
  • KVCache是最简单直接的优化手段，一般模型的默认实现都会自带KVCache因此并不需要额外实现。

Transformer 结构优化（硬件）

Flash Attention

GPU 存储架构

202403012202.png

概述

  • 通过利用更高速的上层存储计算单元SRAM，减少对低速更下层存储器HBM（高带宽内存）的访问次数，来提升模型的训练性能。
  • 主要关注 IO-aware（IO感知），进一步优化 GPU 显存的读写效率

Standard Attention

202403012203.png

  • 标准 Attention 算法的总HBM访问次数为O(Nd + N^2)。当N比较大时，总的HBM访问次数可能会比较昂贵。
  • 标准Attention算法在GPU内存分级存储的架构下，存在以下缺陷：
    • 过多对HBM的访问，如S、P需要在存入HMB后又立即被访问，HBM带宽较低，从而导致算法性能受限
    • S、P需要占用O(N^2)的存储空间，显存占用较高

具体过程

  • Tiling（在向前和向后传递时使用）：将原始的注意力矩阵分解成更小的子矩阵，然后分别对这些子矩阵进行计算，只要这个子矩阵的大小可以在SRAM内存放，就可以在计算过程中只访问SRAM。
  • Recomputation（仅在向后传递中使用）：是一种算力换内存的把戏，就是不要存储那么多梯度和每一层的正向传播的中间状态，而是在计算到反向某一层的时候再临时从头开始重算正向传播的中间状态。

202403012204.png

202403012207.png

  1. 计算分子块的大小
    • 为了让Q、K、V在计算中可以存放在SRAM中，我们需要设定分块的大小尺寸，保证子块大小不超过SRAM的大小即可。
  2. 初始化输出矩阵O
    • 为SRAM上的输出O矩阵赋值为全0，它将作为一个累加器保存softmax的累积分母。l也类似。m用于记录每一行行最大分数，其初始化为-inf。
  3. 切分子块
    • 按步骤一中的块大小将Q, K和V分成块。同时将将O, l, m分割成块(与Q的块大小相同)。
  4. 外循环加载K、V，内循环加载Q子块

202403012205.png

*   外循环：对于每一个Block Key和Value，从HBM加载进SRAM
*   内循环：对于每个Block Query，Oi, li, mi，从HBM加载进SRAM
*   在SRAM上完成Block S的计算
*   这里要注意的是，Oi, li, mi其中存储的可能是上一个循环计算的中间结果。

5. 实现分块SoftMax算法

202403012206.png

  1. 反向传播
    • 利用 SRAM 中的 Q、K、V 重新计算 S 和 P 矩阵。使用更多的 flop 减少 HBM 访问。

Paged Attention

概述

  • PagedAttention 是一种新颖的注意力算法，它将在操作系统的虚拟内存中分页的经典思想引入到 LLM 服务中。在无需任何模型架构修改的情况下，可以做到比 HuggingFace Transformers 提供高达 24 倍的 Throughput。
  • 在自回归 decoder 中，所有输入到 LLM 的 token 会产生注意力 key 和 value 的张量，这些张量保存在 GPU 显存中以生成下一个 token。这些缓存 key 和 value 的张量通常被称为 KV cache，其具有以下特点：
    • 显存占用大：在 LLaMA-13B 中，缓存单个序列最多需要 1.7GB 显存；
    • 动态变化：KV 缓存的大小取决于序列长度，这是不可预测的。因此，这对有效管理 KV cache 挑战较大。研究发现，由于碎片化和过度保留，现有系统浪费了 60% - 80% 的显存。

原理

  • PagedAttention 允许在非连续的内存空间中存储连续的 key 和 value 。具体来说，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。
  • 因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理 key 和 value ，就像在操作系统的虚拟内存中一样：可以将块视为页，将 token 视为字节，将序列视为进程。那么通过一个块表就可以将连续的逻辑块映射到非连续的物理块，而物理块可以根据新生成的token按需分配。

202403012300.png

  • PagedAttention的另外一个好处是高效内存共享。通过块表可以自然地实现内存共享，在PagedAttention中的不同序列通过将逻辑块映射到一样的物理块上可以实现共享块，大大降低了内存使用量。
avatar
文章
阅读
粉丝