开源大模型推理与性能优化实战

我们希望在 LLM 推理中充分利用 GPU 的性能。为此，我们需要了解推理是受计算限制还是受内存限制，以便在正确的地方进行优化。通过计算给定 GPU 上可能的每字节操作数并将其与模型注意力层的算术强度进行比较，我们可以发现瓶颈所在：计算还是内存。使用这些信息，我们可以为模型推理选择合适的 GPU，如果使用场景允许，可以采用批处理等技术来更好地利用 GPU 资源。

在 ML 模型 API 和裸金属 GPU 之间存在许多抽象层。建立对这些抽象的深刻心智模型有助于你在推理过程中控制成本、提高性能，充分利用 GPU 的潜力，获得最大收益。

本指南将帮助你理解 Transformer 推理分析的数学原理。作为具体例子，我们将贯穿全文运行 Llama 2 在 A10 GPU 上的场景。我们将涵盖：

• 阅读关键 GPU 规格以发现硬件的性能指标
• 计算 GPU 的每字节操作数（ops:byte）比率
• 计算 LLM 的算术强度
• 对比 ops:byte 与算术强度，判断推理是受计算限制还是受内存限制
• 优化受内存限制和受计算限制的推理的实际策略

此外，我们还会用真实基准测试来验证我们的计算。阅读完本指南后，你将理解模型服务的主要瓶颈以及如何缓解它们。

一个常见的工作负载是运行 70 亿参数的 LLM，比如 Llama 2 或 Mistral。在生产环境中运行 LLM 需要 GPU，但我们应该选择哪一个呢？

假设我们选择 A10，它是轻量级 T4 和强大但昂贵的 A100 之间的良好权衡。以下是 A10 的关键规格。

在推理方面，我们关注其中的三个数字：

FP16 Tensor Core：这是我们的计算带宽。我们有 125 TFLOPS（万亿次浮点运算/秒）的可用计算能力，用于半精度（也称为 FP16）模型。半精度是一种占用 16 位的二进制数字格式，相对于占用 32 位的全精度二进制格式。对于许多 ML 应用，使用半精度是一个实用的选择，因为它可以减少内存占用而不损失精度。在本文中，我们忽略数据表中与稀疏性相关的值（用星号表示）。

GPU 内存：我们可以通过将参数数量（以十亿为单位）乘以 2 来快速估计模型的大小（以 GB 为单位）。这种方法基于一个简单的公式：在半精度下，每个参数使用 16 位（或 2 字节）的内存，因此内存使用量（GB）大约是参数数量的两倍。例如，一个 70 亿参数的模型将占用约 14 GB 的内存。为什么这很重要？因为 A10 的 24 GB VRAM，我们可以舒适地运行一个 70 亿参数的模型，仍有约 10 GB 的内存作为缓冲区。这个多余的内存在模型执行中起着重要作用，我们稍后会进一步说明。

GPU 内存带宽：我们可以以 600 GB/s 的速度从 GPU 内存（也称为 HBM 或高带宽内存）将数据传输到片上处理单元（也称为 SRAM 或共享内存）。

计算每字节操作数（ops:byte）比率

使用这些数字，我们可以计算硬件的 ops:byte 比率。这告诉我们对于每字节的内存访问，我们能完成多少浮点运算/秒（FLOPS）。

根据规格表中的数字，我们计算 A10 的 ops:byte 比率：

ops_to_byte_A10
    = compute_bw / memory_bw 
    = 125 TF / 600 GB/S  
    = 208.3 ops / byte

这意味着为了充分利用计算资源，我们必须在每字节的内存访问中完成 208.3 次浮点运算。

如果我们发现自己只能完成少于 208.3 次每字节的操作，系统性能就受内存限制。这实质上意味着系统的速度和效率受到数据传输速率或其可处理的输入输出操作的限制。

如果我们想做超过 208.3 次每字节的浮点运算，系统就是受计算限制的。在这种状态下，我们的有效性和性能受限不是由内存，而是由芯片拥有的计算单元数量。

