Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

A卡跑大模型,性能达到4090的80%,价格只有一半:陈天奇TVM团队出品

英伟达 GPU 买不到的问题,就这样解决了?

最近,科技领域有很多人都在为算力发愁。
图片
OpenAI CEO 奥特曼:我整天在都想着 flops。

自预训练大模型兴起以来,人们面临的算力挑战就变得越来越大。为此,人们为大语言模型(LLM)提出了许多训练和推理的解决方案。显然,大多数高性能推理解决方案都基于 CUDA 并针对英伟达 GPU 进行了优化。

但在动辄千亿参数的模型体量,多家科技公司激烈竞争,以及单一供应商的合力作用下,想抢到 GPU 又变成了一件难事。

最近,微软、OpenAI 等公司都表示正在采取必要措施来缓解用于 AI 任务的 H100、A100 专用 GPU 的短缺问题。微软正在限制员工访问 GPU 的时间,Quora 首席执行官表示,硬件短缺掩盖了人工智能应用程序的真正潜力。伊隆・马斯克还开玩笑说,企业级 GPU 比买「药」还难。
图片
马斯克预测 GPT-5 大概需要三到五万块 H100 来训练。

旺盛的需求除了推动英伟达的股价,使其改变生产计划之外,也让人们不得不去寻求其他替代方式。好消息是,图形芯片市场上并不只有 N 卡一家。

昨天,卡耐基梅隆大学博士生侯博涵(Bohan Hou)放出了使用 AMD 显卡进行大模型推理的新方案,立刻获得了机器学习社区的关注。
图片
在 CMU,侯博涵的导师是 TVM、MXNET、XGBoost 的作者陈天奇。对于这项新实践,陈天奇表示,解决 AI 硬件短缺问题的方法还是要看软件,让我们带来高性能、通用部署的开源大模型吧。
图片
知乎上,作者对于实现高性能 LLM 推理进行了详细介绍:
图片
通过这种优化方法,在最新的 Llama2 的 7B 和 13B 模型中,如果用一块 AMD Radeon RX 7900 XTX 速度可以达到英伟达 RTX 4090 的 80%,或是 3090Ti 的 94%。

除了 ROCm 之外,这种 Vulkan 支持还允许我们把大模型的部署推广到其他 AMD 芯片类型上,例如具有 AMD APU 的 SteamDeck。
图片
如果粗略的比较一下规格,我们可以看到 AMD 的 RX 7900 XTX 与英伟达的 RTX 4090 和 RTX 3090 Ti 处于相近级别。

它们的显存都在 24GB,这意味着它们可以容纳相同尺寸的模型,它们都具有相似的内存带宽。

图片

但是在算力上,RTX 4090 的 FP16 性能比 7900 XTX 高两倍,而 3090 Ti 的 FP16 性能比 7900 XTX 高 1.3 倍。如果只考虑延迟敏感的大模型推理,其性能主要受内存限制,因此 FP16 性能不是这里的瓶颈。

而看价格的话,RX 7900 XTX 比 RTX 4090 便宜 40% 还多(京东上看甚至有 50%),在消费级领域里前者几乎是和 RTX 4080 对标的。

3090Ti 的价格则很难比较,毕竟那是上一代产品。但从纯硬件规格的角度来看,AMD 7900 XTX 似乎与 RTX 3090 Ti 相当。

我们知道,硬件层的算力并不一定是 AMD 长期以来在机器学习上落后的原因 —— 主要差距在于缺乏相关模型的软件支持和优化。从生态角度来看,有两个因素已开始改变现状:

  • AMD 正在努力在 ROCm 平台上增加投入。
  • 机器学习编译等新兴技术现在有助于降低跨后端的,更通用软件支持的总体成本。

研究人员深入讨论了 AMD GPU 体系与目前流行的英伟达 GPU 上高性能 CUDA 解决方案相比的表现如何。

用 ROCm 进行机器学习编译

机器学习编译

机器学习编译是一种用于编译和自动优化机器学习模型的新兴技术。MLC 解决方案不是为每个后端(如 ROCm 或 CUDA)编写特定的算子 ,而是自动生成适用于不同后端的代码。在这里,作者利用 MLC-LLM,一种基于机器学习编译的解决方案,提供了 LLM 的高性能通用部署。MLC-LLM 建立在 Apache TVM Unity 之上,后者是一个机器学习编译软件栈,提供了基于 Python 的高效开发和通用部署。MLC-LLM 为各种后端(包括 CUDA、Metal、ROCm、Vulkan 和 OpenCL)提供了最先进的性能,涵盖了从服务器级别 GPU 到移动设备(iPhone 和 Android)。

整体而言,MLC-LLM 允许用户使用基于 Python 的工作流程获取开源的大语言模型,并在包括转换计算图、优化 GPU 算子的张量 layout 和 schedule 以及在感兴趣的平台上本地部署时进行编译。
图片
面向 ROCm 的机器学习编译技术栈。

