在AMD GPU上进行大型语言模型推理优化

Large language model inference optimizations on AMD GPUs — ROCm Blogs

2024年3月15日，作者 Seungrok Jung.

大型语言模型（LLM）在自然语言处理和理解领域取得了巨大进展，促进了多个领域的人工智能应用。LLM有许多有前景的用例，包括人工智能助手、聊天机器人、编程、游戏、学习、搜索和推荐系统。这些应用利用LLM的能力提供个性化和互动性的体验，从而增强用户的参与度。

LLM使用了Transformer架构来解决梯度消失和爆炸问题。这种架构可以轻松并行化自注意力机制，使得能够高效利用多GPU。而其他架构如循环神经网络（RNN）及其变种（如LSTM和GRU）在处理长序列词语时则显得力不从心。

尽管LLM如GPT和Llama表现出令人印象深刻的能力，但在它们被用于商业应用之前，必须进行积极的优化，这是因为它们具有庞大的参数规模和自回归序列处理行为。许多努力已经付诸实践，通过利用GPU的计算能力（TFLOPs）和内存带宽（GB/s）来提高LLM的吞吐量、延迟和内存占用。

我们将讨论这些优化技术，并通过对比AMD的 MI250 和 MI210 GPUs上的Llama-2-7B 和Llama-2-70B 模型的性能指标来加以说明andGPUs。

模型特点：Llama2-7b 和 Llama2-70b

Llama2-7b 和 70b 模型能够处理 32,000 个词汇。这些模型可以处理最长 4,096 个 token 序列。Llama2 通过采用以下新特性优化了其训练和推理性能：

Sigmoid 线性单元 (SiLU) 激活函数：替换了常用的 Rectified Linear Unit (ReLU)，以减少消失梯度问题，实现更平滑的激活。
旋转位置嵌入：降低了经典绝对位置嵌入层的计算成本，同时保留了 token 序列的位置信息。
预归一化：LlamaRMSNorm 模块对输入而不是输出进行归一化，减少了梯度消失和爆炸问题。

在 Llama-2-7b 模型中，自注意力模块有 32 个注意力头，每个头有 128 个维度。多层感知器 (MLP) 模块具有 11,008 的中间大小，并由三层组成：`gate_proj`、`up_proj` 和 down_proj。

基于其行为，LLMs 分为以下几类：

掩码语言模型 (MLM): 在提供的上下文单词之间预测一个新的掩码单词。BERT 是一个 MLM 的例子。
因果语言模型 (CLM): 在提供的上下文单词之后预测下一个单词。一个广为人知的 CLM 示例是 GPT 文本生成。由于其顺序处理行为，CLM 也被称为自回归 token 生成模型。

在本博客中，我们更关注 Llama2 CLM。

CLM 中的 token 生成分为以下两个阶段：

首次 token 时间 (TTFT):生成第一个 token 所需的时间。预填延迟被定义为跨请求的平均 TTFT。在下图中，TTFT 是从输入提示到生成“The largest continent”的“in”所需的时间。
每个输出 token 时间 (TPOT): 以自回归方式生成每个输出 token 所需的时间。输出解码延迟被定义为跨请求的平均 TPOT，通常使用输出解码阶段所需的时间进行估计。在下图中，TPOT 是“the”的解码延迟。

TTFT 和 TPOT 被用来计算 CLM 中的延迟：

latency = TTFT + TPOT * (max_new_tokens - 1)

在预填阶段，经过 token 嵌入后的输入维度与 batch_size 输入序列长度* 成正比。预填阶段的 token 是同时处理的。然而，经过 token 嵌入后的 输出解码 输入与 batch_size 成正比；这个阶段的 token 是顺序处理的。这就是为什么当 batch size 为 1 时，输出解码操作由高而瘦的 GEMM（或 GEMV）组成。

为了简单起见，我们对生成的 token 使用贪婪解码，这种方法已知在输出 logits 方面具有最小的 token 解码开销。在实际的聊天机器人的场景中，如果允许丰富且意外的输出 token 生成，最好考虑基于采样的解码和更高的 beam width。但在贪婪解码方案中，自回归 CLM 会从模型输出 logit 中生成 top-1 token。

设备特性：MI210

AMD的Instinct ™ MI210在FP16数据类型下的最大计算能力为181 TFLOPs。为了充分利用矩阵核性能，GEMM的矩阵维度应该足够大。在大批量情况下的LLM预填充解码阶段使用大型输入矩阵，可以从矩阵核心的高性能中获益。使用MI210，在提示序列长度和批量大小较大的预填充阶段中，GEMM操作是计算受限的。

MI210可以提供最多1.6 TB/s的双倍数据速率（DDR）内存带宽。输出解码按顺序处理令牌。该自回归解码只有一个序列长度维度（使得长且瘦的GEMM或GEMV），这种操作是内存受限的。由于LLM输出令牌生成的这种顺序性，输出解码得益于DDR带宽。