了解我们是受计算限制还是受内存限制是至关重要的，这样我们才能知道在哪里集中优化工作。

计算算术强度

为了确定我们是受计算限制还是受内存限制，我们需要计算 70 亿参数 LLM 的算术强度，然后将其与我们刚刚为 GPU 计算的 ops:byte 比率进行比较。算术强度是算法执行的计算操作数除以它需要的字节访问次数，是一种与硬件无关的度量。

我们 70 亿参数 LLM 中计算成本最高的部分是注意力层，它确保下一个 token 的预测根据前一个 token 的相关性加权。因为注意力层是推理中计算需求最高的部分，我们将在那里计算算术强度。

理解注意力层需要更具体地了解模型在底层的工作方式。从 Transformer 采样时，有两个阶段：

Prefill：在第一个阶段，模型并行地处理你的提示 token，填充键值（KV）缓存。KV 缓存可以被视为模型的状态，嵌套在注意力操作中。在 Prefill 阶段，没有 token 被生成。

自回归采样：在第二个阶段，我们利用当前状态（存储在 KV 缓存中）来采样和解码下一个 token。我们付出少量的存储代价，以避免为每个新 token 重新计算缓存。如果没有 KV 缓存，每个后续 token 的采样会花费更长时间，因为我们必须将所有先前看到的 token 通过模型传递。

自回归采样

在第二阶段，我们利用当前状态（存储在 KV 缓存中）来采样并解码下一个 token。我们付出一点存储空间的代价，以避免为每个新 token 重新计算缓存。如果没有 KV 缓存，每个后续 token 采样会花费更长的时间，因为我们必须将所有之前看到的 token 传入模型。

拆解注意力方程

下面是 LLM 推理中的注意力方程。我们将逐步遍历这个方程，确定哪些需要内存移动，哪些需要计算操作，然后我们可以比较这两者，找到我们所寻找的算术强度。

FlashAttention 论文的作者为标准注意力算法提供了一个很好的实现。这个框架将使我们更容易计算算法中的内存和计算。

眼尖的读者可能会注意到这个算法省略了除以 sqrt(d_k) 的缩放。这是一个次要因素，我们可以安全地忽略。

所有三个步骤都遵循相同的模式：从内存加载值，执行计算，并将该计算的结果存储到内存。在算法中：

N 是 LLM 的序列长度，它设置上下文窗口。对于 Llama 2 7B，N = 4096。

N 是 LLM 的序列长度，它设置上下文窗口。

对于 Llama 2 7B，N = 4096。

对于 Llama 2 7B，N = 4096。

d 是单个注意力头的维度。对于 Llama 2 7B，d = 128。

d 是单个注意力头的维度。

对于 Llama 2 7B，d = 128。

对于 Llama 2 7B，d = 128。

Q、K 和 V 都是用于计算注意力的矩阵。它们的维度是 N 乘以 d，在我们的情况下是 4096x128。

Q、K 和 V 都是用于计算注意力的矩阵。

它们的维度是 N 乘以 d，在我们的情况下是 4096x128。

它们的维度是 N 乘以 d，在我们的情况下是 4096x128。

S 和 P 都是在方程中计算的矩阵。它们的维度是 N 乘以 N，在我们的情况下是 4096x4096。

S 和 P 都是在方程中计算的矩阵。

它们的维度是 N 乘以 N，在我们的情况下是 4096x4096。

它们的维度是 N 乘以 N，在我们的情况下是 4096x4096。

O 是具有注意力计算结果的输出矩阵。O 是一个 N 乘以 d 的矩阵，在我们的情况下是 4096x128。

O 是具有注意力计算结果的输出矩阵。

O 是一个 N 乘以 d 的矩阵，在我们的情况下是 4096x128。

O 是一个 N 乘以 d 的矩阵，在我们的情况下是 4096x128。

