0
  • 聊天消息
  • 系统消息
  • 评论与回复
登录后你可以
  • 下载海量资料
  • 学习在线课程
  • 观看技术视频
  • 写文章/发帖/加入社区
会员中心
创作中心

完善资料让更多小伙伴认识你,还能领取20积分哦,立即完善>

3天内不再提示

利用Arm Kleidi技术实现PyTorch优化

Arm社区 来源:Arm社区 2024-12-23 09:19 208次阅读

作者:Arm 基础设施事业部高级产品经理 Ashok Bhat

PyTorch 是一个广泛应用的开源机器学习 (ML) 库。近年来,Arm 与合作伙伴通力协作,持续改进 PyTorch 的推理性能。本文将详细介绍如何利用 Arm Kleidi 技术提升 Arm Neoverse 平台上的 PyTorch 推理表现。Kleidi 技术可以通过 Arm Compute Library (ACL) 和 KleidiAI 库获取。

PyTorch 提供两种主要的执行模式:即时执行模式 (Eager Mode) 和图模式 (Graph Mode)。即时执行模式是一种动态执行模式,操作会以 Python 代码编写的方式立即执行,该模式非常适用于实验与调试。而图模式则是在执行前将一系列 PyTorch 操作编译成静态计算图,从而实现性能优化和高效的硬件加速。通过使用 torch.compile 函数,可以方便地将 PyTorch 代码转换为图模式,通常能够显著提升执行速度。

PyTorch 即时执行模式

CPU 推理性能提升高达三倍

PyTorch 即时执行模式使用 oneDNN,针对具有 ACL 内核的 Arm Neoverse 处理器进行了优化。可以通过以下的 PyTorch 软件栈图进行了解。

b5a46902-bead-11ef-9310-92fbcf53809c.jpg

图 1:PyTorch 软件栈

PyTorch 中的 FX Graph 是用于可视化和优化 PyTorch 模型的一种中间表示形式。

ATen 是支撑 PyTorch 框架的基础张量库。它提供了核心张量类别和大量数学运算,构成了 PyTorch 模型的基本组件。

oneDNN 是一个性能库,为包括 Arm 和 x86 在内的各种硬件架构提供常见深度学习原语的优化实现方案。在这些架构上,ATen 使用 oneDNN 作为性能增强后端。这意味着当 PyTorch 遇到支持的操作时,会将计算委托给 oneDNN,后者可以使用针对特定硬件的优化来提高执行效率。

Arm Compute Library 于 2016 年首次发布,提供针对 Arm 进行优化的关键 ML 原语,包括卷积、池化、激活函数、全连接层、归一化。这些原语利用 Arm Neoverse 核心上针对特定 ML 和特定硬件的功能和指令来实现高性能。我们已将 Arm Compute Library 集成到 oneDNN 中,以便在 Arm 平台上加速 ATen 操作。

Arm Neoverse CPU 包含有助于加速 ML 的硬件扩展,其中包括 Neon、SVE/SVE2、BF16 和 I8MM,通过有效地进行向量处理、BF16 运算和矩阵乘法来加速 ML 任务。

b6220aa6-bead-11ef-9310-92fbcf53809c.jpg

图 2:各种模型在即时执行模式下实现的性能提升

PyTorch 图模式(使用 torch.compile)

比 PyTorch 即时执行模式进一步提高两倍

PyTorch 2.0 引入了 torch.compile,与默认的即时执行模式相比,可提高 PyTorch 代码的速度。与即时执行模式不同,torch.compile 将整个模型预编译成针对特定硬件平台优化的单图。从 PyTorch 2.3.1 开始,官方 AArch64 安装包均包含 torch.compile 优化。在基于亚马逊科技 (AWS) Graviton3 的 Amazon EC2 实例上,对于各种自然语言处理 (NLP)、计算机视觉 (CV) 和推荐模型,这些优化可以为 TorchBench 模型推理带来比即时执行模式高出两倍的性能。

b63c4b14-bead-11ef-9310-92fbcf53809c.jpg

图 3:各种模型在编译模式下实现的性能提升

下一步通过 KleidiAI 库

提升生成式 AI 推理性能

目前,我们已经研究了 Arm Compute Library 如何在即时执行模式和编译模式下提升 PyTorch 推理性能。接下来,我们来看 PyTorch 即将推出什么新功能。Arm 目前正在努力提升 PyTorch 中的大语言模型 (LLM) 推理性能,并以 Llama 和 Gemma 为主要 LLM 示例。

