大型语言模型的推理演算_大模型推理过程

作者｜kipply

翻译｜杨婷、徐佳渝、贾川‍‍

本文详细阐述了大型语言模型推理性能的几个基本原理，不含任何实验数据或复杂的数学公式，旨在加深读者对相关原理的理解。此外，作者还提出了一种极其简单的推理时延模型，该模型与实证结果拟合度高，可更好地预测和解释Transformer模型的推理过程。

为了更好地阅读本文，读者需了解一些Transformer模型的相关先验知识，比如《图解Transformer》的大部分内容。另外，了解与本文相关的参数计数文章也能更好地帮助读者理解本文内容。本文主要包括以下内容：

• kv 缓存 (kv cache) 解释了在推理过程中缓存自注意力向量所带来的性能优化效果，以及可能导致的权衡（tradeoff）以及容量成本问题。

• 容量（capacity）考虑了kv缓存的存储成本以及模型权重的存储成本之间的联系，并解释了容量大小对模型性能的影响。

• 模型并行可帮助我们理解张量并行，以明确通信成本。

• 时延计算需要从其他概念中获得理解，并创建用于确定推理速度底线（floorline）的方程。

• 批大小（batch size）对性能的影响以及最优批大小为多少。

• 通过transformer blocks执行flops（每秒浮点运算次数）计数操作，可以识别对flops速度有实质性贡献的操作。

• 中间内存成本涵盖了激活（即激活函数的输出结果）占用额外内存，以及一些真实基准测试中的内存带宽成本。

• 对比真实基准测试是指将计算出的内容与英伟达 FasterTransformer基准测试结果进行对比，并确定其中的差异。

（本文经授权后由OneFlow编译发布，译文转载请联系OneFlow获得授权。原文：https://kipp.ly/blog/transformer-inference-arithmetic/）

1

kv 缓存

采样时，Transformer模型会以给定的prompt/context作为初始输入进行推理（可以并行处理），随后逐一生成额外的token来继续完善生成的序列（体现了模型的自回归性质）。在采样过程中，Transformer会执行自注意力操作，为此需要给当前序列中的每个项目（无论是prompt/context还是生成的token）提取键值（kv）向量。这些向量存储在一个矩阵中，通常被称为kv缓存或者past缓存（开源GPT-2的实现称其为past缓存）。past缓存通常表示为：[batch, 2, num_heads, seq_len, features]。

kv缓存是为了避免每次采样token时重新计算键值向量。利用预先计算好的k值和v值，可以节省大量计算时间，尽管这会占用一定的存储空间。每个token所存储的字节数为：

 

第一个因子2表示k和v这两个向量。在每一层中我们都要存储这些k，v向量，每个值都为一个矩阵。然后再乘以2，以计算每个向量所需的字节数（在本文中，我们假设采用16位格式）。

我们乘以token嵌入（token embeddings）得到的权重为，其中每个token嵌入为。这样，我们就可以算出所有层的k和v需进行的浮点运算次数为：

将乘以需要进行次浮点运算。另一个2表示我们需要重复两次这样的操作，一次用于计算k和一次用于计算v，然后再重复所有层数。

矩阵乘法（matmul）中的浮点运算次数为多少？

矩阵-向量（matrix-vector）乘法的计算公式是，其中，。对于矩阵-矩阵（matrix-matrix）乘法，计算公式是 ，其中，。因子十分重要，因为它反映了矩阵乘法中由乘法和加法组成的组合方式，即“乘法(1)-加法(2) 操作组合”。更多内容见讲义（lecture notes）。

这意味着对于一个520亿参数的模型来说 （以Anthropic中的模型为例， ，），其浮点运算次数为：

假设有一个A100 GPU，其每秒可执行的浮点运算次数为，其内存带宽可达字节/秒。以下数字仅涉及kv权重及计算的数值：

Flops vs 内存有界性（Boundedness）

英伟达使用了数学带宽这个术语，我觉得这个术语真的很可爱。从技术上讲，这种描述存在于每个内核中，但可以抽象为操作组。 

Flops vs 内存有界性是Transformer推理和深度学习优化的常见问题。为了完成所需计算，通常需要加载权重，加载过程会占用内存带宽。假设通过加载权重已经得到了很好的优化，那么我们可以在加载权重的同时开始计算。

在这种情况下，flop bound意味着一段时间内存中没有任何数据传输；memory bound则意味着没有进行任何计算操作。

英伟达使用数学带宽（math bandwidth）来描述该情况，我觉得相当有意思。从技术上讲，这种划分通常是指每个