HBM 是高带宽内存。从数据手册，我们知道 A10 上有 24 GB 的 HBM，以 600 GB/s 的速率运行。

HBM 是高带宽内存。

从数据手册，我们知道 A10 上有 24 GB 的 HBM，以 600 GB/s 的速率运行。

从数据手册，我们知道 A10 上有 24 GB 的 HBM，以 600 GB/s 的速率运行。

考虑到这些数字，让我们为实现中的每一行分解标准注意力算法，然后求和找到运行该算法的总计算和内存成本。

我们通过对第一列和第三列求和来计算总内存移动（从内存加载和存储到内存）。

total_memory_movement_in_bytes:
    = (2 * 2 * (N * d)) + (2 * (N * N)) + (2 * ((N*N) + (N * d))) + (2 * (N * N)) + (2 * (N * N)) + (2 * (N * d))
    = 8N^2 + 8Nd bytes

并通过对第二列求和来计算总计算（对加载数据的计算）。

total_compute_in_floating_point_ops: 
    = ((2 * d) * (N * N)) + (3 * (N * N)) + ((2 * N) * (N * d))
    = 4(N^2)d + 3N^2 ops

算术强度可以按如下方式计算。

arithmetic_intensity_llama 
    ~= total compute / total memory movement
    = 4d(N^2) + 3N^2 ops / 8N^2 + 8Nd bytes
    = 62 ops/byte for Llama 2 7B

发现我们的推理瓶颈

我们的 Llama 2 7B 的算术强度是 62 ops/byte，这远小于我们 A10 的 208.3 ops/byte 比率。

因此，在自回归阶段，我们的模型是内存受限的。换句话说，在我们将单个字节从内存移动到计算的时间内，我们可以完成许多倍于仅在该字节上进行的计算。

这是一个问题。我们花费好钱来保持 GPU 运行，但没有充分利用可用的计算资源。

在 GPU 上对内存受限的过程进行批处理

一个解决方案是利用片上内存来批量运行通过模型的前向传递。换句话说，我们可以等待几百毫秒来积累几个请求，然后在一次传递中一起运行它们，而不是贪心地处理到达的请求。这使我们能够重用已经加载到 GPU 的 SRAM 中的模型部分。

批处理通过对相同数量的内存加载和存储执行更多计算来增加模型的算术强度，这反过来减少了模型受内存限制的程度。

我们的批次最大可以有多大？回想一下，在我们的 A10 上加载 7B 参数模型后，我们还有 10 GB 的内存：

24 GB - (2 * 7GB) = 10GB

现在，问题是我们一次可以在备用 GPU 内存中容纳多少个序列？

要计算这个值，我们需要回到 KV 缓存。回想一下，在注意力层的 prefill 步骤中，我们基于提示或输入序列填充 KV 缓存。

KV 缓存包含我们在注意力计算期间使用的矩阵 K 和 V。要计算 KV 缓存的大小，我们需要回顾一些之前定义的值和几个新的变量：

d，可以记为 d_head，是单个注意力头的维度。对于 Llama 2 7B，d = 128。

d，可以记为 d_head，是单个注意力头的维度。

对于 Llama 2 7B，d = 128。

对于 Llama 2 7B，d = 128。

n_heads 是注意力头的数量。对于 Llama 2 7B，n_heads = 32。

n_heads 是注意力头的数量。

对于 Llama 2 7B，n_heads = 32。

对于 Llama 2 7B，n_heads = 32。

n_layers 是注意力块出现的次数。对于 Llama 2 7B，n_layers = 32。

n_layers 是注意力块出现的次数。

对于 Llama 2 7B，n_layers = 32。

对于 Llama 2 7B，n_layers = 32。

d_model 是模型的维度。d_model = d_head * n_heads。对于 Llama 2 7B，d_model = 4096。

d_model 是模型的维度。d_model = d_head * n_heads。

对于 Llama 2 7B，d_model = 4096。

对于 Llama 2 7B，d_model = 4096。