经优化的 INT4 内核

今年早些时候,Arm 软件团队和合作伙伴共同优化了 llama.cpp 中的 INT4 和 INT8 内核,以利用更新的 DOT 和 MLA 指令。在 AWS Graviton3 处理器上,这些内核在即时评估方面比现有 GEMM MMLA 内核提升了 2.5 倍,并且在文本生成方面比默认的 vec_dot 内核提升了两倍。这些经优化的新内核也是 Arm KleidAI 库的一部分。

今年早些发布的 KleidiAI 库是一个开源库,具有针对 Arm CPU 上的 AI 任务进行优化的微内核。对于微内核,可将它视为能够提升特定 ML 操作性能的软件。开发者可以通过包含相关的 .c 和 .h 文件及公共头文件来使用这些微内核。无需包含库的其余部分。

Kleidi 与 PyTorch 的集成

b65c141c-bead-11ef-9310-92fbcf53809c.jpg

图 4:Kleidi 技术与 PyTorch 的集成

我们引入了两种新的 ATen 操作:torch.ops.aten._kai_weights_pack_int4() 和 torch.ops.aten._kai_input_quant_mm_int4(),两者均使用 KleidiAI 库中高度优化的打包技术和 GEMM 内核。gpt-fast 利用这些 PyTorch 算子来 (1) 使用对称的每通道量化将权重量化为 INT4,并添加包含量化尺度的额外数组;(2) 动态量化激活矩阵并使用 AArch64 I8MM 扩展来执行激活矩阵和权重的 INT8 矩阵乘法。

b674764c-bead-11ef-9310-92fbcf53809c.jpg

图 5:通过在 PyTorch 中集成 KleidiAI

来提升 4 位量化 LLM 模型性能

通过这种方法,与目前默认的 PyTorch 实现方案相比,我们可以将 Llama 的推理性能提高 18 倍,将 Gemma 的推理性能提高 14 倍。

结论

Arm 及其合作伙伴利用 Arm Compute Library 中的 Kleidi 技术提高了 Arm Neoverse 平台上的 PyTorch 推理性能。在即时执行模式下可实现高达两倍的性能提升,在图模式下(使用 torch.compile)可再提升两倍。此外,我们还在努力将生成式 AI 模型(Llama 和 Gemma)的推理性能提升高达 18 倍。

Arm 通过部署 Kleidi 技术来实现PyTorch 上的优化,以加速在基于 Arm 架构的处理器上运行 LLM 的性能。Arm 技术专家在基于 Neoverse V2 的 AWS Graviton4 R8g.4xlarge EC2 实例上运行 Llama 3.1 展示了所实现的性能提升。如果你对这一演示感兴趣,可阅读《Arm KleidiAI 助力提升 PyTorch 上 LLM 推理性能》了解。

声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉
  • ARM
    ARM
    +关注

    关注

    134

    文章

    9111

    浏览量

    368047
  • cpu
    cpu
    +关注

    关注

    68

    文章

    10882

    浏览量

    212241
  • 机器学习
    +关注

    关注

    66

    文章

    8425

    浏览量

    132776
  • pytorch
    +关注

    关注

    2

    文章

    808

    浏览量

    13256
  • Neoverse
    +关注

    关注

    0

    文章

    9

    浏览量

    4612

原文标题:如何在 Arm Neoverse 平台上使用 Kleidi 技术加速 PyTorch 推理?

文章出处:【微信号:Arm社区,微信公众号:Arm社区】欢迎添加关注!文章转载请注明出处。

