KubeAI 大模型推理加速实践｜得物技术

一、背景

最近我们在生产环境批量部署了大模型专用推理集群，并成功让包括 70B 在内的大模型推理速度提升 50%，大幅缩减部署成本，稳定应用于生产环境。本文基于我们在部署大模型推理集群时的一些经验，分享一些有效提升大模型的推理速度方法。最后，我们在结尾处推荐了几个经过我们评测且表现优异的大模型推理框架。希望这些建议能帮助读者在项目中选择适合自己的推理框架。

OpenAI 的科学家 Hyung Won Chung 在 2023 年的公开演讲《Large Language Models》[8]中指出，大模型的某些能力仅在达到特定规模时才能显现，可见未来大模型的参数量肯定会越来越大，这也是大模型的发展趋势。随着参数量的增加，对大模型的推理速度要求越来越高，有哪些方法可以提高大模型的推理速度或吞吐量？

首先我们将探讨大模型的加速优化方向，随后文章将依据时间线，介绍一些业界内较为经典的实用大模型加速技术，包括但不限于“FlashAttention[1]”和“PageAttention[3]”等技术。

以下为按时间顺序业界的一些经典大模型推理加速技术，本文试图为读者提供一个按时间发展顺序的大模型加速方法综述。

除了上面提到的技术外，提高大模型推理速度的还有大模型的量化技术等，这里先不探讨，后面有机会，我们会单独发文章来介绍。

二、大模型发展面临的挑战

未来大模型的参数量肯定会越来越大，这也是大模型的发展趋势，对推理加速的要求会越来越高。

OpenAI 在其论文《Scaling Laws for Neural Language Models》[7]中介绍了大模型的扩展规则，这些规则阐释了模型能力与其规模之间的关系。具体来说，模型的能力强烈依赖于其规模，包括模型参数的数量，数据集的大小，以及训练过程中所需的计算量。此外，OpenAI 的科学家 Hyung Won Chung 在 2023 年的公开演讲《Large Language Models》[8]中指出，大模型的某些能力仅在达到特定规模时才能显现。

上图摘自 Hyung Won Chung 演讲中的 ppt[8]。图中主要表达一个观点，随着模型规模的增大，比如 GPT3 到 GPT4，模型的能力变的越来越强，甚至会出现新的能力。

但是随着模型的规模增大，大模型的推理速度将会逐渐降低，这是因为更多的参数量需要更多的 GPU 计算。推理速度的下降进一步带来更差的用户体验，因此如何对大模型推理加速变得越来越重要了。

三、大模型推理加速的优化方向

Llama2 的模型结构

我们先简单了解一下 Llama 2 模型系列的结构，参考自 Llama 2 的论文[9]。目前，像 Llama 系列这样的大多数生成式语言模型，主要采用了 Transformer 架构中的 Decoder 模块。在 Huggingface 平台上，这类模型结构通常被称作 CausalLM，即因果语言模型。

上图为 Llama2 大模型的结构，其中最核心的是注意力计算(Llama Attention)。这也是整个推理过程中最耗费时间的模块，后面的优化大部分都是基于 Attention 去实施的。为了更好的理解 Llama 2 大模型的结构，我们先简单对 Llama2 模型的整个推理过程进行拆解，不感兴趣同学可以直接跳过。

1. 用户向模型提交 Prompt 后，模型首先进行的操作是预测下一个字符(Token)，并将预测出的字符添加到输入中继续进行预测。这个过程会一直持续，直到模型输出一个停止符号(STOP token)，此时预测停止，模型输出最终结果。

2. 在生成下一个字符(Token)的过程中，模型需要执行 N 次的 Llama 解码器层(Llama Decoder Layer)计算。具体来说，Llama-2-7B 模型执行 32 次计算，而 Llama-2-13B 模型执行 40 次。

3. Llama 解码器层(Llama Decoder Layer)中最关键的计算环节是注意力(Llama Attention)的计算。大部分推理时间都消耗在 Attention 的计算上，因此多种优化技巧都旨在提高 Attention 计算的效率。

大模型推理的加速方向有哪些

从 Llama 2 模型的结构分析中，我们可以总结出大模型在推理计算过程中表现出以下特点：

1. 在整个推理过程中，最耗时的部分为注意力(Attention)计算。针对 Attention 的计算进行速度优化，可以显著提高整体推理性能。

2. 注意力(Attention)计算过程中，键值对缓存(KV Cache)占用了大量显存资源。以 13B 模型为例，处理一个 Prompt 序列大约需要 3GB 额外显存，并且这部分显存会被频繁地分配和释放，产生大量碎片，如果能减少显存碎片，也能提升大模型的吞吐。