值得注意的是，d_model 与 N（上下文窗口长度）相同是巧合。正如 Llama 论文所示，其他大小的 Llama 2 有更大的 d_model（见"维度"列）。

在半精度 (FP16) 下，每个浮点数占用 2 个字节。有 2 个矩阵，要计算 KV 缓存大小，我们将两者都乘以 n_layers 和 d_model，得到以下方程：

kv_cache_size
    = (2 * 2 * n_layers * d_model) bytes/token
    = (4 * 32 * 4096) bytes/token
    = 524288 bytes/token
    ~ 0.00052 GB/token

鉴于 KV 缓存每个 token 需要 524288 字节，我们的 KV 缓存最多能容纳多少个 token？

kv_cache_tokens
    = 10 GB / 0.00052 GB/token
    = 19,230 tokens

我们的 KV 缓存可以容纳 19,230 个 token。因此，对于 Llama 2 标准序列长度的 4096 个 token，我们的系统有足够的带宽同时处理 4 个序列的批次。

总之，为了充分利用我们为之付费的计算能力，我们希望在推理期间一次批处理 4 个请求，以填充我们的 KV 缓存。这将增加我们的吞吐量。

如果你使用 LLM 异步处理大量文档队列，批处理是一个好主意。你将处理队列的速度比单独处理每个元素快得多，并可以安排推理调用，使它们快速填充批次，最大程度地减少对延迟的影响。

为 LLM 推理评估 GPU

在某些情况下，批处理可能没有意义。例如，如果你正在构建面向用户的聊天机器人，你的产品对延迟的敏感性要高得多，因此在运行推理之前无法等待批次填充。在这种情况下我们应该做什么？

一个选择是认识到我们将无法充分利用 GPU 的片上内存，并改用规格更小的 GPU。例如，我们可以迁移到 T4 GPU，它有 16 GB 的 VRAM。这仍然可以容纳我们的 7B 参数模型，但剩余容量要少得多——仅 2 GB——用于批处理和 KV 缓存。

然而，T4 GPU 通常比 A10 慢。而 A100，虽然更强大，但也更贵。我们可以通过计算推理时间的一些简单下界来量化这种差异。

在每个 GPU 上生成单个 token

回想一下，在生成的自回归部分中，如果我们的批大小为 1，我们受内存带宽限制。让我们使用以下方程快速计算生成单个 token 需要多长时间：

time/token = total number of bytes moved (the model weights) / accelerator memory bandwidth

在 T4 上：(2 * 7B) bytes / (300 GB/s) = 46 ms/token

在 T4 上：(2 * 7B) 字节 / (300 GB/s) = 46 ms/token

在 A10 上：(2 * 7B) 字节 / (600 GB/s) = 23 ms/token

在 A10 上：(2 * 7B) 字节 / (600 GB/s) = 23 ms/token

在 A100 SXM 80 GB 上：(2 * 7B) 字节 / (2039 GB/s) = 6 ms/token

在 A100 SXM 80 GB 上：(2 * 7B) 字节 / (2039 GB/s) = 6 ms/token

这些估计值表明 A10 的速度是 T4 的两倍，而 A100 的速度几乎是 A10 的四倍（因此是 T4 的八倍）。

这些数字只是近似值，因为它们假设在推理过程中 GPU 内的通信开销为零、每次前向传播的开销为零，且计算期间的并行化完美。

在每个 GPU 上使用批处理提示令牌进行预填充

我们也可以计算预填充阶段所需的时间，假设我们将所有提示令牌批处理为单次前向传播。假设提示有 350 个令牌（为了简单起见），且限制瓶颈是计算而不是内存。

预填充时间 = 令牌数量 * (参数数量 / 加速器计算带宽)

在 T4 上：350 * (2 * 7B) FLOP / 65 TFLOP/s = 75 ms

在 T4 上：350 * (2 * 7B) FLOP / 65 TFLOP/s = 75 ms

