https://blog.mlc.ai/2023/08/09/Making-AMD-GPUs-competitive-for-LLM-inference
MLC-LLM 可以编译 LLM 并将其部署在使用具有竞争性能的 ROCm 的 AMD GPU 上。 更具体地说,对于 Llama2-7B/13B,AMD Radeon™ RX 7900 XTX 的速度是 NVIDIA® GeForce RTX™ 4090 的 80%,是 NVIDIA® GeForce RTX™ 3090Ti 的 94%。 除了 ROCm 之外,我们的 Vulkan 支持还允许我们将 LLM 部署推广到其他 AMD 设备,例如具有 AMD APU 的 SteamDeck。
背景
自开源大型语言模型(LLM)兴起以来,出现了许多大型语言模型(LLM)推理解决方案。 大多数高性能推理解决方案都基于 CUDA 并针对 NVIDIA GPU 进行了优化。 与此同时,随着对计算可用性的高需求,为更广泛的硬件加速器提供支持是很有用的。 AMD 是一个潜在的候选者。
硬件和软件讨论
从规格比较中,我们可以看到AMD的RX 7900 XTX与NVIDIA的RTX 4090和RTX 3090 Ti非常匹配。
- 全部都有 24GB 内存,这意味着它们可以容纳相同尺寸的模型。
- 它们都具有相似的内存带宽。
- 4090 的 FP16 性能比 7900 XTX 高 2 倍,而 3090 Ti 的 FP16 性能比 7900 XTX 高 1.3 倍。 延迟敏感的 LLM 推理主要受内存限制,因此 FP16 性能不是这里的瓶颈。
- RX 7900 XTX 比 RTX 4090 便宜 40%。
3090Ti 的价格很难比较,因为那是上一代产品。 我们将其放在这里作为参考点以提供更多信息。 从高层次来看,我们可以发现从硬件规格的角度来看,AMD 7900 XTX 与 RTX 3090 Ti 相当。
硬件并不一定是AMD过去落后的原因。 主要差距是由于缺乏相关模型的软件支持和优化。 生态系统中有两个因素开始改变现状:
- AMD 正在努力追赶 ROCm 堆栈的投资。
- 机器学习编译等新兴技术有助于降低跨后端的更通用软件支持的总体成本。
在这篇文章中,我们将深入研究 AMD GPU 与目前 NVIDIA GPU 上的高性能 CUDA 解决方案相比的表现如何。
ROCm 的机器学习编译
什么是机器学习编译(MLC)。 机器学习编译是一种新兴技术,可编译并自动优化机器学习工作负载。 MLC 解决方案不会为每个单独的后端(如 ROCm 或 CUDA)制作特定的内核,而是会自动为不同的后端生成代码。 在这里,我们利用 MLC-LLM,这是一种基于 ML 编译的解决方案,可为 LLM 提供高性能通用部署。 MLC-LLM 构建在 Apache TVM Unity 之上,这是一个机器学习编译堆栈,可提供高效的 Python 优先开发和通用部署。 MLC-LLM 为各种后端带来了最先进的性能,包括 CUDA、Metal、ROCm、Vulkan 和 OpenCL,涵盖服务器级 GPU 到移动设备(iPhone 和 Android)。 在较高层面上,该框架允许用户采用开放语言模型,并使用基于 Python 的工作流程进行编译,包括用于转换计算图的 API、优化 GPU 内核的布局和调度,并将其本地部署在感兴趣的平台上。
适用于 AMD GPU 和 APU 的 MLC。 支持 AMD GPU 的方法有多种:ROCm、OpenCL、Vulkan 和 WebGPU。 ROCm堆栈是AMD最近推动的,有很多类似于CUDA堆栈的相应构建块。 Vulkan 是最新的图形标准,提供跨 GPU 设备最广泛的支持。 WebGPU 是最新的网络标准,允许计算在网络浏览器上运行。 尽管有很多可能的方法,但很少有 ML 软件解决方案是为 CUDA 以外的解决方案构建的,这主要是由于为新硬件或 GPU 编程模型复制堆栈的工程成本。 我们支持自动代码生成,无需为每个代码重新设计 GPU 内核,并为所有这些方式提供支持。 话虽如此,性能仍然取决于低级 GPU 运行时的性能及其在每个平台中的可用性。 我们为 Radeon 7900 XTX 选择 ROCm,为 Steamdeck 的 APU 选择 Vulkan。 我们发现 ROCm 堆栈开箱即用。 得益于 TVM Unity 中基于 Python 的高效开发流程,我们花了几个小时来进一步推出优化版本。 我们做了以下事情来获得 ROCm 支持:
- 为现有目标(例如 CUDA 和 Metal)重用整个 MLC 管道,包括内存规划、算子融合等。
- 重用以 TVM TensorIR 编写的通用 GPU 内核优化空间,并将其重新定位到 AMD GPU。
- 重用 TVM 的 ROCm 代码生成流程,通过 LLVM 生成低级 ROCm 内核。
- 最后,将生成的代码导出为可由 CLI、Python 和 REST API 调用的共享或静态库。
使用 MLC Python 包进行基准测试
我们使用 4 位量化对 Llama 2 7B 和 13B 进行了基准测试。 我们通过设置单个提示令牌并生成 512 个令牌来测量解码性能。 所有结果都是针对单批次推理进行测量的。
对于单批次推理性能,随着ROCm 5.6的发布,可以达到NVIDIA 4090速度的80%。
比较注意事项:我们的 CUDA 基线有多强? 据我们所知,这是这项任务的最先进技术。 我们相信仍有改进的空间,例如 通过更好的注意力优化。 一旦这些优化进入 MLC,我们预计 AMD 和 NVIDIA 的数据都会有所改善。 如果这样的优化只在NVIDIA这边实现,那么差距就会从20%拉大到30%。 因此,我们建议在查看此处的数字时设置 10% 的误差线
自己尝试一下
我们提供预制轮子和说明,以便在您自己的设备上重现我们的结果。 要运行这些基准测试,请确保您拥有在 Linux 上运行且具有 ROCm 5.6 或更高版本的 AMD GPU。 按照此处的说明安装启用了 ROCm 的预构建 MLC 包:运行下面的 Python 脚本,该脚本使用我们的 MLC 包来重现性能数据:
- 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()
使用具有统一内存的 Vulkan 在 SteamDeck 上运行
我们还可以研究更广泛的 AMD 设备,更具体地说,是配备 AMD APU 的 SteamDeck。 虽然 ROCm 中可用的 GPU VRAM 在 BIOS 中上限为 4GB,但 Mesa Vulkan 驱动程序具有强大的支持,允许缓冲区使用高达 16GB 的统一内存来超出上限,这足以运行 4 位量化的 Llama-7B。
这些结果揭示了如何为更多样化的消费者群体提供广泛的 AMD 设备支持。
讨论和未来的工作
硬件可用性已成为生成式人工智能时代的一个紧迫问题。 机器学习编译可以通过在硬件后端实现高性能通用部署来提供帮助。 鉴于所提供的证据,加上合适的价格和可用性,我们认为 AMD GPU 可以开始对 LLM 推理有用。
目前我们的研究重点是消费级 GPU。 根据我们过去的经验,消费级 GPU 模型的 MLC 优化通常可以推广到云 GPU(例如从 RTX 4090 到 A100 和 A10g)。 我们相信该解决方案可以在云以及消费级 AMD 和 NVIDIA GPU 上推广,并且一旦我们能够使用更多 GPU,我们也将更新我们的研究。 我们还鼓励社区在 MLC 通用部署流程之上构建解决方案。
这篇文章是通过 MLC 实现高性能通用部署的持续努力的一部分。 我们还积极致力于几个可以推广我们研究的领域。
- 启用批处理和多 GPU 支持;
- 与 PyTorch 生态系统集成;
- 赋能更多量化和模型架构;
- 在更多硬件后端引入更多自动优化。
我们最后的结论是,机器学习系统工程是一个持续存在的问题。 NVIDIA 凭借持续的创新仍处于该领域的领先地位,我们预计,随着 H100 等新硬件以及更重要的是软件的发展,这一格局将会发生变化。 因此,关键问题不仅在于现在构建正确的解决方案,还在于如何赶上并不断将机器学习工程引入新平台。 机器学习工程的生产力是这里的关键。 得益于 Python 优先的 ML 编译开发流程,我们在几个小时内就获得了 ROCm 优化的支持。 随着我们探索更多想法来实现通用部署并解决硬件可用性问题,我们预计相关方法将变得更加有用。
链接
请参阅我们的项目页面(https://mlc.ai/mlc-llm/docs/),了解有关如何尝试 MLC LLM 部署的详细指南。 MLC LLM 的源代码可在我们的官方 GitHub 存储库上获取(https://github.com/mlc-ai/mlc-llm/)。 也非常欢迎您加入 Discord 频道进行进一步讨论(https://discord.com/invite/9Xpy2HGBuD)。
致谢
整个 MLC 项目的实现要归功于我们所站的开源生态系统的肩膀。 我们很乐意继续开发和支持开源机器学习社区。 我们要感谢 Apache TVM 社区和 TVM Unity 编译器的开发人员。 开源机器学习社区成员公开了这些模型。 PyTorch 和 Hugging Face 社区使这些模型变得易于访问。 我们要感谢 RedPajama、Dolly、Vicuna、SentencePiece、LLaMA 和 Alpaca 背后的团队。 我们还要感谢支持该项目的 OpenCL、Vulkan、C++、Python 和 Rust 社区。