NLP（十八）：LLM 的推理优化技术纵览

NLP（十八）：LLM 的推理优化技术纵览

紫气东来

目录

收起

一、子图融合（subgraph fusion）

1.1 FasterTransformer by NVIDIA

1.2 DeepSpeed Inference by Microsoft

1.3 MLC LLM by TVM

二、模型压缩（Model Compression）

2.1 稀疏(Sparsity)

2.2 量化(Quantization)

2.3 蒸馏(Distillation)

三、并行化（Parallelism）

3.1 数据并行 (Data Parallelism, DP)

3.2 张量并行(Tensor Parallelism, TP)

3.3 流水线并行(Pipeline Parallelism, PP)

四、Transformer 结构优化

4.1 FlashAttention

4.2 PagedAttention

4.3 FLAT Attention

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

5.1 ORCA

5.2 FastServe

5.3 vLLM

5.4 Text Generation Inference

5.5 LMDeploy

六、KV cache 优化

七、投机采样与MoE

7.1 投机采样

7.2 超大模型的 MoE

八、硬件升级

8.1 NVIDIA H100 PCIe

8.2 AMD MI300

8.3 Apple M2 Ultra

参考资料

推理是 LLM 应用的重要一环，在部署服务环节影响重大，本文将讨论主流的 LLM 的推理优化技术。

一、子图融合（subgraph fusion）

图融合技术即通过将多个 OP（算子）合并成一个 OP（算子），来减少Kernel的调用。因为每一个基本 OP 都会对应一次 GPU kernel 的调用，和多次显存读写，这些都会增加大量额外的开销。

1.1 FasterTransformer by NVIDIA

FasterTransformer(FT) 是一个用于实现基于Transformer的神经网络推理的加速引擎。FT框架是用C++/CUDA编写的，依赖于高度优化的 cuBLAS、cuBLASLt 和 cuSPARSELt 库，与 NVIDIA TensorRT 等其他编译器相比，FT 的特点是它支持以分布式方式推理 Transformer 大模型。

图融合是FT 的一个重要特征，将多层神经网络组合成一个单一的神经网络，将使用一个单一的内核进行计算。 这种技术减少了数据传输并增加了数学密度，从而加速了推理阶段的计算。 例如， multi-head attention 块中的所有操作都可以合并到一个内核中。

除此之外，FT还对部分大模型分别支持：