3. 推理过程 GPU 需要处理和计算大量的参数。7B 模型拥有 70 亿参数，而 13B 模型则包含 130 亿参数，最新全球最强大模型 DBRX 更是高达 1300 亿参数，这需要高效地处理这些参数。这里也可以有优化空间。

针对上述三个特性，目前业界提出了多种有效的优化方法，典型如下：

1. FlashAttention-Attention 计算速度优化

FlashAttention[1]在不改变 Attention 算子的计算结果的前提下，提升 Attention 算子的计算速度。FlashAttention 在各种模型和任务上展示了显著的性能提升。例如，在 BERT-large、GPT-2 等模型上，相比于基线实现，FlashAttention 能够实现 15%到 3 倍的端到端加速。

2. PageAttention-KV Cache 显存管理优化

PageAttention[3]的目标是减少显存碎片，基于 PageAttention 的 VLLM 系统能够将流行的大型语言模型（LLM）的吞吐量提高到 10 倍以上，同时保持耗时分布平稳。

3. MOE-缩减推理时模型参数

MOE(Mixture of Experts)[4]目标是减少模型推理时参与计算的参数量。

实验效果：Mixtral 模型在多数基准测试中表现优于 Llama 2 70B 模型，并且其推理速度比后者快了 6 倍。该模型支持多种语言，具有强大的代码生成能力，并可以细化配置以遵循具体指令，从而在 MT-Bench 基准测试中取得了高分。

后面我们将针对上面的每个方向详细介绍。

四、FlashAttention-Attention 算子计算优化

FlashAttention 先后发表了两篇论文阐述对 Attention 算子的优化，包括 FlashAttention-1[1]与 FlashAttention-2[2]，我们以 FlashAttention-1[1]为例了解下他的优化原理。

我们先了解下 GPU 的内存分层结构，参考下图，图片来自论文 FlashAttention-1[1]。

GPU 的内存层次结构由三个主要部分组成：SRAM、HBM 和 DRAM，下面为 A100GPU 的参考配置。

SRAM（静态随机访问存储器）具有最快的访问速度（19TB/s），但其容量相对较小（仅 20MB）。

HBM（高带宽存储器）提供较大的存储空间（40GB）和高速的数据访问（1.5TB/s）。

DRAM（动态随机访问存储器），在这里特指 GPU 外部的主存，容量最大（超过 1TB），但访问速度最慢（12.8GB/s）。

从上述配置中可以看出，内存容量越小，处理速度就越快。

在传统的 Attention 计算过程中，大量的输入/输出操作都是通过访问 HBM 来完成的。FlashAttention 算法通过优化 Attention 计算流程，减少了对 HBM 的访问次数，以提高计算效率，所以它是一种 IO 感知的优化算法。

下图为 FlashAttention 的加速方法，来自论文 FlashAttention-1[1]

FlashAttention 利用了一个聪明的技巧来快速且内存高效地计算注意力机制，即它通过将输入数据分块（tiling）来避免一次性处理整个巨大的注意力矩阵，这通常需要大量的内存和计算资源。想象一下，我们有一个巨大的图书馆(矩阵)，而 FlashAttention 的方法就像是把图书馆里的书分成几个小堆，然后每次只处理一堆书。这样，我们就不需要一次性把所有书都拿出来放在桌子上（这需要很大的桌子和很多时间）。

具体来说，FlashAttention 在做矩阵计算的时候，通过将数据分块并利用 GPU 上的快速但容量较小的存储（SRAM）去计算，有效减少了对慢速但容量大的存储（HBM）的访问。这样不仅加快了计算速度，而且大幅减少了显存的需求。

通过减少对慢速存储的依赖，FlashAttention 能够显著提高模型训练的速度，同时保持或甚至提高模型的性能。例如，让 BERT-large 的训练比 MLPerf 1.1 的记录快 15%，GPT-2 训练速度是 HuggingFace 和 Megatron-LM 基线的三倍，长序列领域训练速度提升至 2.4 倍。

下图来自 huggingface 对 flash attention 介绍的 blog[14]，可以更好的理解 Flash Attention 对矩阵拆分的方式。

既然 Flash Attention 可以加速计算，那么支持 Flash Attention 计算的框架包括都有哪些，文章后半部我们会推荐一些比较优秀的推理框架。

五、PageAttention-显存管理优化