MI250由两个来自MI210的图形计算芯片（GCD）组成。因此，MI250的计算能力、内存大小和内存带宽是MI210的两倍。LLM可以在MI250的两个硅芯片上进行张量并行（TP）、流水线并行（PP）或数据并行（DP）风格的模型并行。数据并行可以通过两倍的模型参数复制开销来加倍LLM的吞吐量。由于没有开销，张量并行被广泛使用，因为它能够通过一些集合操作同步开销将更大的LLM适配到高容量的MI250 DDR上。

在前图中，注意到MI250单个GCD的房顶线与MI210的相似。

软件设置

在主机上安装ROCm

要在主机上安装ROCm 6.0，请参见 install instructions.

设置Docker

要设置官方的[PyTorch ROCm Docker容器PyTorch ROCm Docker container, 请使用以下命令：

docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G --name llm_optimization rocm/pytorch:rocm6.0_ubuntu22.04_py3.9_pytorch_2.0.1

请使用以下命令：

运行以下命令来安装 PyTorch 2.3 nightly版:

pip3 uninstall torch torchvision torchaudio pytorch-triton-rocm -y
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0

有关库的设置，请参阅 Hugging Face 的 transformers.

有关工具包的设置，请参阅 Text Generation Inference (TGI).

MI210 上 Llama-2-7b 的优化对比

预填充延迟

输出解码延迟

默认机器学习框架

PyTorch 2 支持两种运行模式：eager 模式和编译模式。Eager 模式是默认的 PyTorch 模式，在这种模式下，模型的操作会在运行时顺序执行。编译模式涵盖在大语言模型推理优化技术中。

为了运行一个 LLM 解码器模型（例如，Llama-2），Hugging Face 提供了 transformers 库，以在 PyTorch 上运行模型。

transformers 库使用各种标记生成选项作为其APIs中的参数APIs。在这篇博客中，为了公平比较每种优化的性能，这些选项被部署如下：

预填充: 使用一个长度为 2048 的序列作为提示，预填充延迟会随批量大小的增加而增加，因为在预填充期间的大型 GEMM 计算是计算密集型的。
输出解码: 当批量大小增加时，输出解码延迟不会显著增加，因为在此阶段的 GEMM 算术强度仍然受到内存带宽的限制。

LLM推理优化技术

在这里，我们讨论各种LLM（大型语言模型）推理优化技术。

PyTorch编译模式

在PyTorch编译模式中，一个模型会被综合成一个图，然后被降低到最优算子。这些算子使用TorchInductor进行编译，后者使用OpenAI的Triton作为GPU加速的基础构件。PyTorch编译模式的一个优势是其GPU内核是用Python编写的，这使得修改和扩展它们更加容易。PyTorch编译模式通常提供更高的性能，因为模型操作在运行时之前就已经融合，这使得高性能内核的部署变得简单。

要在PyTorch编译模式下运行LLM解码器模型（例如，Llama2），必须明确指定模型的特定层作为编译目标。在LLM解码器模型中使用PyTorch编译模式需要对动态输入形状进行重新编译，因为输入的批量大小和序列长度在运行时可能变化。要支持动态输入形状的重新编译，可以设置参数 dynamic=True。

for i in range(model.config.num_hidden_layers):
    model.model.decoder.layers[i].self_attn = torch.compile(model.model.decoder.layers[i].self_attn, backend="inductor", mode="default", dynamic=True)

预填充: 预填充的延迟显著减少。但是，对于LLM解码器模型，由于不同批量大小和输入序列长度的重新编译开销，初始预填充仍然面临巨大开销。在初始重新编译（预热）之后测量预填充延迟。
输出解码: 输出解码的延迟略有改善，因为模型已部分编译。然而，由于部分键/值缓存的动态形状问题，图形回退到即时模式。很多工作正在致力于解决这个问题（称为静态键/值缓存）。静态键/值缓存与`torch.compile`结合使用，显示出更好的输出解码性能，但我们在本博客中不讨论这一点。

Flash Attention v2

The Flash Attention (flash_attention) 算法旨在解决 transformer中多头注意力（Multi-head Attention，MHA）模块查询、键和值组件所需的大量内存移动问题。这个问题通过对部分查询进行分块并存储在较快的缓存内存中来解决，从而避免在MHA计算过程中不断从较慢的外部DDR内存读取数据。`flash_attention v2`（https://arxiv.org/abs/2307.08691）最大化了长输入序列长度上的并行性，相比于原生MHA带来了显著的性能提升。

您可以在最新的Hugging Face的transformers库中无缝使用适用于ROCm的flash_attention v2 模块。

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="flash_attention_2")

预填充: 对于大批量和长序列长度，Flash Attention模块显著减少了预填充处理的延迟，因为MHA矩阵维度与这两者成正比。这使得flash attentions在这种场景下的增益更为显著。
输出解码:在输出解码阶段，`flash_attention`并不有效，因为此时的序列长度仅为1。

