在 PyTorch/XLA 上使用 LLaMA 65B 实现超低推理延迟的路径

背景与最新技术

在自然语言处理 (NLP) 领域，语言模型旨在根据过去输入的标记序列生成一个标记（例如单词）。大型语言模型 (LLM) 是该领域最新的深度学习创新，旨在以类似人类的方式生成文本。这些模型通常使用Transformer来提高它们对大量输入标记序列的注意力。

LLaMA，由Meta AI开源，是一个强大的基础 LLM，训练数据超过 1 万亿标记。LLaMA 与许多同类最佳模型（如GPT-3、Chinchilla、PaLM）具有竞争力。LLaMA (13B) 优于 GPT-3 (175B)，这突出显示了它从每个模型参数中提取更多计算能力的能力。

在这篇博客文章中，我们使用 LLaMA 作为示例模型来演示 PyTorch/XLA 在 LLM 推理方面的能力。我们讨论了本文中讨论的计算技术和优化如何将由 Google Cloud TPU v4 (v4-16) 提供支持的 65B 参数 LLaMA 模型的推理延迟缩短 6.4 倍。

模型概述

我们演示了 PyTorch/XLA 在 Meta 最新 LLM LLaMA 上的性能。我们展示了对一系列常见 LLaMA 配置的性能优化。请注意，175B 参数模型配置在公共领域中不存在。对于下面提到的 175B 参数模型，我们将OPT 175B 模型配置应用于 LLaMA 代码库。除非另有说明，在所有配置中，我们使用max_seq_len=256和dtype=bfloat16表示权重和激活。

表 1：本文探讨的模型配置

LLaMA	模型超参数
参数数量	维度	N 头	N 层	最大序列长度
7B	4,096	32	32	256
33B	6,656	52	60	256
65B	8,192	64	80	256
175B	12,288	96	96	256

LLM 的性能挑战

LLM 具有一些特性，使其难以进行编译器优化。(a) LLM 使用自回归解码来生成基于先前标记的下一个标记；这意味着提示张量和缓存具有动态形状。(b) LLM 必须处理可变输入提示长度，而不会因输入张量形状变化而触发重新编译；输入张量必须正确分桶和填充以避免重新编译。(c) LLM 通常需要比单个 TPU（或 GPU）设备支持的内存更多。需要模型分片方案才能将模型适应分布式计算架构。例如，一个具有 65B 参数的 LLaMA 模型可以适应 v4-16 Cloud TPU，这与 8 个 A100 GPU 相当。(d) 在生产中运行 LLM 可能很昂贵；提高每总拥有成本性能 (Perf/TCO) 的一种方法是通过量化；量化可以潜在地降低硬件要求。

PyTorch/XLA 中的推理技术栈

我们的目标是为 AI 社区提供高性能推理堆栈。PyTorch/XLA 与TorchDynamo、PjRt、OpenXLA以及各种模型并行方案集成。TorchDynamo 消除运行时跟踪开销，PjRt 实现高效的主机-设备通信；PyTorch/XLA 可追溯集合通过TorchDynamo在 LLaMA 上实现模型和数据并行。要尝试我们的结果，请使用我们的自定义torch、torch-xla轮子来重现我们的LLaMA 推理解决方案。PyTorch/XLA 2.1 将默认支持本文讨论的功能。

并行计算

FairScale分片

LLaMA 使用 FairScale 模型分片 API (fairscale.nn.model_parallel.layers)。我们使用 PyTorch/XLA 通信集合 (CC) 操作（例如all-reduce）构建了此 API 的等效表示，以在加速器之间通信程序状态（例如激活）。TorchDynamo 目前不完全支持捕获 CC 操作（又名可追溯集合）。如果没有此支持，TorchDynamo FX 图将在每次设备通信处（即每个模型层）被切断。图切断会导致性能损失，因为底层 XLA 编译器失去了完整的图优化机会。为了解决这个问题，我们通过将调度程序集合集成到我们现有的 CC API 中来提供 PyTorch/XLA 可追溯集合。不同之处在于，鉴于 PyTorch/XLA 的惰性执行性质，我们不需要在集合之后插入c10d.wait()操作。通过支持可追溯集合，PyTorch/XLA 允许在 TorchDynamo 中生成奇异 FX 图。

PyTorch/XLA 上的自回归解码

LLM 需要自回归解码以将前一个单词作为提示来预测下一个标记。自回归解码导致无边界的动态形状问题，这反过来又导致每个提示的重新编译。我们优化了 LLaMA 自回归解码器，使其使用固定形状操作，在每次标记生成期间就地更新 KV 缓存、输出序列和注意力掩码。通过结合填充、掩码和索引操作，我们避免了过度的图重新编译，从而实现了高效的自回归解码。

KV 缓存优化

LLaMA 使用 KV 缓存实现自回归解码。对于每个生成的标记，KV 缓存存储每个 Transformer 层的注意力键/值激活。因此，在解码新标记时，不再需要重新计算先前标记的键/值。

