迈向100倍加快：全栈Transformer推理优化

作者 | 符尧，爱丁堡大学博士生

OneFlow编译

翻译｜宛子琳、杨婷

假定有两家公司，它们拥有相同强大的模型。公司A能够用1个GPU为10个用户供给模型，而公司B能够用1个GPU为20个用户供给模型。从长远来看，谁会在竞赛中获胜呢？

答案是公司B，由于它的成本更低。

假定一位研讨人员提出了一种超级聪明的解码办法：这种办法拥有巧妙的算法和厚实的数学根底，但无法与FlashAttention兼容。它能在出产环境中运用吗？

或许不行，由于FlashAttention对大规划模型部署至关重要。

对Transformer推理的深化了解对研讨和出产极为有益。但是在现实中，大规划出产一般与前沿研讨的关联并不亲近，了解算法的人或许并不了解MLSys，反之亦然。

本文评论了全栈Transformer推理优化，从A100内存层次结构等硬件规范，到FlashAttention和vLLM等MLSys办法，再到专家混合等模型架构，以及估测性解码（Speculative Decoding）及其变体等解码算法。咱们确认了一个最基本的现实：Transformer推理受限于内存，且大部分优化（无论来自MLSys仍是建模）都根据/利用了这一现实。就像在角色扮演游戏中添加buff相同，能够看到Transformer推理是怎么逐渐扩展和加快的。

（本文由OneFlow编译发布，转载请联系授权。原文：yaofu.notion.site/Towards-100…

1

硬件：GPU上的推理

首先，咱们将探讨GPU架构，特别是其内存层次结构。咱们确认了两个重要形式：核算约束（compute bound）和内存约束（memory bound），并评论了大规划Transformer推理受内存约束的原因。大部分优化都根据Transformer推理受内存约束这一基本现实，例如只要咱们进步FLOP利用率，就能进步功率。

1.1 初步推理

1.1.1 GPU架构

GPU架构整体如下图所示：

迈向100倍加快：全栈Transformer推理优化

  • 根底部分：DRAM（动态随机存取存储器）、L2缓存和SM（流处理器单元）
  • 与CPU对比
    • SM类似于CPU中心，但具有更高级的并行性
    • L2缓存和DRAM类似于CPU的L2缓存和DRAM
    • 在Flash Attention论文中，L2缓存被称为SRAM（静态随机存取存储器）
  • A100 80G SXM

    • 108个SM，DRAM容量为80GB，有40M L2缓存

SM内部包括什么？

迈向100倍加快：全栈Transformer推理优化

  • L1缓存：指令和数据
  • 张量中心：进行矩阵乘法运算的当地。回想一下，神经网络核算基本上便是巨大批量的矩阵乘法。

1.1.2 GPU编程根底

在履行model.generate(prompt)时，咱们进行以下操作：

  • 内存拜访:
    • 从高带宽内存（HBM）加载模型权重到L2缓存，然后传输到SM（流处理器单元）
  • 核算:
    • 在SM中履行矩阵乘法，SM请求张量中心履行核算
  • A100:
    • 108个SM，DRAM容量为80G，40M L2缓存
    • bf16张量中心：每秒312万亿浮点运算（TFLOPS）
    • DRAM内存带宽为2039GB/秒 = 2.039T/秒
  • 假设模型很大，咱们将其分割到多个GPU上，比方两个由NVLink连接的GPU
    • NVLink 300GB/秒 = 0.3T/秒
    • 咱们大致观察了速度层次结构。虽然不能直接比较，但它们的数量级差异是咱们需要优化的主要方面：
    • 312T（SM核算） > 2.03T（DRAM内存拜访） > 0.3T=300G（NVLink跨设备通讯） > 60G（PCIe跨设备通讯）
  • 这意味着，假设咱们期望速度更快，咱们应该尽力：

    • 充分利用SM

    • 削减单个GPU的内存拜访（由于它比核算慢得多），

      • 削减GPU之间的通讯（由于它甚至比内存拜访还要慢）。

1.1.3 核算约束与内存约束

怎么确认咱们是否充分利用了SM呢？咱们经过以下办法检查是否核算或内存约束：

  • 界说每字节GPU操作= flop / 内存带宽
    • A100 = 312 / 2.039
  • 界说核算强度= 核算 / 内存拜访
  • 假设核算强度大，说明程序更会遭到核算约束；假设核算强度较小，则更受内存约束。