Memory-efficient multi-head-attention

Memory-efficient, multi-head-attention (Xformers) 是一个来自 Meta 的可定制模块集合，用于优化 transformers。Xformers 的一个关键特性是它的 memory-efficient MHA 模块，该模块可以显著减少 MHA 处理过程中的内存流量。此模块使用类似于 flash_attention 的算法来减少 DDR 读取或写入带宽。

你可以无缝地将 Xformers 的 ROCM 版本的 memory-efficient MHA 模块集成到 Hugging Face 的 transformers 库中。

Prefill（预填充）: Memory-efficient MHA 也显著减少了在处理大批量和长序列长度时的预填充处理延迟，原因与 flash_attention v2 相同。
Output decoding（输出解码）: Xformers 在输出解码阶段并不有效，因为此时序列长度只有 1。

分页注意力机制(vLLM)

vLLM推理系统中的分页注意力机制 (paged_attention) 是一种算法，可以在输出解码阶段高效地减少内存消耗，并将延迟降低两到四倍。分页注意力机制通过使用虚拟内存和分页技术管理输出解码阶段的键和值缓存（K-V缓存），以减少内存碎片化。传统的K-V缓存会为最大输出标记长度（根据模型不同为2,048或4,096）预先分配内存，如果实际解码长度较短，则可能导致内存碎片化。该算法利用分页的虚拟内存，可以在束搜索较大且并行运行多个请求时节省K-V缓存内存。

为ROCm的vLLM的分页注意力机制目前已经可用。

预填充阶段: 分页注意力机制在预填充阶段无效果，因为此阶段不使用K-V缓存。
输出解码阶段: 分页注意力机制可以显著减少解码延迟。

PyTorch TunableOp

PyTorch TunableOp 允许你使用高性能的 rocblas 和 hipblaslt 库进行 GEMM（通用矩阵乘法）。它会对大型语言模型（LLM）进行分析，并为每个多头注意力（MHA）和多层感知器（MLP）模块准备性能最优的 GEMM 核心程序。在运行时，这些最优的 GEMM 核心程序将替代 PyTorch 内置的 GEMM 核心程序进行调用。

PyTorch TunableOp 现在已经可用。

预填充Prefill: 与 flash_attention v2 结合使用时，PyTorch TunableOp 在不同批处理大小下显示出了显著的性能提升。
输出解码Output decoding: 使用 paged attention 功能时，PyTorch TunableOp 对细长矩阵（tall and skinny GEMM 或 GEMV）也显著减少了延迟。因此，输出解码性能最能从 rocBLAS 和 hipBLASLt GEMM 中受益。

多GPU大语言模型（LLM）推理优化

预填充延迟

输出解码延迟

Hugging Face文本生成推理

多GPU推理和训练的扩展需要模型并行技术，例如张量并行TP, 管道并行PP, or 数据并行DP. 张量并行（TP)广泛应用，因为它不会导致流水线气泡；数据并行（DP)则能提供较高的吞吐量，但需要将参数的副本载入GPU显存。

在这篇博客中，我们使用TP将模型分割到多个GPU，并使用Hugging Face的TGI 工具来测量多GPU大语言模型（LLM）的推理性能。 Hugging Face’的TGI实现支持ROCm启用的`flash_attention`和`paged_attention`，兼容PyTorch的TunableOp，以及未来支持ROCm启用的量化方法（如GPTQ），使其成为一个不错的选择。
一台服务器配置了4个MI250 GPU，总共有8个图形计算组件（GCD）。每个GCD有64GB的高带宽内存（HBM）。

为了充分利用多个MI250 GPU，你需要考虑GPU GCD之间的互连带宽，因为GCD之间的连接具有非均匀的吞吐量。例如，在TP=4的情况下，使用GCD#0、1、4、6结合在一起可以获得最佳性能，因为集合操作（例如全规约或全收集）在TP中引起的同步开销较少。

启用了非均匀内存访问（NUMA)平衡后，GPU必须等待源自页面错误的预取内存管理单元（MMU)通知更改。因此，建议禁用NUMA平衡，以防止周期性自动平衡干扰GPU操作。执行以下命令禁用NUMA平衡：

echo 0 > /proc/sys/kernel/numa_balancing

预填充和输出解码: 在TP=8（8 GCD）的情况下，预填充和输出解码延迟比TP=4（4 GCD）更好。延迟提升并未增加一倍，因为为了每层同步多头注意力（MHA）和多层感知器（MLP）进行的集合操作也是一个巨大的延迟瓶颈。

总结

在这篇博客中，我们介绍了在AMD CDNA2 GPU上部署最新的LLM（大型语言模型）的几种软件优化技术。这些技术包括PyTorch 2编译、Flash Attention v2、`paged_attention`、PyTorch TunableOp和多GPU推理。这些技术已经被AI社区广泛采用。使用这些优化措施，根据批处理大小和输入序列长度，你可以体验最多高达三倍的开箱即用加速效果。