AMD GPU 在 LLM 推理中的性能优化与成本竞争

2023 年 8 月 9 日 • MLC 社区

MLC-LLM 让我们能够编译 LLM，并使用 ROCm 将其部署到 AMD GPU 上，同时获得有竞争力的性能。更具体地说，在运行 Llama2-7B/13B 时，AMD Radeon™ RX 7900 XTX 的速度可达到 NVIDIA® GeForce RTX™ 4090 的 80%，以及 NVIDIA® GeForce RTX™ 3090 Ti 的 94%。除了 ROCm，我们对 Vulkan 的支持还让 LLM 能够部署到其他 AMD 设备上，例如搭载 AMD APU 的 Steam Deck。

自开源 LLM 蓬勃发展以来，业界已经出现了许多 LLM 推理解决方案。其中，大多数高性能推理方案都基于 CUDA，并针对 NVIDIA GPU 进行了优化。与此同时，随着算力需求不断攀升，让更多类型的硬件加速器获得支持也变得很有价值。AMD 就是一个潜在选择。

关于硬件与软件的讨论

从规格对比可以看出，AMD RX 7900 XTX 与 NVIDIA RTX 4090 和 RTX 3090 Ti 处于相近水平。

• 三者都配备了 24GB 显存，这意味着它们可以容纳相同规模的模型。

• 三者的内存带宽也较为接近。

• RTX 4090 的 FP16 性能是 RX 7900 XTX 的 2 倍，而 RTX 3090 Ti 的 FP16 性能是 RX 7900 XTX 的 1.3 倍。对延迟敏感的 LLM 推理主要受内存带宽限制，因此 FP16 性能在这里并不是瓶颈。

• RX 7900 XTX 比 RTX 4090 便宜 40%。

由于 RTX 3090 Ti 属于上一代产品，两者的价格很难直接比较。我们将其列在这里，是为了提供一个包含更多信息的参考点。总体来看，从硬件规格的角度来说，AMD RX 7900 XTX 与 RTX 3090 Ti 大致相当。

硬件本身未必是 AMD 过去落后的原因。主要差距来自软件支持不足，以及缺少针对相关模型的优化。如今，生态系统中有两个因素开始改变这一局面：

• AMD 正在加大对 ROCm 软件栈的投入，努力追赶。

• 机器学习编译等新兴技术，有助于降低跨不同后端提供通用软件支持的整体成本。

本文将深入研究：截至目前，与 NVIDIA GPU 上的高性能 CUDA 方案相比，AMD GPU 究竟能做到什么程度。

面向 ROCm 的机器学习编译

什么是机器学习编译（MLC）

机器学习编译是一项新兴技术，可以对机器学习工作负载进行编译并自动完成优化。MLC 解决方案不需要为 ROCm 或 CUDA 等每一种后端分别编写特定 kernel，而是可以自动为不同后端生成代码。

这里，我们使用 MLC-LLM。它是一套基于机器学习编译的解决方案，能够为 LLM 提供高性能的通用部署能力。MLC-LLM 构建于 Apache TVM Unity 之上。Apache TVM Unity 是一套机器学习编译软件栈，提供以 Python 为先的高效开发体验和通用部署能力。

MLC-LLM 可以在多种后端上实现业界领先的性能，包括 CUDA、Metal、ROCm、Vulkan 和 OpenCL，覆盖从服务器级 GPU 到移动设备（iPhone 和 Android）的广泛平台。从高层来看，这套框架允许用户获取开放语言模型，通过基于 Python 的工作流进行编译，其中包括用于转换计算图、优化 GPU kernel 布局与调度的 API，并将模型原生部署到目标平台。

面向 AMD GPU 和 APU 的 MLC

支持 AMD GPU 有多种可行方式，包括 ROCm、OpenCL、Vulkan 和 WebGPU。ROCm 软件栈是 AMD 最近重点推动的方向，其中包含许多与 CUDA 软件栈相似的基础组件。Vulkan 是最新一代图形标准，能够覆盖最广泛的 GPU 设备。WebGPU 则是最新的 Web 标准，允许计算任务在浏览器中运行。

尽管存在如此多的实现方式，但很少有机器学习软件方案会构建 CUDA 之外的解决方案。这主要是因为，要为一种新的硬件或 GPU 编程模型重新实现整套软件栈，需要付出很高的工程成本。

我们支持自动代码生成，不必针对每种方式重新编写 GPU kernel，便可同时为上述方案提供支持。话虽如此，最终性能仍取决于底层 GPU runtime 的质量，以及它们在各个平台上的可用性。

对于 Radeon RX 7900 XTX，我们选择 ROCm；对于 Steam Deck 的 APU，我们选择 Vulkan。我们发现，ROCm 软件栈开箱即用。得益于 TVM Unity 高效的、基于 Python 的开发流水线，我们又花了几个小时，进一步实现了一个经过优化的版本。

为了提供 ROCm 支持，我们完成了以下工作：

• 复用面向现有目标平台（例如 CUDA 和 Metal）的整套 MLC 流水线，包括内存规划、算子融合等。

• 复用一套使用 TVM TensorIR 编写的通用 GPU kernel 优化空间，并将其重新定向到 AMD GPU。