迈向100倍加快：全栈Transformer推理优化

  • 添加批次巨细会将行为从内存约束变为核算约束。
  • 内核融合：削减了内存拜访操作，由于咱们将多个操作合并为一个操作。

1.2 Transformer推理根底

1.2.1 预填充和解码

调用model.generate(prompt)时有两个过程：

  • 预填充：
    • 为提示核算键值（kv）缓存。
    • 这一过程受核算约束，由于咱们并行核算了一系列词元。
  • 解码：
    • 自回归采样下一个词元。
    • 这一过程受内存约束，由于咱们仅核算一个词元，未充分利用SM。

1.2.2 Transformer推理受内存约束

添加批处理巨细能够将形式从内存约束变为核算约束，正如Kipply博客《大型言语模型的推理演算》中的下图所示。

迈向100倍加快：全栈Transformer推理优化

  • 由于解码每次只采样一个词元。
  • 添加批次巨细进步了硬件功率。
    • 由于咱们一起核算多个词元。
    • 大型批次从内存约束变为核算约束。
  • 但是，咱们或许无法运用太大型的批次巨细，由于GPU内存有限，现在最大为80GB。

1.2.3 内存布局

迈向100倍加快：全栈Transformer推理优化

正如咱们所看到的，为了在bf16格式下运转一个13B模型，咱们大约只有10GB的内存来存储kv缓存。这意味着：

  • 不能运用太大型的批次（虽然咱们期望运用更大的批次巨细以进步功率）
  • 也不能处理太长的序列，虽然咱们的确期望能够处理长度为100k的序列。

1.2.4 在线与离线推理，吞吐量vs时延

  • 离线：吞吐量优化
    • 咱们重视这种状况，由于期望能够离线评价模型，例如，在100个基准测验上运转一个中心的预练习检查点（checkpoint），用于验证预练习是否状况杰出。
    • 添加批次巨细是有协助的，但需要记住当前单个设备的最大内存为80GB。
  • 在线：时延和吞吐量权衡
    • 当批次巨细较大型时（假定依然适配内存），会遭到核算吸纳之，随后时延会相应添加。
    • 时延不该慢于人类的阅览速度，不然用户会诉苦，或许转而运用竞赛对手的模型。
    • 但再次着重，咱们的确期望运用更大的批次巨细以进步功率。

1.2.5 上下文长度扩展

到现在为止，咱们的假定都是根据提示不长（低于4k）的状况。接下来，咱们将考虑提示超过100k的状况，期望模型阅览多个PDF文件，然后进行文档问答。

  • 预填充
  • 这一次，预填充需要花费更长时刻，由于输入长度很长
  • 这种状况下，生成首个词元的时延很重要，由于用户期望在10秒内看到模型生成的呼应。
  • 大型KV缓存
    • 由于上下文长度很长，咱们需要一个大型的键值（KV）缓存。
  • 现在就改善这方面的推理而言，人们做的作业好像并不多。

2

MLSys：Flash Attention和vLLM

本节评论了怎么充分利用GPU内存层次结构。vLLM供给了一种进行GPU内存办理的办法，就像在操作体系中办理CPU虚拟内存相同；Flash Attention展现了怎么经过在SM上履行大部分操作来有效削减内存IO，然后显着削减内存拜访核算开支。

2.1 vLLM和Paged Attention

由于GPU内存有限，咱们期望合理运用它来存储键值（kv）缓存。但是，GPU或PyTorch自身并不会主动供给将KV缓存置入内存的最佳办法，而其默许战略实际上适当糟糕。这激发了vLLM中用于GPU内存办理的Paged Attention：

迈向100倍加快：全栈Transformer推理优化

  • Paged Attention基本构建了一个类似于CPU内存办理的内存办理体系，以削减内存碎片并充分利用内存吞吐量
  • 现在这成为了Transformer推理的首要挑选

2.2 Flash attention

对每位从业者来说都是必备的，请简略记住完整的算法。

迈向100倍加快：全栈Transformer推理优化