• INT8 低精度量化推理
• Ampere 架构的 GPU 硬件部分支持稀疏化
• Hopper 架构支持 FP8 推理
• Tensor 并行
• Pipeline 并行

  1.2 DeepSpeed Inference by Microsoft

  对于 Transformer layer，可分为以下4个主要部分：

  1. Input Layer-Norm plus Query, Key, and Value GeMMs and their bias adds.
  2. Transform plus Attention.
  3. Intermediate FF, Layer-Norm, Bias-add, Residual, and Gaussian Error Linear Unit (GELU).
  4. Bias-add plus Residual.

  如图所示，每一部分可分别进行融合，与未融合相比，以上几个部分的加速比可分别达到 1.5x, 2.9x, 3x, 1.2x 。

  

  除此之外，DeepSpeed Inference 的优化点还有以下几点：

  • 多 GPU 的并行优化
  • INT8 模型量化
  • 推理的 pipeline 方案

    更多详细介绍及实践可参考笔者之前的文章：

    紫气东来：NLP（十二）：DeepSpeed Inference 在 LLM 推理上的优化探究192 赞同 · 13 评论文章​编辑

    1.3 MLC LLM by TVM

    之前介绍的推理方案主要是基于GPU的优化，而 MLC LLM 提供了可应用于移动端（例如 iPhone）、消费级电脑端（例如 Mac）和 Web 浏览器的轻设备解决方案。

    MLC LLM 的主要工作流基于 Apache TVM Unity，通过扩展 TVM 后端使模型编译更加透明和高效。其中以编译代码转换、融合、内存规划和库卸载（library offloading）为代表的可组合的 ML 编译优化是其中重要的优化特性。

    

    除此之外，MLC LLM 还具有以下特性：

    • Dynamic shape：避免了对最大输入长度进行额外填充的需要，并减少了计算量和内存使用量。
    • 量化：MLC LLM 利用低位量化来压缩模型权重，并利用 TVM 的 loop-level TensorIR 为不同的压缩编码方案快速定制代码生成。
    • 运行时（Runtime）：TVM 编译生成的库能够通过 TVM runtime 在设备的原生环境中运行，TVM runtime 支持 CUDA/Vulkan/Metal 等主流 GPU 驱动以及 C、JavaScript 等语言的绑定。

      除了上述3种方案外，其他也支持图融合的方案还包括 NVIDIA TensorRT， Tencent TurboTransformers 等。

      二、模型压缩（Model Compression）

      模型压缩的基本动机在于当前的模型是冗余的，可以在精度损失很小的情况下实现模型小型化，主要包括3类方法：稀疏(Sparsity)、量化(Quantization)、蒸馏(Distillation)。

      2.1 稀疏(Sparsity)

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

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

        关于剪枝稀疏的基本原理，可参考笔者之前的文章：

        紫气东来：NLP（八）：大语言模型的稀疏化技术48 赞同 · 6 评论文章​编辑

        除了上文介绍的稀疏方法外，还有其他的稀疏化方法，包括但不限于：

        • SparseGPT：该方法的工作原理是将剪枝问题简化为大规模的稀疏回归实例。它基于新的近似稀疏回归求解器，用于解决分层压缩问题，其效率足以在几个小时内使用单个 GPU 在最大的 GPT 模型（175B 参数）上执行。同时，SparseGPT 准确率足够高，不需要任何微调，剪枝后所损耗的准确率也可以忽略不计。
        • LLM-Pruner：遵循经典的“重要性估计-剪枝-微调”的策略，能够在有限资源下完成大语言模型的压缩，结果表明即使剪枝 20％ 的参数，压缩后的模型保留了 93.6％ 的性能。
        • Wanda: 该方法由两个简单但必不可少的组件构成——剪枝度量和剪枝粒度。剪枝度量用来评估权重的重要性，然后按照剪枝粒度进行裁剪。该方法在 65B 的模型上只需要 5.6 秒就可以完成剪枝，同时达到SparseGPT相近的效果。

          以上主要实现了稀疏的方法，那么对于稀疏后的模型如何加速呢？NVIDIA Ampere 架构对与结构化稀疏做了专门的稀疏加速单元，下图展示了结构化稀疏的物理表示：

          

          2:4 结构化稀疏表示

          下图展示了稀疏单元GEMM计算与标准GEMM计算的区别（详细解释参见https://arxiv.org/pdf/2104.08378.pdf）

          

          Sparse VS Dense GEMM

          2.2 量化(Quantization)

          常见量化有两种常见方法：

          • 训练后量化（Post-Training Quantization，PTQ）：模型首先经过训练以达到收敛，然后我们将其权重转换为较低的精度，而无需进行更多训练。 与训练相比，实施起来通常相当便宜。
          • 量化感知训练（Quantization-Aware Training，QAT）：在预训练或进一步微调期间应用量化。 QAT 能够获得更好的性能，但需要额外的计算资源和对代表性训练数据的访问。

            实际上，由于 GPU 内核缺乏对某些类型的矩阵乘法（例如 INT4 x FP16）的支持，理论最优量化策略与硬件内核支持之间的差距，并非以下所有方法都能加速实际推理。

            关于量化的基本原理和实现细节，可参考笔者之前的文章：

            紫气东来：NLP（十一）：大语言模型的模型量化(INT8/INT4)技术278 赞同 · 15 评论文章​编辑

            许多关于 Transformer 模型量化的研究都有相同的观察结果：简单的低精度（例如 8 bit）训练后量化会导致性能显着下降，这主要是由于动态的 activation 和静态的 weight 量化策略无法保持一致。

            

            为了不损失精度而提高性能，可以考虑 WeightOnly 量化技术，即只把 Weight 量化成 int8 格式，以降低访存压力。到实际 Kernel 内部再 Dequantize 回 fp16，进行矩阵乘计算。这种方法在 BS 较小是比较有效(因为此时的瓶颈在IO)，BS 较大时(因为此时的瓶颈在计算)效果变差。

            

            WeightOnly 量化的典型案例是 AWQ: Activation-aware Weight Quantization，即只对 weight 进行量化以实现压缩和加速的效果。

            2.3 蒸馏(Distillation)

            知识蒸馏是一种构建更小、更便宜的模型（“student 模型”）的直接方法，通过从预先训练的昂贵模型中转移技能来加速推理（“ teacher 模型”）融入 student。 除了与 teacher 匹配的输出空间以构建适当的学习目标之外，对于如何构建 student 架构没有太多限制。

            

            知识蒸馏基本框架

            给定数据集，训练 student 模型通过蒸馏损失来模仿 teacher 的输出。 通常神经网络有一个softmax层； 例如，LLM 输出 token 的概率分布。 将 softmax 之前的 logits 层表示为 �� 和 �� , 分别表示 teacher 和 student 模型。 蒸馏损失最小化两个 softmax 输出之间的差异（温度 � ）。 当标签 � 已知，可以将其与student 的 logits 之间计算交叉熵，最后将两个损失相加，如下：

            �KD=�distll (softmax⁡(��,�),softmax⁡(��,�))+��CE(�,��)

            在 Transformer 中一个典型案例是DistilBERT，模型参数减少 40%，速度提升71%。在大模型时代，蒸馏可以与量化、剪枝或稀疏化技术相结合，其中 teacher 模型是原始的全精度密集模型，而 student 模型则经过量化、剪枝或修剪以具有更高的稀疏级别，以实现模型的小型化。

            三、并行化（Parallelism）

            当前的推理的并行化技术主要体现在3个维度上，即 3D Parallelism:

            • Data Parallelism(DP)
            • Tensor Parallelism(TP)
            • Pipeline Parallelism(PP)

              

              3D Parallelism 的3个维度

              3.1 数据并行 (Data Parallelism, DP)

              在推理中，DP 主要是增加设备数来增加系统整体 Throughput，其中最经典的即DeepSpeed的Zero系列

              

              另外 FSDP 也比较高效和易用

              

              3.2 张量并行(Tensor Parallelism, TP)

              在推理中，TP 主要是横向增加设备数通过并行计算来减少 latency，其实现原理及细节可参考笔者之前的文章

              紫气东来：NLP（六）：GPT 的张量并行化（tensor parallelism）方案50 赞同 · 11 评论文章​编辑

              当前也有一些方便易用的 TP 方案，如 BlackSamorez/tensor_parallel ，使用起来非常简单：

              import transformers
              import tensor_parallel as tp
              tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/opt-13b")
              model = transformers.AutoModelForCausalLM.from_pretrained("facebook/opt-13b")  # use opt-125m for testing
              model = tp.tensor_parallel(model, ["cuda:0", "cuda:1"])  #