增强在 AMD GPU 上的 vLLM 推理

最新推荐文章于 2025-03-12 19:37:10 发布

109702008

最新推荐文章于 2025-03-12 19:37:10 发布

阅读量2.8k

点赞数 46

分类专栏： # ROCm 文章标签：人工智能学习

本文链接：https://blog.csdn.net/eidolon_foot/article/details/144069126

版权

ROCm 专栏收录该内容

122 篇文章

订阅专栏

Enhancing vLLM Inference on AMD GPUs — ROCm Blogs

2024 年 10 月 11 日，作者 Clint Greene.

在这篇博客中，我们将展示在 ROCm 上使用 AMD Instinct 加速器进行 vLLM 推理的最新性能增强。简而言之，vLLM 优化了 GPU 内存的利用，使得大语言模型（LLM）可以在现有硬件约束内更高效地处理，最大化吞吐量并最小化延迟。我们将首先简要解释像 Llama 3 和 ChatGPT 这样的因果语言模型如何生成文本，从而激发提高吞吐量和减少延迟的需求。如果你是 vLLM 的新手，我们也推荐阅读我们的介绍在 AMD GPU 上进行 vLLM 推理和服务。ROCm 6.2 引入了以下 vLLM 功能支持，我们将在这篇博客中使用这些功能。

FP8 KV Cache: Store Key-Value (KV) pair data in FP8 (8-bit floating point) to enhance efficiency.
GEMM Tuning: Achieve significant performance boosts by tuning matrix multiplications.
FP8 Quantization: Enable inference on models quantized to FP8, improving speed without compromising accuracy.

理解文本生成

为了充分理解这些功能的优势，我们首先来看看像 Llama 3 和 ChatGPT 这样的因果语言模型如何生成文本。这些大型语言模型（LLM）经过预训练，可以基于前面的令牌预测序列中的下一个令牌。例如，如果给模型提供了“The color of the sky varies”这一序列，它会学会预测下一个令牌可能是“depending”或其他合适的词。当 LLM 收到一个提示时，它会通过逐步预测序列中最可能的下一个令牌来生成文本。预测的令牌会追加到输入中，生成过程会持续进行，直到达到最大令牌数或输出了句子结束（EOS）令牌，表示序列结束。例如，如果 LLM 以“The color of the sky varies”作为提示，它可能会逐步生成“depending”“on”“the”“time”“of”“day”，然后是“<EOS>”，结束生成过程。通常，生成过程（推理）分为两个关键阶段：预填充和解码。

预填充阶段对标记化的输入提示进行编码、嵌入，并计算键值对。在变压器架构中，self-attention 机制依赖于键值对，其中每个输入令牌都与一个键向量和一个值向量关联。键用来通过与查询向量取点积来计算注意力分数，使模型能够确定每个令牌在当前输入中的相关性。而这些值提供了由这些注意力分数加权后的上下文信息。在推理过程中，这些键值对是必不可少的，因为它们需要被计算或检索以告知模型的预测。对于预填充阶段，大多数计算已经得到了优化，可以在 GPU 上一次性执行。

相比之下，负责生成输出的解码阶段由于要逐个生成令牌的顺序性质，显著慢于预填充阶段。模型基于输入提示和由此计算出的键值对生成第一个令牌。这个令牌会追加到现有输入中，更新后的序列被用来计算新的一组键值对。这个过程会重复进行以生成后续的令牌，直到达到停止标准。由于解码阶段并行化的挑战，GPU 吞吐量受限，导致成本性能降低。为了增强吞吐量和减少延迟，可以应用多种优化。ROCm 6.2 支持的 vLLM 主要优化包括 FP8 键值缓存、GEMM 调优和 FP8 模型量化。

前提条件

为了跟随本博客的内容，您需要:

Linux: 请参阅支持的Linux发行版。
ROCm: 请参阅安装说明。
MI300+ GPUs: FP8支持仅在MI300系列上可用。

安装

要访问ROCm 6.2中的最新vLLM功能，克隆vLLM仓库，将Dockerfile.rocm中的`BASE_IMAGE`变量修改为`rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.3.0`，并使用以下命令构建Docker镜像。根据您的系统，构建过程可能需要较长时间。

git clone https://github.com/vllm-project/vllm.git
cd vllm
DOCKER_BUILDKIT=1 docker build -f Dockerfile.rocm -t vllm-rocm .