要点：

  • 不要将完整的注意力矩阵存储在HBM中，而是对点积进行分块核算（blockwise computation），以便一切核算都在L2高速缓存中履行。

主要优势：

  • 大幅削减内存运用，能够运用蛮力法（即“brutal force”，经过穷举一切或许性来解决问题）处理包括100k上下文长度的状况——没错，并没有运用杂乱算法来处理100k上下文长度，而是采用了蛮力法。
    • 在原始论文中，作者们最多只测验了16k，但其实完全能够处理100k的上下文长度。
  • 大幅进步吞吐量，尤其是关于小模型，其中大部分浮点运算用于点积操作。

2.3Flash解码

迈向100倍加快：全栈Transformer推理优化

要点：不再运用一个查询来扫描整个KV缓存，而是复制查询，以便能够并行扫描KV缓存的不同块（chunk）。

3

建模：架构与解码算法

现在咱们进入一个大家更熟悉的范畴。咱们将从蒸馏（distillation）和量化（quantization）等规范且广为人知的技能开端，然后深化探讨专家混合（MoE）和估测性解码（speculative decoding）等进阶主题。

3.1 经典办法

3.1.1 蒸馏和稀少注意力

  • 蒸馏：利用较大模型的输出/logits对小模型进行微调
    • 对输出进行微调：现在，大家都很擅长从GPT进行蒸馏，所以这部分略过
    • 对logits/分布进行微调的探索较少，一些结果表明，经过这种办法收敛速度更快，质量更好
  • 稀少和局部注意力，尤其适用于长上下文
    • 在小模型年代，这一范畴得到了深化研讨（拜见Yi Tay的超卓调研，arxiv.org/abs/2009.06… ），但不确认这些结果是否适用于更大规划的模型
    • 在大模型年代，除了Mistral的滑动窗口注意力，这一范畴几乎没有相关研讨

3.1.2 量化

github.com/TimDettmers…

  • 最基本的办法，将模型权重量化为int 8。
  • 量化现在是部署大模型的必备过程。好消息是，它实际上并不会对功能造成本质损害。

迈向100倍加快：全栈Transformer推理优化

  • Yi-34B谈天模型的4位量化仅需17G内存。而基准测验的功能几乎相同。
  • 实际上它的速度适当快，你能够在Hugging Face（huggingface.co/spaces/01-a… ）上进行测验。

3.1.3 多查询注意力和组查询注意力

迈向100倍加快：全栈Transformer推理优化

  • 一般状况下，多查询注意力经过一起削减内存和核算，显着加快了练习和推理速度。
  • 多查询注意力也是一个很好的例子，小模型与大模型在此方面的差异并不显着：关于7B这样的小模型，多查询注意力不如全注意力，但当模型到达70B时，多查询注意力的功能基本上与全注意力相同。LLaMA2 7B运用全注意力，而70B运用组查询注意力。
  • 现有的顶尖大模型默许运用多查询注意力。

3.2 进阶技能

3.2.1 专家混合模型

从头开端预练习

  • 假定咱们有7B激活单元，一共34B参数。
  • 能否实现以下目标？
    • 与34B类似的功能
    • 优于34B的吞吐量
    • 与7B类似的时延 最近发布的Mistral MoE使上述目标成为了或许，详见推特上的评论（twitter.com/Francis_YAO… ）

迈向100倍加快：全栈Transformer推理优化

如上表所示：

  • 功能：50B的MoE Mistral模型，其7B的密布部分在功能上接近于34B的Yi模型和67B的DeepSeek模型
  • 推理功率：MoE的密布部分为7B，其中包括前2的激活单元

Mistral MoE展现了下降更小模型成本的一起，取得与更大模型适当的功能的或许性。

迈向100倍加快：全栈Transformer推理优化

由Dmytro的Twitter帖子中的图表可知：

  • 当并发性较低时，大部分时刻都用于将两个激活的专家模型加载到内存中，这比密布层小，因而需要的内存拜访时刻更少，然后下降了时延。

换句话说，关于单个查询，MoE的时延低于密布层，由于咱们需要从内存中读取的参数更少。

  • 当并发性较高时，咱们进入了flop约束的状态，但由于MoE的激活少于密布层，因而吞吐量更高。

换句话说，关于许多并发查询，MoE具有更高的吞吐量，由于：