在 LLaMA 中，KV 缓存张量切片是就地更新的；这会导致每次生成标记时发生重新编译事件。为了解决这个问题，我们使用索引张量和tensor.index_copy()操作来替换就地切片更新。注意力掩码和输出序列也受益于相同的优化。

输入提示优化

可变长度输入提示在 LLM 应用程序中很常见。此属性导致输入张量形状动态性，进而导致重新编译事件。在处理提示以填充 KV 缓存时，我们要么 (a) 逐标记处理输入提示，要么 (b) 在一次迭代中处理整个提示。每种方法的优缺点如下：

预编译 1 个图并逐标记处理提示
- 实用：在预热期间编译 1 个图
- 慢：处理输入提示长度L为O(L)——对于长提示来说是缺点
预编译所有图，输入长度范围从 1 到 max_seq_len（例如 2,048）
- 不实用：在预热期间预编译和缓存max_seq_len图
- 快：1 次图执行即可处理完整的提示

我们引入了提示长度分桶，这是一种在两种替代方案之间取得平衡的优化。我们定义了一组递增的桶大小，(b₀,b₁,b₂,…,b_B-1)，然后根据这些桶值预编译程序图，(G₀,G₁,G₂,…,G_B-1)；B是桶的数量。对于给定的输入提示，我们将提示长度向上舍入到最接近的桶值b_n，填充序列，并使用G_n在一次迭代中处理提示。填充标记上的计算被丢弃。对于大于最大桶大小的提示，我们分段处理它们。

最佳桶大小应由目标应用程序中的提示长度分布决定。在这里，我们采用的桶长度为：128、256、384、512。任何输入提示（最多 2,047 个标记）最多需要 4 次图执行。例如，一个 1,500 个输入提示，生成长度为 256，需要 260 次图执行——4 次处理输入，256 次生成输出。

量化

量化减少了表示值所需的位数；它减少了通过多个加速器节点（通过集合）通信数据的带宽，并降低了服务特定模型大小的硬件要求。

通常，对于BF16权重，一个 175B 参数模型将消耗大约 351GB 内存，因此需要 v4-32 实例才能容纳该模型。通过将权重量化为INT8，我们将模型大小减少了大约 50%，使其可以在较小的 v4-16 实例上运行。由于 LLaMA 对模型激活进行分片，因此量化提供了可忽略不计的通信增益。

在我们的实验中，我们对线性层进行了量化。由于 LLaMA 模型检查点未公开可用，并且我们的目标是评估性能，因此量化模型使用随机权重进行初始化。最近的文献，如AWQ和Integer or Floating Point?，提供了关于 LLaMA 在各种低位量化方案下的性能特性的见解。

批量大小对量化性能的影响

TPU v4被编程为当模型批量大小 (BS) > 1 时在矩阵乘法单元 (MXU) 上运行matmul。当 BS = 1 时，matmul在向量处理器单元 (VPU) 上运行。由于 MXU 比 VPU 更高效，INT8量化在 BS>1 时获得性能。有关详细信息，请参阅性能分析部分。

操作支持

有时，新模型会引入新的数学运算，需要 PyTorch/XLA 扩展其支持的编译操作集。对于 LLaMA，我们支持：多项式。

方法

LLaMA 在 PyTorch/XLA 上开箱即用，支持 LazyTensorCore。我们将此配置作为我们后续分析的基准。所有实验都假设 256 个长输入提示。在没有公开可用模型检查点的情况下，我们使用随机张量初始化进行此推理堆栈优化工作。模型检查点预计不会改变此处讨论的延迟结果。

模型尺寸

假设N是参数数量，dimensions是隐藏大小，n_layers是层数，n_heads是注意力头数，可以使用以下方程来近似模型大小。有关详细信息，请参阅模型概述部分。

N = (dimensions)^2 * n_layers * 12

n_heads不影响N，但以下方程适用于开源模型配置。

dim = 128 * n_heads

缓存大小

模型参数和注意力块中的缓存层都会导致内存消耗。由于默认 LLaMA 模型使用BF16权重，因此本节中的内存消耗计算基于BF16权重。

缓存层的大小由cache_size = max_batch_size * max_seq_len * dimensions计算。在以下计算中，max_batch_size = 1和max_seq_len = 256用作示例配置。每个注意力块中有 2 个缓存层。因此，LLaMA 总缓存大小（以字节为单位）为total_cache_size = n_layers * 2 * cache_size * (2 bytes)。

TPU v4 硬件尺寸

每个 TPU v4 芯片具有 32GB 的可用高带宽内存 (HBM)。表 2 详细说明了内存消耗和容纳 LLaMA 模型所需的 TPU 芯片数量。

表 2：LLaMA TPU v4 HBM 要求（即 TPU v4 芯片要求）

参数数量	参数 (MB)	缓存 (MB)	总计 (GB)	TPU v4 芯片的最小数量
7B	14,000	134	14.128	1
33B	66,000	408	66.41	3
65B	130,000	671	130.67	5
175B	350,000	1,208	351.21	11