针对 AMD GPU 和 APU 的 MLC

人们对于 A 卡用于机器学习的探索其实并不鲜见,支持 AMD GPU 有几种可能的技术路线:ROCm、OpenCL、Vulkan 和 WebGPU。ROCm 技术栈是 AMD 最近推出的,与 CUDA 技术栈有许多相应的相似之处。Vulkan 是最新的图形渲染标准,为各种 GPU 设备提供了广泛的支持。WebGPU 是最新的 Web 标准,允许在 Web 浏览器上运行计算。

虽然有这么多可能的路线,但很少有解决方案支持除了 CUDA 之外的方法,这在很大程度上是因为复制新硬件或 GPU 编程模型的技术栈的工程成本过高。MLC-LLM 支持自动代码生成,无需为每个 GPU 算子重新定制,从而为以上所有方法提供支持。但是,最终性能仍然取决于 GPU 运行时的质量以及在每个平台上的可用性。

在这个案例中,作者选择 Radeon 7900 XTX 的 ROCm 和 Steamdeck 的 APU 的 Vulkan,可以发现 ROCm 技术栈是开箱即用的。由于 TVM unity 中具有高效的基于 Python 的开发流程,花费了若干小时来进一步提供 ROCm 的性能优化。具体来说,研究人员采取了以下措施来提供 ROCm 支持:

  • 重用现有后端(如 CUDA 和 Metal)的整个 MLC 流水线,包括内存规划、算子融合等。
  • 重用 TVM TensorIR 中的通用 GPU 算子优化空间,并将其后端选为 AMD GPU
  • 重用 TVM 的 ROCm 代码生成流程,通过 LLVM 生成 ROCm 代码。
  • 最后,将生成的代码导出为可以由 CLI、Python 和 REST API 调用的共享或静态库。

使用 MLC Python 包进行性能测试

作者使用 4 bit 量化对 Llama 2 7B 和 13B 进行了性能测试。通过设置 prompt 长度为 1 个 token 并生成 512 个 token 来测量 decoding 的性能。所有结果都是在 batch size=1 的情况下测试。

图片
AMD RX 7900 XTX 与 NVIDIA RTX 4090 和 3090 Ti 的性能对比。

基于 ROCm5.6,AMD 7900 XTX 可以达到 NVIDIA 4090 速度的 80%。

关于 CUDA 性能说明:在这里 CUDA baseline 的性能如何?据我们所知,MLC-LLM 是 CUDA 上大语言模型推理的最优解决方案。但作者相信它仍然有改进的空间,例如通过更好的 attention 算子优化。一旦这些优化在 MLC 中实现,预计 AMD 和 NVIDIA 的数据都会有所改善。

如果这些优化仅在 N 卡那里实施,将使差距从 20% 增加到 30%。因此,在查看这些数字时,作者建议放置 10% 的误差。

自行尝试