迈向100倍加快：全栈Transformer推理优化

随之而来的一个问题是，假设我已经练习了一个大型密布模型，是否能够将其转换为MoE模型？

  • MoE化（MoEfication）：将一个密布模型分解为MoE模型，使其
    • 具有小型模型的高效
    • 一起拥有大型密布模型的强大功能。
  • 参考文献：

3.2.2 提前退出

此处的要点是：关于简略词元，咱们不需要核算一切的Transformer层，核算部分层就够了，由于它们是简略词元。

迈向100倍加快：全栈Transformer推理优化

  • 关于一切词元，运用一个门控（gate）来确认是否提前退出或继续核算

参考文献：

3.2.3 分块解码

估测性解码

关键在于，逐词元解码不能充分利用SM的核算能力，因而咱们期望能够一次解码多个词元。以下文章值得参考：

  • Leviathan et. al. 2022. Fast Inference from Transformers via Speculative Decoding
  • Chen et. al. 2023. Accelerating Large Language Model Decoding with Speculative Sampling
  • Liu et. al. 2023. Online Speculative Decoding
  • Leviathan等，2022年。Fast Inference from Transformers via Speculative Decoding
  • Chen等，2023年。Accelerating Large Language Model Decoding with Speculative Sampling
  • Liu等，2023年。Online Speculative Decoding

迈向100倍加快：全栈Transformer推理优化

运用小型草稿模型（draft model）存在以下缺点：

  • 功能较弱，因而在一些具有挑战性的范畴（如编码）中，回绝率（rejection rate）或许较高。
  • 咱们需要在GPU中放入两个模型，但这已经耗尽了GPU内存。

因而，咱们期望将大模型作为自身的提案模型（proposal model），由于

  • 功能较强，因而能够下降回绝率。
  • 只需在内存中保存一个模型。

这便是咱们推广Medusa的原因：

迈向100倍加快：全栈Transformer推理优化

  • 运用多头一起解码多个词元
  • 运用大模型自身作为草稿模型

4

本文尚未涵盖的内容

  • 深水（Deep water）硬件和MLSys
    • 我对MLSys还很陌生，因而不得不向朋友询问自己的了解是否正确。因而，虽然我测验覆盖最重要的根底知识，但我对MLSys的了解依然仅是皮裘。
    • 还有许多涵盖硬件和MLSys的精彩文章，我将在文末添加相关参考文献
  • 超大模型的分布式推理。拜见Pope等人的Efficiently Scaling Transformer Inference（*arxiv.org/abs/2211.05…
    • 关键技能：模型分片（model sharding）、流水并行和张量并行
    • 继续批处理和类似的库DeepSpeed MII（*github.com/microsoft/D…
  • 对蒸馏的具体评论
  • 更进阶的分块解码

5

个人挑选

不幸的是，关于大部分模型架构论文而言，其收益和真实重要的作业大多来自MLSys，而非来自论文。实际上，许多花哨的论文并不切合实际，它们无法兼容SOTA MLSys的最新进展，如模型并行和Flash Attention。因而，亲爱的研讨同行们，是时候动手实践，仔细研讨真实的代码了。

但是，的确有一些来自建模方面的功能收益，以下是我最喜欢的技能，以备你的不时之需：

迈向100倍加快：全栈Transformer推理优化

6

定论

本文回忆了从GPU架构到MLsys办法，从模型架构到解码算法的全栈Transformer推理优化办法。能够看出，大部分功能提高都来自于一个准则的利用：Transformer推理受内存约束，因而咱们能够开释额外的核算能力/flops。其次，优化要么来自于优化内存拜访，比方Flash Attention和Paged Attention，要么来自于开释核算能力，比方Medusa和前向解码。

咱们信任MLSys和建模仍有许多改善空间。在行将到来的2024年，跟着模型变得更大、上下文变得更长以及跟着更多开源MoE（混合专家模型）、更高内存带宽和更大内存容量的硬件，以及具有更大DRAM和专用核算引擎的移动设备的露脸，将呈现更强大且人人可操作、可拜访的AI。一个新年代行将到来。

参考文献

欢迎 Star、试用 OneFlow 最新版别：

github.com/Oneflow-Inc…