收藏 人收藏

    相关推荐

    Arm KleidiAI助力提升PyTorch上LLM推理性能

    热门的深度学习框架尤为突出,许多企业均会选择其作为开发 AI 应用的库。通过部署 Arm Kleidi 技术Arm 正在努力优化
    的头像 发表于 12-03 17:05 852次阅读
    <b class='flag-5'>Arm</b> KleidiAI助力提升<b class='flag-5'>PyTorch</b>上LLM推理性能

    ARM程序设计优化策略与技术

    程序优化是指软件编程结束后,利用软件开发工具对程序进行调整和改进,让程序充分利用资源, 提高运行效率, 缩减代码尺寸的过程。按照优化的侧重点不同, 程序
    发表于 07-07 11:06

    Pytorch模型训练实用PDF教程【中文】

    ?模型部分?还是优化器?只有这样不断的通过可视化诊断你的模型,不断的对症下药,才能训练出一个较满意的模型。本教程内容及结构:本教程内容主要为在 PyTorch 中训练一个模型所可能涉及到的方法及函数,并且
    发表于 12-21 09:18

    在Ubuntu 18.04 for Arm上运行的TensorFlow和PyTorch的Docker映像

    ,并做出优化实现尽可能高的性能。我们希望这些 Docker 镜像和创建它们的方法对希望在 AArch64 上使用 TensorFlow 和 PyTorch 的人有所帮助。包括什么?构建
    发表于 10-14 14:25

    解读最佳实践:倚天 710 ARM 芯片的 Python+AI 算力优化

    编者按: 在刚刚结束的 PyCon China 2022 大会上,龙蜥社区开发者朱宏林分享了主题为《ARM 芯片的 Python+AI 算力优化》的技术演讲。本次演讲,作者将向大家介绍他们在倚天
    发表于 12-23 16:02

    ARM编译器优化版本1.0

    ARM编译器armcc可以优化您的代码以实现小代码和高性能。 本教程介绍了编译器执行的主要优化技术,并解释了如何控制编译器
    发表于 08-28 07:11

    基于ARM的EPA通信协议栈优化技术的研究与实现

    本文叙述基于ARM的EPA通信协议栈优化技术的研究与实现的案例分析。
    发表于 10-13 16:48 32次下载
    基于<b class='flag-5'>ARM</b>的EPA通信协议栈<b class='flag-5'>优化</b><b class='flag-5'>技术</b>的研究与<b class='flag-5'>实现</b>

    如何利用ARM Mali在机器视觉技术实现应用优势

    谈GPU计算的优势,并阐述生态系统合作伙伴是如何利用ARM Mali 在各种图像处理,多媒体和计算机视觉使用案例来实现这些优势。
    的头像 发表于 07-02 13:13 2869次阅读

    Pytorch实现MNIST手写数字识别

    Pytorch 实现MNIST手写数字识别
    发表于 06-16 14:47 7次下载

    PyTorch教程12.1之优化和深度学习

    电子发烧友网站提供《PyTorch教程12.1之优化和深度学习.pdf》资料免费下载
    发表于 06-05 15:08 0次下载
    <b class='flag-5'>PyTorch</b>教程12.1之<b class='flag-5'>优化</b>和深度学习

    PyTorch教程19.1之什么是超参数优化

    电子发烧友网站提供《PyTorch教程19.1之什么是超参数优化.pdf》资料免费下载
    发表于 06-05 10:25 0次下载
    <b class='flag-5'>PyTorch</b>教程19.1之什么是超参数<b class='flag-5'>优化</b>

    PyTorch教程19.2之超参数优化API

    电子发烧友网站提供《PyTorch教程19.2之超参数优化API.pdf》资料免费下载
    发表于 06-05 10:27 0次下载
    <b class='flag-5'>PyTorch</b>教程19.2之超参数<b class='flag-5'>优化</b>API

    PyTorch教程19.4之多保真超参数优化

    电子发烧友网站提供《PyTorch教程19.4之多保真超参数优化.pdf》资料免费下载
    发表于 06-05 10:45 0次下载
    <b class='flag-5'>PyTorch</b>教程19.4之多保真超参数<b class='flag-5'>优化</b>

    2.0优化PyTorch推理与AWS引力子处理器

    2.0优化PyTorch推理与AWS引力子处理器
    的头像 发表于 08-31 14:27 630次阅读
    2.0<b class='flag-5'>优化</b><b class='flag-5'>PyTorch</b>推理与AWS引力子处理器

    Arm推出AI优化Arm终端CSS以及新的Arm Kleidi软件

    Arm 终端计算子系统 (CSS) 作为新的计算解决方案,结合了 Armv9 架构的优势,以及基于三纳米工艺节点,经过验证和证实为生产就绪的新 Arm CPU 和 GPU 实现,可赋能芯片合作伙伴快速创新,
    的头像 发表于 05-30 11:33 680次阅读
    <b class='flag-5'>Arm</b>推出AI<b class='flag-5'>优化</b>的<b class='flag-5'>Arm</b>终端CSS以及新的<b class='flag-5'>Arm</b> <b class='flag-5'>Kleidi</b>软件