指标

以下是衡量推理速度的有用指标。假设T是总时间，B是批量大小，L是解码序列长度。

延迟定义

延迟是获取目标长度L解码结果所需的时间，与批量大小B无关。延迟表示用户等待生成模型响应的时间。

Latency = T (s)

每标记延迟

自回归解码的一步为批处理中的每个样本生成一个标记。每标记延迟是该一步的平均时间。

Per-token latency = T / L (s/token)

吞吐量

吞吐量衡量每单位时间生成的标记数。虽然它不是评估在线服务的有用指标，但它对于衡量批处理速度很有用。

Throughput = B * L / T (tokens/s)

为了最大程度地减少混淆和误解，最好避免使用T / (B * L)等指标，因为它混合了延迟和吞吐量。

结果

图 1 显示了 LLaMA 7B 到 175B 模型的每标记延迟结果。在每种情况下，模型都在一系列 TPU v4 配置上运行。例如，LLaMA 7B 在 v4-8 和 v4-16 上分别显示 4.7ms/标记和 3.8ms/标记。有关更多比较，请访问 HuggingFace LLM 性能排行榜。

在没有本文讨论的功能的情况下，在 v4-32 上运行的 LLaMA 65B 提供了 120ms/标记而不是这里获得的 14.5ms/标记，从而实现了8.3 倍的加速。如前所述，鼓励开发人员尝试我们的自定义torch、torch-xla轮子，这些轮子可重现此处共享的LLaMA 推理结果。

Figure 1: LLaMA Inference Performance on TPU v4 hardware

图 1：TPU v4 硬件上的 LLaMA 推理性能

PyTorch/XLA:GPU 性能优于 PyTorch:GPU eager，并且与 PyTorch Inductor 相似。PyTorch/XLA:TPU 性能优于 PyTorch/XLA:GPU。在不久的将来，XLA:GPU 将提供优化，使其与 XLA:TPU 保持一致。单个 A100 配置仅适用于 LLaMA 7B，而 8-A100 不适用于 LLaMA 175B。

Figure 2: LLaMA Inference Performance on GPU A100 hardware

图 2：GPU A100 硬件上的 LLaMA 推理性能

随着批量大小的增加，我们观察到每标记延迟呈次线性增长，这突出了硬件利用率和延迟之间的权衡。

Figure 3: LLaMA Inference Performance across different batch sizes

图 3：不同批量大小下的 LLaMA 推理性能

我们的研究表明，最大序列输入长度（max_seq_len）对推理延迟的影响相对较小。我们将其归因于标记生成的顺序和迭代性质。性能上的微小差异可能是由于 KV 缓存访问延迟随存储大小增加而变化所致。

Figure 4: LLaMA Inference Performance across different prompt lengths

图 4：不同提示长度下的 LLaMA 推理性能

LLM 通常是内存密集型应用程序；因此，通过量化模型参数，我们可以实现在 MXU 上单位时间加载和执行更大的张量（即 HBM ⇒ CMEM 和 CMEM ⇒ MXU 数据移动）。图 5 显示，INT8仅权重量化提供了 1.6 倍至 1.9 倍的加速，从而允许在给定硬件上运行更大的模型。

当 BS=1 时，INT8 张量被调度到 VPU，VPU 小于 MXU（参见TPU v4 论文）；否则，使用 MXU。因此，当 BS=1 时，量化内存带宽增益被 MXU 利用率不足所抵消。然而，当 BS>1 时，内存增益在量化模型上提供了卓越的延迟。例如，在 175B 参数 LLaMA 的情况下，带量化的 v4-16 和不带量化的 v4-32 提供相似的性能。请注意，我们不提供FP8比较，因为 PyTorch 尚未提供此数据类型。

Figure 5: LLaMA Inference Performance vs. weight-only quantization. The missing blue bars suggest the model size doesn’t fit in the specified TPU hardware.

图 5：LLaMA 推理性能与仅权重量化。缺失的蓝色条表示模型大小不适合指定的 TPU 硬件。

图 6 演示了 PyTorch/XLA 随着输入提示长度从 10 个标记增加到 1,500 个标记而保持稳定的性能优势。这种强大的扩展能力表明 PyTorch/XLA 的重新编译事件最少，从而支持广泛的实际应用。在此实验中，最大长度为 2,048，最大生成长度为 256。

Figure 6: LLaMA Inference Performance vs. Input Prompt Length

图 6：LLaMA 推理性能与输入提示长度

最终想法

我们对 PyTorch/XLA 的未来感到兴奋，并邀请社区加入我们。PyTorch/XLA 完全在开源中开发。因此，请在GitHub上提交问题、提交拉取请求和发送 RFC，以便我们公开协作。您还可以亲自试用 PyTorch/XLA，在各种 XLA 设备（包括 TPU 和 GPU）上。

祝好，
Google 的 PyTorch/XLA 团队
#PyTorch驱动

在 PyTorch/XLA 上实现 LLaMA 65B 超低推理延迟的路径