该项目提供了预构建的安装包和使用说明,以便用户在自己的设备上复现新的结果。要运行这些性能测试,请确保你的 Linux 上有安装了 ROCm 5.6 或更高版本的 AMD GPU。按照这里的说明(https://mlc.ai/mlc-llm/docs/get_started/try_out.html)安装启用了 ROCm 的预构建 MLC pacakge。

运行以下 Python 脚本,需要使用 MLC package 来复现性能数据:

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 benchmarksoutput = 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。请按照这里的说明(https://mlc.ai/mlc-llm/docs/deploy/cli.html#option-2-build-mlc-runtime-from-source)从源代码构建 CLI。

使用统一内存在 SteamDeck 上运行 Vulkan

作者表示,还将考察更广泛的 AMD 设备,更具体地说,是搭载了 AMD APU 的 Steam Deck。虽然在 BIOS 中,ROCm 中可用的 GPU VRAM 被限制为 4GB,但 Mesa Vulkan 驱动程序具有强大的支持,允许缓冲区超越上限,使用统一内存最多可达 16GB,足以运行 4 位量化的 Llama-7B。
图片
在 Steam Deck 上与大语言模型进行交互。

这些结果为支持更多不同类型的消费者提供了一些启示。

讨论和未来的方向

身处生成式 AI 的时代,硬件可用性已经成为一个迫切需要解决的问题。ML 编译可以通过在硬件后端之间提供高性能的通用部署,从而提高硬件的可用性。

鉴于本文所展现的数据,作者认为在适当的价格和可用性条件下,AMD GPU 可以开始用于 LLM 推理。

陈天奇团队,研究目前的重点关注消费级 GPU。作者表示,根据过往经验,针对消费级 GPU 型号的 MLC 优化通常可以推广到云 GPU(例如从 RTX 4090 到 A100 和 A10g),有信心让该解决方案在云和消费级 AMD 和 NVIDIA GPU 之间具有普适性,并将在获得更多 GPU 访问权限后更新研究。与此同时,作者期待研究社区在 MLC 通用部署流程的基础上构建解决方案。

本文是通过 MLC 来支持高效通用的机器学习部署的研究的一个阶段性努力,研究人员也正积极地在以下几个方向上努力推广成果:

  • 启用 batching 和多 GPU 支持(对服务器端推理尤为重要);
  • 与 PyTorch 生态系统集成;
  • 支持更多量化和模型架构;
  • 在更多硬件后端上进行更多自动优化。

机器学习系统工程是一个持续的问题,在持续创新方面,英伟达仍然处于领先地位,作者预计随着新硬件(如 H100)以及更重要的软件演变,整个领域将发生变化。因此,关键问题不仅是现在构建正确的解决方案,还包括如何不断赶上并将机器学习工程引入新平台。在这个过程中,机器学习工程的生产力是关键。

由于基于 Python 的 ML 编译开发流程,我们可以在几小时内获得 ROCm 优化的支持。预计此次提出的新方法,在探索更多关于通用部署的想法并解决硬件可用性问题时会变得更加有用。

相关资源

该项目已经在 GitHub 上发布。有关如何尝试 MLC LLM 部署的详细指南,请参阅项目页面。MLC LLM 的源代码可在官方 GitHub 上找到。

  • 项目页面:https://mlc.ai/mlc-llm/docs/
  • GitHub:https://github.com/mlc-ai/mlc-llm/

参考内容:
https://blog.mlc.ai/2023/08/09/Making-AMD-GPUs-competitive-for-LLM-inference
https://zhuanlan.zhihu.com/p/649088095

MLC-机器学习编译
MLC-机器学习编译

MLC-机器学习编译社区专栏,专注于机器学习系统和机器学习编译技术,不定期更新包括TVM、MLC-LLM等相关内容介绍。

https://mlc.ai
产业英伟达
2
相关数据
陈天奇人物

陈天奇,华盛顿大学计算机系博士生,此前毕业于上海交通大学ACM班,研究方向为大规模机器学习。陈天奇曾获得KDD CUP 2012 Track 1第一名,并开发了SVDFeature,XGBoost,cxxnet等著名机器学习工具,是最大开源分布式机器学习项目DMLC的发起人之一。

机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

人工智能技术

在学术研究领域,人工智能通常指能够感知周围环境并采取行动以实现最优的可能结果的智能体(intelligent agent)

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

规划技术

人工智能领域的「规划」通常是指智能体执行的任务/动作的自动规划和调度,其目的是进行资源的优化。常见的规划方法包括经典规划(Classical Planning)、分层任务网络(HTN)和 logistics 规划。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

XGBoost技术

XGBoost是一个开源软件库,为C ++,Java,Python,R,和Julia提供了渐变增强框架。 它适用于Linux,Windows,MacOS。从项目描述来看,它旨在提供一个“可扩展,便携式和分布式的梯度提升(GBM,GBRT,GBDT)库”。 除了在一台机器上运行,它还支持分布式处理框架Apache Hadoop,Apache Spark和Apache Flink。 由于它是许多机器学习大赛中获胜团队的首选算法,因此它已经赢得了很多人的关注。

京东机构

京东(股票代码:JD),中国自营式电商企业,创始人刘强东担任京东集团董事局主席兼首席执行官。旗下设有京东零售、京东物流、京东科技子集团、印尼&泰国海外合资跨境电商等核心业务。2013年正式获得虚拟运营商牌照。2014年5月在美国纳斯达克证券交易所正式挂牌上市。 2016年6月与沃尔玛达成深度战略合作。

https://www.jd.com
相关技术
知乎机构

知乎,中文互联网综合性内容平台,自 2010 年成立以来,知乎凭借认真、专业、友善的社区氛围,独特的产品机制,以及结构化、易获得的优质内容,聚集了中文互联网科技、商业、影视、时尚、文化等领域最具创造力的人群,已成为综合性、全品类,在诸多领域具有关键影响力的内容平台。知乎将AI广泛应用与社区,构建了人、内容之间的多元连接,提升了社区的运转效率和用户体验。知乎通过内容生产、分发,社区治理等领域的AI应用,也创造了独有的技术优势和社区AI创新样本。

zhihu.com
AMD机构

超威半导体(中国)有限公司专门为计算机、通信和消费电子行业设计和制造各种创新的微处理器(CPU、GPU、主板芯片组、电视卡芯片等),以及提供闪存和低功率处理器解决方案,公司成立于1969年。AMD致力为技术用户——从企业、政府机构到个人消费者——提供基于标准的、以客户为中心的解决方案。

https://www.amd.com/zh-hans
语言模型技术

统计式的语言模型是借由一个几率分布,而指派几率给字词所组成的字串。语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。

自动代码生成技术

自动代码生成(Code Generation),指生成程序的程序,更详细的是指通过读取工程中设计的各种文档或者模型,按照一定的领域规则,生成规范的计算机能理解的源代码的过程。

量化技术

深度学习中的量化是指,用低位宽数字的神经网络近似使用了浮点数的神经网络的过程。

推荐文章
A卡能不能多卡训练呢?