在构建vLLM ROCm Docker镜像后，您可以使用以下命令运行它。要在容器中使用LLM文件夹，请将`<path/to/model>替换为实际的文件夹路径。如果您没有要挂载的模型，可以删除-v <path/to/model>:/app/models`选项。

docker run -it --network=host --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --device /dev/kfd --device /dev/dri -v <path/to/model>:/app/models vllm-rocm

FP8 KV缓存

KV缓存是通过避免对相同输入token的KV值进行冗余计算，从而加速LLM推理的关键。当处理一个提示时，会计算KV值并生成一个token。这个新token会被附加到输入序列中，需要计算新的KV值。然而，原始token的KV值是不变的，因此重新计算它们是低效的。KV缓存通过将先前token的KV值存储在GPU的高带宽内存（HBM）中来解决此问题，以防止冗余计算。这提升了吞吐量和效率。

然而，使用KV缓存也有一定的限制。通常它以float16存储，其大小随着序列长度、批处理大小、注意力层数和嵌入维度线性增加。因此，较大的模型尺寸和较长的序列需要更多的内存，这可能会限制可行的批处理大小。此外，GPU内存限制对序列长度也施加了实际管理上的限制。使用ROCm 6.2，vLLM中的KV缓存现在可以以FP8格式存储，显著减少了内存占用。这个增强功能有效地使您在保持其他参数不变的情况下将序列长度或批处理大小加倍。

现在让我们探讨如何在vLLM中访问这一新功能。要以FP8存储KV值，只需在`vllm serve`命令中包含`--kv-cache-dtype fp8`。

vllm serve meta-llama/Meta-Llama-3-8B-Instruct --kv-cache-dtype fp8

如果您观察终端输出，对于MI300X系统中的单个GPU，您将看到以下输出：

INFO 08-29 19:06:51 gpu_executor.py:121] # GPU blocks: 157794, # CPU blocks: 4096

如果您现在运行 vllm serve 且不指定`--kv-cache-dtype fp8`，您将看到如下输出:

INFO 08-29 19:05:21 gpu_executor.py:121] # GPU blocks: 78897, # CPU blocks: 2048

要估算可服务的最大token数，将GPU块数乘以默认的块大小16。使用FP8格式的KV缓存，最大容量约为2,524,704个token。相比之下，使用FP16格式的KV缓存，容量约为1,262,352个token。这一FP8带来的token容量加倍展示了KV缓存内存占用的显著减少。

GEMM 调优

通用矩阵乘法（GEMM）操作支撑着许多神经网络计算，例如卷积和全连接层，并且在训练和推理中占据了很大一部分计算负载。优化 GEMM 可以显著提高吞吐量并减少生成式 AI 应用的延迟。GEMM 调优调整了诸如块大小、内存访问模式和线程块配置等因素，以最大限度地利用 GPU 的并行处理能力。

通过 ROCm 6.2，得益于 PyTorch 的 TunableOp 集成，现在在 vLLM 中优化 GEMM 变得更加容易。有关详细信息，请访问我们的博客关于使用 TunableOps 调优 GEMM。要优化 vLLM 工作负载，请设置环境变量 PYTORCH_TUNABLEOP_ENABLED=1并照常运行您的工作负载。调优完成后，会生成一个名为 tunableop_results0.csv 的 CSV 文件来保存结果。未来的运行将自动加载并应用这些调优。如果输入或输出长度发生变化，将会触发新的调优运行。

例如，要对 Llama3-8B 进行 GEMM 调优，设定输入和输出长度为 512 个 token，并对其延迟进行基准测试，请在终端中运行以下命令：

export PYTORCH_TUNABLEOP_ENABLED=1
python3 benchmarks/benchmark_latency.py --input-len 512 --output-len 512 --num-iters 10 --model meta-llama/Meta-Llama-3-8B-Instruct

当基准测试完成后，您将看到输出平均延迟。

Avg latency: 4.3067230121872855 seconds

接下来，让我们在没有 GEMM 调优的情况下进行基准测试并比较延迟。

PYTORCH_TUNABLEOP_ENABLED=0 python3 benchmarks/benchmark_latency.py --input-len 512 --output-len 512 --num-iters 10 --model meta-llama/Meta-Llama-3-8B-Instruct

在没有 GEMM 调优的情况下，平均延迟超过 4.60 秒！GEMM 调优使延迟减少了大约 6.5%。

FP8 量化