• 复用 TVM 的 ROCm 代码生成流程，通过 LLVM 生成底层 ROCm kernel。

• 最后，将生成的代码导出为共享库或静态库，以便通过 CLI、Python 和 REST API 调用。

使用 MLC Python Package 进行基准测试

我们对经过 4-bit 量化的 Llama 2 7B 和 13B 进行了基准测试。我们通过设置单个 prompt token，并生成 512 个 token 来测量解码性能。所有结果均针对单 batch 推理测得。

对于单 batch 推理，在 ROCm 5.6 发布后，其性能可以达到 NVIDIA RTX 4090 的 80%。

关于此次对比需要说明的是：我们的 CUDA 基线有多强？据我们所知，它是这项任务当前最先进的实现。我们认为仍有进一步优化的空间，例如改进 attention 优化。当这些优化进入 MLC 后，我们预计 AMD 和 NVIDIA 的性能数据都会得到提升。如果这些优化只在 NVIDIA 侧实现，那么双方的差距将从 20% 扩大到 30%。因此，在查看这里的数据时，我们建议预留 10% 的误差范围。

我们提供了预构建的 wheel 和操作说明，方便你在自己的设备上复现结果。要运行这些基准测试，请确保你拥有 AMD GPU，并且在 Linux 上运行 ROCm 5.6 或更高版本。请按照相应说明安装启用了 ROCm 的预构建 MLC package，然后运行下面使用 MLC package 的 Python 脚本，复现性能数据：

from mlc_chat import ChatModule

# Create a ChatModule instance that loads from `./dist/prebuilt/Llama-2-7b-chat-hf-q4f16_1`
cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1")

# Run the benchmarks
output = cm.benchmark_generate("Hi", generate_length=512)
print(f"Generated text:\n{output}\n")
print(f"Statistics: {cm.stats()}")

# Reset the chat module by
# cm.reset_chat()

MLC-LLM 还提供了 CLI，允许你与模型进行交互式对话。对于 ROCm，需要从源码构建 CLI。请按照相应说明从源码构建 CLI。

使用 Vulkan 和统一内存在 Steam Deck 上运行

接下来，我们再看看范围更广的 AMD 设备，具体来说，就是搭载 AMD APU 的 Steam Deck。

虽然 BIOS 将 ROCm 可用的 GPU VRAM 限制为 4GB，但 Mesa Vulkan 驱动提供了可靠支持，允许 buffer 借助统一内存突破这一上限，最高使用 16GB 内存。这足以运行经过 4-bit 量化的 Llama-7B。

这些结果让我们看到，覆盖范围广泛的 AMD 设备可以如何获得支持，进而服务更多不同类型的用户。

讨论与未来工作

在生成式 AI 时代，硬件可用性已经成为一个亟待解决的问题。机器学习编译可以在不同硬件后端上提供高性能的通用部署能力，从而帮助缓解这一问题。根据本文给出的证据，我们认为，只要价格和供应情况合适，AMD GPU 已经可以开始用于 LLM 推理。

目前，我们的研究主要关注消费级 GPU。根据过去的经验，面向消费级 GPU 模型的 MLC 优化通常也可以推广到云端 GPU，例如从 RTX 4090 推广到 A100 和 A10g。我们有信心，这套解决方案可以同时覆盖云端和消费级的 AMD 与 NVIDIA GPU。等我们获得更多 GPU 后，也会继续更新这项研究。我们同样鼓励社区基于 MLC 的通用部署流程构建更多解决方案。

本文是 MLC 持续推动高性能通用部署工作的一部分。我们还在积极研究多个可以进一步拓展本次研究的方向：

• 支持 batching 和多 GPU；

• 与 PyTorch 生态系统集成；

• 支持更多量化方案和模型架构；

• 在更多硬件后端上引入更多自动优化能力。

我们最终的结论是：机器学习系统工程是一个需要持续解决的问题。NVIDIA 仍然凭借不断创新引领这一领域。随着 H100 等新硬件的出现，以及更重要的软件演进，我们预计竞争格局还会不断变化。

因此，关键问题不仅在于现在构建出正确的解决方案，还在于如何持续跟进发展，并不断将机器学习工程能力带到新的平台。机器学习工程的生产力才是这里的关键。得益于以 Python 为先的机器学习编译开发流程，我们只用了几个小时就实现了针对 ROCm 优化的支持。随着我们继续探索更多思路，以实现通用部署并解决硬件可用性问题，我们预计相关方法会变得愈发实用。

关于如何尝试部署 MLC LLM，请参阅我们的项目页面，其中提供了详细指南。MLC LLM 的源代码可以在我们的官方 GitHub repository 中获取。也非常欢迎你加入 Discord Channel，参与进一步讨论。

整个 MLC 项目之所以能够实现，离不开我们所依托的开源生态系统。我们希望继续为开源机器学习社区提供开发与支持。感谢 Apache TVM 社区和 TVM Unity compiler 的开发者。开源机器学习社区成员让这些模型得以公开发布，PyTorch 和 Hugging Face 社区则让这些模型变得触手可及。

我们要感谢 RedPajama、Dolly、Vicuna、SentencePiece、LLaMA 和 Alpaca 背后的团队。也要感谢 OpenCL、Vulkan、C++、Python 和 Rust 社区，是它们让这个项目成为可能。