在 A10 上：350 * (2 * 7B) FLOP / 125 TFLOP/s = 39 ms

在 A10 上：350 * (2 * 7B) FLOP / 125 TFLOP/s = 39 ms

在 A100 SXM 80 GB 上：350 * (2 * 7B) FLOP / 312 TFLOP/s = 16 ms

在 A100 SXM 80 GB 上：350 * (2 * 7B) FLOP / 312 TFLOP/s = 16 ms

估计每个 GPU 上的总生成时间

假设我们允许 150 个完成令牌（忽略任何停止令牌），总生成时间如下所示。

总生成时间 = 预填充时间 + 令牌数量 * 每令牌时间

在 T4 上 = 75 ms + 150 令牌 * 46 ms/token = 6.98 s

在 T4 上 = 75 ms + 150 令牌 * 46 ms/token = 6.98 s

在 A10 上 = 39 ms + 150 令牌 * 23 ms/token = 3.49 s

在 A10 上 = 39 ms + 150 令牌 * 23 ms/token = 3.49 s

在 A100 SXM 80 GB 上：16 ms + 150 令牌 * 6 ms/token = 0.92s

在 A100 SXM 80 GB 上：16 ms + 150 令牌 * 6 ms/token = 0.92s

考虑 GPU 价格因素，我们可以查看推理速度与成本之间的近似权衡。具体数值会因计算中使用的令牌数量而略有不同。

NVIDIA GPU 上 Llama 2 7B Chat 的基准测试

在物理课堂上，许多家庭作业问题都设定在无摩擦的真空环境中。我们所做的估计也处于类似的理论场景中。现在让我们将它们与现实世界的数据对比。

我们计算的数字是推理时间的低估，因为它们没有考虑 GPU 通信成本、网络延迟和不完美利用等因素。在现实世界中，这些因素可能使推理时间增加到两倍。

实际基准在 Baseten 上使用 Llama 2 7B Chat，测量适用 GPU 上的端到端延迟，除了使用 ExLlama V2 外没有任何其他优化。

我们使用一个标准化提示，根据 Llama 2 分词器的精确输入长度为 350 个令牌。提示内容如下所示。要总结的文本取自 Paul Graham 的《How to Do Great Work》。

Write a concise summary of the following:

"The first step is to decide what to work on…[excerpt continues long enough to make the prompt exactly 350 tokens]"

CONCISE SUMMARY:

我们还设置最多 150 个输出令牌，没有停止令牌，这保证了每个查询总共将消费 500 个令牌。因此我们可以使用上面的计算来预测这些查询需要多长时间。

预测结果以深绿色显示，而测量的生成时间以较浅的色调显示。预期的计算结果在正确的范围内，但总是偏低，因为它们在工作并行化方面做了保守的假设，并忽视了 GPU 内的通信延迟。

使用 Transformer 数学优化 LLM 模型推理

我们希望在 LLM 推理期间充分利用计算容量，但当我们受内存限制时，无法做到这一点。通过计算给定 GPU 上可能的每字节操作数，并将其与模型注意力层的算术强度进行比较，我们可以判断自己是受内存限制还是受计算限制。

当受内存限制时，批处理可以让我们充分利用计算容量，尽管对许多延迟敏感的用例来说，批处理是不可行的。当我们有严格的延迟要求时，可以使用类似的计算来估计哪些 GPU 能够满足我们的需求。

尽管这些理论计算很有用，但将它们与实际基准对比始终是必要的，以考虑通信成本和网络延迟等因素。

深入研究 LLM 推理的工作原理是令人着迷的，总有更多的内容值得探索。以下是一些优秀的资源，供你了解更多：

Transformer inference arithmetic on kipp.ly

Why GPT-3.5 is (mostly) cheaper than Llama 2 on cursor.sh

Making Deep Learning Go Brrrr From First Principles on horace.io

GPU Performance Background User's Guide by NVIDIA

订阅我们的新闻通讯

及时了解模型性能、推理基础设施等最新信息。