引言
在最近的 PyTorch 大会上，Arm 强调了其技术从云端到边缘设备的广泛影响，并强调致力于将其先进的 AI 计算能力无缝地提供给全球数百万开发者。
在演示中，Arm 强调其肩负着巨大责任，即为 2000 多万开发者和数十亿用户提供先进的 AI 计算功能，且力求顺畅无阻。实现这一目标需要软件和硬件合作伙伴组成的庞大生态系统之间的关键软件协作。
就在几个月前，Arm 推出了 Arm Kleidi，这是一系列开发者赋能技术和资源，旨在推动整个机器学习（ML）堆栈的技术协作和创新。其中包括 KleidiAI 软件库，它提供优化的软件例程，集成到 XNNPACK 等关键框架后，可为 Arm Cortex-A CPU 上的开发者实现自动 AI 加速。
今天，我们很高兴地宣布 AI 开源社区的一个新里程碑，这使得 Arm 更接近实现其愿景：通过 XNNPACK 将 KleidiAI 集成到 ExecuTorch 中，从而提升 Arm 移动 CPU 上的 AI 工作负载性能！
得益于 Arm 和 Meta 工程团队的协作努力，AI 开发者现在可以部署量化的 Llama 模型，这些模型在具有 i8mm ISA 扩展的 Arm Cortex-A v9 CPU 上运行时速度可提高 20%。
还有更多激动人心的消息——ExecuTorch 团队已正式发布 Beta 版本！
这标志着我们合作关系中的一个重要里程碑。在本博客中，我们渴望分享关于 ExecuTorch 功能、新的 Meta Llama 3.2 模型、带块级量化的整数 4 位量化以及在某些 Arm CPU 上记录的令人印象深刻的性能的更多细节。值得注意的是，正如以下截图所示，我们在 Samsung S24+ 设备上使用量化的 Llama 3.2 1B 模型在预填充阶段实现了每秒超过 350 个 token 的速度。
现在，让我们深入了解实现前面图像中所示演示创建的关键组件。首先是：新的 Llama 3.2 模型！
Meta Llama 3.2
Meta 最近发布了首批轻量级量化 Llama 模型，旨在运行在流行的移动设备上。Meta 使用了两种技术来量化 Llama 3.2 1B 和 3B 模型：使用 LoRA 适配器（QLoRA）的量化感知训练（QAT）以及最先进的训练后量化方法 SpinQuant。使用 PyTorch 的 ExecuTorch 框架作为推理引擎，并以 Arm CPU 作为后端，对量化模型进行了评估。
这些指令调优模型保留了原始 1B 和 3B 模型的质量和安全性，同时实现了 2-4 倍的加速，并且与原始 BF16 格式相比，模型大小平均减少 56%，内存占用平均减少 41%。
在本篇博客文章中，我们将展示我们在实验中观察到的性能提升。
ExecuTorch
ExecuTorch 是一个专为在设备上部署 AI 模型而设计的 PyTorch 原生框架，可增强隐私性并降低延迟。它支持部署最先进的开源 AI 模型，包括 Llama 系列模型以及像 Segment Anything 和 Seamless 这样的视觉和语音模型。
这为移动电话、智能眼镜、VR 头戴设备和智能家居摄像头等边缘设备开启了新的可能性。传统上，将 PyTorch 训练的 AI 模型部署到资源有限的边缘设备是一项具有挑战性且耗时的工作，通常需要转换为其他格式，这可能导致错误和次优的性能。硬件和边缘生态系统中多样化的工具链也降低了开发者体验，使得通用解决方案难以实现。
ExecuTorch 通过提供可组合的组件来解决这些问题，这些组件包括核心运行时、算子库和委托接口，从而实现可移植性和可扩展性。可以使用 torch.export() 导出模型，生成与 ExecuTorch 运行时原生兼容的图，能够在大多数带有 CPU 的边缘设备上运行，并可扩展到 GPU 和 NPU 等专用硬件以提高性能。
与 Arm 合作，ExecuTorch 现在利用 Arm KleidiAI 库中优化的低位矩阵乘法核，通过 XNNPACK 提高设备上的大语言模型（LLM）推理性能。我们也感谢 Google 的 XNNPACK 团队为此提供的支持。
在本文中，我们将重点介绍 ExecuTorch 中提供的此集成
演进 AI 工作负载的架构
在 Arm，自深度学习浪潮早期以来，我们一直致力于投资开源项目并推进处理器中的新技术，重点是使 AI 工作负载高性能且更具能效。
例如，Arm 从 Armv8.2-A 架构开始引入了 SDOT 指令，用于加速 8 位整数向量的点积运算。这项功能现已广泛应用于移动设备中，显著加快了量化 8 位模型的计算速度。在 SDOT 指令之后，Arm 又引入了 BF16 数据类型和 MMLA 指令，进一步提升了 CPU 上的浮点和整数矩阵乘法性能，最近还宣布了可扩展矩阵扩展（SME），标志着机器学习能力的重大飞跃。
下图展示了 Arm CPU 在过去十年中在 AI 领域持续创新的几个示例
鉴于 Arm CPU 的广泛使用，AI 框架需要在关键算子中充分利用这些技术以实现性能最大化。认识到这一点，我们看到了构建一个开源库来分享这些优化软件例程的需求。然而，我们也注意到了将新库集成到 AI 框架中的挑战，例如关于库大小、依赖关系和文档的担忧，以及避免增加开发者额外负担的需求。因此，我们采取了额外措施，收集合作伙伴的反馈，并确保顺畅的集成过程，该过程不需要为 AI 开发者引入额外的依赖关系。这项努力催生了 KleidiAI，这是一个开源库，为针对 Arm CPU 的人工智能（AI）工作负载提供了优化的性能关键型例程。您可以在此处了解更多关于 KleidiAI 的信息。
与 Meta 的 ExecuTorch 团队合作，Arm 为他们新颖的带块级量化的 4 位量化方案提供了软件优化，该方案用于加速 Llama 3.2 量化模型在 Transformer 层 torch.nn.linear 算子中的矩阵乘法核。ExecuTorch 的这种灵活的 4 位量化方案在模型精度和针对设备上 LLM 的低位矩阵乘法性能之间取得了平衡。
带块级量化的整数 4 位量化
在 KleidiAI 中，我们引入了为此新的 4 位整数量化方案优化的微核（matmul_clamp_f32_qai8dxp_qsi4c32p）
如下图所示，这种 4 位量化对权重（RHS 矩阵）使用块级量化策略，对激活（LHS 矩阵）使用行级 8 位量化
正如您在前面图像中看到的，权重矩阵中的每个输出特征图（OFM）被划分为大小相等的块（组大小），每个块都有一个以 BF16 格式存储的缩放因子。BF16 的优势在于，它以一半的位数大小保持了 32 位浮点（FP32）格式的动态范围，并且通过简单的移位操作可以轻松地在 BF16 和 FP32 之间进行转换。这使得 BF16 非常适合节省模型空间、保持精度，并确保与缺少 BF16 硬件加速的设备向后兼容。您可以在这篇 Arm 社区博客文章中了解更多关于 BF16 格式的信息。
为了完整起见，此 4 位量化方案以及我们在 KleidiAI 中的实现允许用户配置线性权重（RHS）的组大小，从而允许他们在模型大小、模型精度和模型性能之间进行权衡（如果模型由用户量化）。
此时，我们准备揭示在 Arm CPU 上使用 ExecuTorch 运行 Llama 3.2 1B 和 Llama 3.2 3B 时记录的惊人性能。首先，让我们回顾一下我们将用于评估 LLM 推理性能的指标。
LLM 推理性能指标
通常用于评估 LLM 推理性能的指标包括
- 首个 Token 生成时间 (TTFT)：衡量用户提供提示后生成首个输出 token 所需的时间。此延迟或响应时间对于良好的用户体验至关重要，尤其是在手机上。TTFT 也与提示或提示 token 的长度有关。为了使此指标独立于提示长度，我们在此使用每秒预填充 token 数作为代理。两者之间的关系是反向的：TTFT 越低，每秒预填充 token 数越高。
- 解码性能：指每秒生成的平均输出 token 数，以 Tokens/Second（每秒 token 数）报告。它与生成的 token 总数无关。对于设备上推理，保持此值高于用户的平均阅读速度非常重要。
- 运行时峰值内存：此指标反映了以预期性能运行模型所需的 RAM 量，通常以兆字节 (MiB) 报告。考虑到 Android 和 iOS 设备上有限的 RAM 量，这是设备上部署 LLM 的关键指标之一。它决定了可以在设备上部署的模型类型。
结果
量化的 Llama 3.2 1B 模型，包括 SpinQuant 和 QLoRA，旨在在各种 RAM 有限的手机上高效运行。在本节中，我们展示了量化的 Llama 3.2 1B 模型在预填充阶段可以实现每秒超过 350 个 token 的速度，在解码阶段可以实现每秒超过 40 个 token 的速度。仅使用 Arm CPU，这种性能水平足以支持设备上文本摘要功能并提供合理的用户体验。举例来说，平均 50 条未读消息约包含 600 个 token。有了这种性能，响应时间（即首个生成的词出现在屏幕上的时间）大约为两秒。
我们展示了在运行原生 Android 的 Samsung S24+ 上的测量结果。我们在此实验中使用了 Llama 3.2 1B 参数模型。尽管我们仅展示了使用 1B 模型的结果，但对于 3B 参数模型也可以预期类似的性能提升。实验设置包括一次热身运行，序列长度为 128，提示长度为 64，使用 8 个可用 CPU 中的 6 个，并通过 adb 测量结果。
使用 GitHub 上的 ExecuTorch 主分支，我们首先使用已发布的检查点为每个模型生成了 ExecuTorch PTE 二进制文件。然后，使用同一仓库，我们生成了适用于 Armv8 的 ExecuTorch 运行时二进制文件。在本节的其余部分，我们将使用使用 KleidiAI 构建的二进制文件，比较不同量化 1B 模型与 BF16 模型的性能。我们还将比较使用 KleidiAI 和不使用 KleidiAI 的二进制文件在量化模型上的性能提升，以提炼出 KleidiAI 的影响。
量化模型性能
与基线 BF16 相比，Llama 3.2 量化模型（SpinQuant 和 QLoRA）在提示预填充和文本生成（解码）方面表现显著更好。我们观察到解码性能提高了 2 倍以上，预填充性能提高了 5 倍以上。
此外，量化模型的大小，即 PTE 文件的大小（以字节为单位），不到 BF16 模型的一半，分别为 2.3 GiB 和 1.1 GiB。虽然 int4 的大小是 BF16 的四分之一，但模型中的某些层使用 int8 进行量化，使得 PTE 文件大小的比例变大。我们观察到运行时峰值内存占用显著减少了近 40%，从 BF16 模型的 3.1 GiB 减少到 SpinQuant 模型的 1.9 GiB，这是在最大序列长度为 2048 时以常驻集大小（RSS）测量的。
凭借全面的改进，新的量化 Llama 3.2 模型是针对 Arm CPU 进行设备上部署的理想选择。有关精度详情，请查阅 Meta Llama 3.2 博客。
KleidiAI 的影响
ExecuTorch 依赖 Arm KleidiAI 库为具有先进 Armv8/9 ISA 功能的最新 Arm CPU 提供高性能的低位矩阵乘法核。这些核用于 ExecuTorch 中的设备上量化 Llama 3.2 模型推理。如下图所示，与非 KleidiAI 核相比，ExecuTorch 在 S24+ 上使用 KleidiAI 时，预填充性能平均提高了 20% 以上，同时保持了相同的精度。这种性能优势不限于特定模型或设备，预计将惠及所有在 Arm CPU 上使用低位量化矩阵乘法的 ExecuTorch 模型。
为了评估 Kleidi 的影响，我们生成了针对 Arm Cortex-A CPU 的两个 ExecuTorch 运行时二进制文件并比较了它们的性能。
- 第一个 ExecuTorch 运行时二进制文件是通过 XNNPACK 库使用 Arm KleidiAI 库构建的。
- 第二个二进制文件是在没有 Arm KleidiAI 仓库的情况下构建的，使用了 XNNPACK 库中的原生核。
亲自动手试试吧！
准备好亲身体验性能提升了吗？以下是您在项目中试用利用 KleidiAI 优化功能的 ExecuTorch 的方法：此处提供了 Arm 的学习路径链接，可帮助您开始使用 ExecuTorch 和 KleidiAI 开发您自己的使用 LLM 的应用程序。
我们期待收到您的反馈！