随着大型语言模型（LLM）的参数规模增长到数千亿，如何高效地部署这些模型变得越来越关键和具挑战性。通常，LLM 的组件——例如权重、激活值和 KV 缓存——使用 16 位浮点数表示，因为这在输出质量和速度之间提供了良好的平衡。进一步提升推理速度的一种广泛使用的技术是量化，它通过降低 LLM 部分或全部组件的数值精度来实现。传统上，这涉及将模型参数的数据类型从 FP16 转换为 INT8。然而，INT8 量化会显著降低模型输出质量，特别是对于小于 70 亿参数的更小的 LLM。

一个有前途的替代方法是 FP8 格式，它提供了类似于 8 位整数量化的性能优势，而不会影响输出质量。FP8 提供比 INT8 更大的精度和动态范围，使其非常适合对 LLM 中关键性能的组件进行量化，包括权重、激活值和 KV 缓存。

要理解 FP8，有必要回忆一下浮点数由三部分组成：

符号位: 一个比特，指示数字是正数还是负数
指数（范围）: 数的幂
尾数（精度）: 数的有效数字

FP8 有两种变体，适用于不同的用例：E4M3 和 E5M2：

E4M3: 具有 1 个符号位，4 个指数位和 3 个尾数位，能够表示高达 ±448 和 NaN（非数字）的值。
E5M2: 具有 1 个符号位，5 个指数位和 2 个尾数位，能够存储高达 ±57344、±无穷大和 NaN 的值。

通常，E4M3 用于前向传播或推理，因为激活值和权重需要更高的精度。相比之下，E5M2 则用于反向传播，因为梯度对精度损失的敏感性较小，但会从更高的动态范围中受益。相比之下，INT8 主要关注于尾数（精度），可能包含或不包含符号位，但缺乏指数位，意味着它不支持广泛的范围。

在 ROCm 6.2 中，你现在可以使用 vLLM 以 E4M3 FP8 格式部署模型，只需在终端中运行命令 vllm serve <model-name>。这里，`<model-name>` 可以是一个 E4M3 FP8 量化模型的本地路径或者是 Hugging Face 上 NeuralMagic 的 FP8 量化模型名称。有关 vLLM 支持的 FP8 NeuralMagic 模型的完整列表，请点击这里。

例如，要以 FP8 格式服务 Meta-Llama-3-8B 模型，可以运行以下命令：

vllm serve neuralmagic/Meta-Llama-3-8B-Instruct-FP8

然后可以在另一个终端窗口中使用 curl 命令查询它。

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "neuralmagic/Meta-Llama-3-8B-Instruct-FP8",
        "prompt": "Write a haiku about artificial intelligence",
        "max_tokens": 128,
        "top_p": 0.95,
        "top_k": 20,
        "temperature": 0.8
      }'

接下来，我们使用相同的参数但采用 FP8 量化格式对 Llama3-8B 进行延迟基准测试。

python3 benchmarks/benchmark_latency.py --input-len 512 --output-len 512 --num-iters 10 --model neuralmagic/Meta-Llama-3-8B-Instruct-FP8

这将产生 4.13 秒的平均延迟。开箱即用的 FP8 量化比之前的 FP16 基准降低了大约 10% 的延迟。接下来，我们将基于 GEMM 调优进行延迟基准测试。

export PYTORCH_TUNABLEOP_ENABLED=1
python3 benchmarks/benchmark_latency.py --input-len 512 --output-len 512 --num-iters 10 --model neuralmagic/Meta-Llama-3-8B-Instruct-FP8

通过 GEMM 调优，平均延迟现在下降到 3.40 秒以下，相比之前的 FP16 基准大约降低了 26%。

总结

在这篇博文中，我们简要讨论了 LLM 如 Llama 3 和 ChatGPT 如何生成文本，强调了 vLLM 在提升吞吐量和减少延迟方面的作用。我们介绍了如何在 KV 缓存中以 FP8 格式存储值，优化矩阵乘法以实现更快的计算，以及如何在 FP8 中执行完整的推理。通过这些最新的增强功能，我们展示了 ROCm 6.2 如何显著加速你的 vLLM 工作负载。

FP8 量化是一种能够在不显著牺牲模型输出质量的情况下，显著提高操作效率和减少延迟的技术，尤其适合那些更小参数规模的 LLM 以及对性能有高要求的应用场景。对于有资源限制的部署环境，FP8 提供了一种有效的解决方案，可以更好地平衡计算资源的利用率和模型的准确性。