PageAttention[3]的概念最初由 VLLM 的作者 Woosuk Kwon 提出，它也是 VLLM 推理框架的最主要的优化策略。Woosuk Kwon 在其论文中介绍了如何通过 PageAttention 来解决大型语言模型（LLM）服务中的一个关键问题——在不增加延迟的情况下有效管理内存以提升吞吐量。

我们先了解下大模型在推理的情况下的内存结构分布，下图来自论文[3]。

这是一个在 NVIDIA A100 上服务一个拥有 13B 参数的大型语言模型的内存布局，13B LLM 推理显存占用分部，13B LLM 的参数占用 26G 显存，每个请求，KV Cache 会占用 12G 显存，随着 QPS 的增加，KVCache 会快速上升，并且会被频繁的分配与释放，系统会产生大量的显存碎片，如果不加处理，系统就会慢慢崩掉。

那么 VLLM 是如何通过 PageAttention 解决显存碎片的问题的呢？下图来自文章[14]，为 VLLM 的显存管理技术。

PageAttention 的工作原理是通过将键值缓存（KV 缓存）分割成固定大小的块（或“页面”），并允许这些块在内存中非连续地存储。这种方法灵感来源于操作系统的虚拟内存和分页技术，目的是为了更灵活和高效地管理内存资源。

在传统的注意力机制中，一个请求的 KV 缓存需要在内存中连续存储，这会导致两个主要问题：内存碎片化和无法高效共享内存。内存碎片化限制了批处理的大小，而无法共享内存则导致重复数据，浪费宝贵的内存资源。

PageAttention 通过以下步骤工作来解决这些问题：

1. 分割 KV 缓存：将每个请求的 KV 缓存划分为多个较小的块，这些块的大小是固定的，可以根据模型和硬件的具体需求进行调整。

2. 非连续存储：与传统 KV 缓存块在内存中连续存储不同，PageAttention 允许这些块在物理内存中非连续地分布。这样，就可以根据实际需要动态地分配和回收内存块，减少内存浪费。

3. 动态管理：通过类似于操作系统中虚拟内存管理的方式，PageAttention 动态地管理这些内存块。系统可以根据当前的内存使用情况，按需分配或释放 KV 缓存块，从而优化内存使用。

4. 内存共享：PageAttention 还支持在不同请求之间或同一个请求中的不同序列之间共享 KV 缓存块。这种共享是灵活的，可以基于块级别进行，进一步减少内存使用和提高效率。

通过这种方式，PageAttention 允许 LLM 服务系统在保持相同延迟的情况下，通过减少内存浪费和提高内存共享，显著提高处理请求的吞吐量。

通过 PageAttention 的优化，VLLM 对 LLaMA 7B 与 13B 的吞吐量提升了 10 倍以上，下图来自文章[11]。

六、MOE-缩减推理时模型参数

最近发布的全球最强开源大模型 1300 亿参数的 DBRX，以及 Mistral 的 8x7B 开源大模型都是基于 MOE 架构的。为什么参数量越大的模型越要使用 MOE 架构呢？我们以 Mistral 的 8x7B 开源大模型为例，介绍下 MOE 架构在性能方面的优势。

说到 MOE 大模型，我们先对比下普通大模型与 MOE 大模型在结构上的区别，参考上图。在 MOE 大模型中，把大模型的参数分成了 8 个小组外加一个路由器，每个小组我们称作专家组。当请求过来的时候，MOE 大模型则先有路由器从 8 个专家组中选择两个，只有这两个专家组参与了计算。而对比普通大模型，则需要所有参数都参加 GPU 计算。

所以 MOE 大模型要比同等级的普通大模型推理速度快四倍左右。

我们来看下 Mistral MOE 的实现，Mistral MOE 是由 mistral.ai 发布的 8*7B 大模型[12]，下图来自论文[12]，是其 8*7B 大模型的专家层的结构。

Mixtral 8x7B 是一个稀疏混合专家（Sparse Mixture of Experts, SMoE）语言模型，它基于 Mistral 7B 的架构，但每一层都由 8 个前馈块（即专家）组成。在处理每个令牌时，每层的一个路由网络会选择两个专家来处理当前状态并结合它们的输出。虽然每个令牌只与两个专家交互，但在每个时间步骤中选取的专家可以不同，因此每个令牌可以接触到 47B 的参数，但在推理过程中只使用 13B 的活跃参数。

Mixtral 在多项基准测试上展示了其卓越的性能，尤其是在数学、代码生成和多语言理解方面。相比于 Llama 2 70B 和 GPT-3.5，Mixtral 在大多数评估指标上表现出类似或更优的性能。具体来说，Mixtral 使用的活跃参数（13B）比 Llama 2 70B（70B）少 5 倍，但在几乎所有类别中的表现都更好或相当。

MOE 大模型实现了增加参数量的同时，推理速度并不降低，是未来大模型的发展趋势。

七、Tensor parallelize-张量并行

如果你有多卡 GPU，可以采用张量并行进一步加速大模型的推理速度。

想象一下，你有一本非常厚的书，你想一次性复印整本书，但是你的复印机一次只能复印几页。这时，你可以把这本书分成几个部分，每个部分分别复印，最后再把所有复印好的部分按顺序拼接起来，这样就完成了整本书的复印。

在张量并行中，我们要处理的大模型就像是那本厚书，而 GPU 则像是复印机。因为单个 GPU 无法一次处理整个大模型，我们就需要把模型（在这个例子中是权重张量）分成几个部分，让不同的 GPU 分别处理（相当于复印书的不同部分）。在处理输入数据时，就像是把书的每一页分别复印，然后再把复印好的各个部分拼接起来，形成完整的输出结果。

这样，通过分担工作，多个 GPU 协同完成了一个本来单个 GPU 无法完成的大任务。这就是张量并行的工作方式，它让我们能够处理那些非常大的模型。

图片来自文章[13]

张量并行技术用于将大模型分布式地部署在多个 GPU 上。以矩阵乘法来打个比方，当输入张量与第一个权重张量进行矩阵相乘时，这个操作可以视作先将权重张量按列进行分割，接着将分割后的每列分别与输入张量相乘，然后把这些乘积结果合并。这些合并后的输出会被从 GPU 中导出，并聚合以形成最终的输出结果，过程上图，参考文章[13]。

八、推理框架推荐

在前文中，我们探讨了几种加速和优化技术，诸如 Flash Attention、Page Attention、MOE 以及张量并行技术。接下来，基于我们自身的实际操作和评估，我们将向您推荐一些当前表现较为出色的推理框架。

九、总结与展望

在本文中，我们深入探讨了一系列旨在提升大模型推理速度的技术和方法，包括但不限于 Flash Attention、Page Attention、MOE 以及张量并行技术。通过在生产环境中批量部署专用大模型推理集群，我们成功地将包括 70B 规模模型在内的推理速度降低了 50%，稳定地应用这些技术于生产环境，从而证明了这些优化方法的有效性和实用性。

随着大型模型在各个领域的应用越来越广泛，如何有效地提升推理速度、降低推理成本成为了一项挑战。我们的实践不仅展示了目前可用的一些加速技术，还基于我们的经验，推荐了几款经过评测表现优秀的大模型推理框架。这些建议旨在帮助读者在面对众多选择时，能够挑选出最适合自己需求的推理框架。

展望未来，随着技术的不断进步和新算法的不断涌现，我们相信还会有更多的加速优化技术被开发出来，进一步推动大模型推理效率的提升。最后，我们也期待未来有机会深入探讨和介绍更多提升大模型推理速度的新技术和方法。

参考资料

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness(https://arxiv.org/abs/2205.14135)

[2] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning(https://arxiv.org/abs/2307.08691)

[3] Efficient Memory Management for Large Language Model Serving with PagedAttention(https://arxiv.org/abs/2309.06180)

[4] mixtral-of-experts(https://mistral.ai/news/mixtral-of-experts/)

[5] Mixtral of Experts(https://arxiv.org/abs/2401.04088)

[6] MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads(https://arxiv.org/pdf/2401.10774.pdf)

[7] Scaling Laws for Neural Language Models(https://arxiv.org/pdf/2001.08361.pdf)

[8] Hyung Won Chung(OpenAI), Large Language Models (in 2023) , talked at Seoul National University

[9] Llama 2: Open Foundation and Fine-Tuned Chat Models(https://arxiv.org/abs/2307.09288)

[10] Attention Is All You Need(https://arxiv.org/pdf/1706.03762.pdf)

[11] https://blog.vllm.ai/2023/06/20/vllm.html

[12] https://arxiv.org/pdf/2401.04088.pdf

[13] https://huggingface.co/docs/text-generation-inference/en/conceptual/tensor_parallelism

[14] https://huggingface.co/docs/text-generation-inference/en/conceptual/flash_attention

[15] https://blog.vllm.ai/2023/06/20/vllm.html

*文/ linggong

本文属得物技术原创，更多精彩文章请看：得物技术

未经得物技术许可严禁转载，否则依法追究法律责任！