简介
在最近的 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 上记录的令人印象深刻的性能的更多详细信息。值得注意的是，我们在配备量化 Llama 3.2 1B 模型的 Samsung S24+ 设备上，在预填充阶段实现了每秒超过 350 个 token 的速度，如下面的屏幕截图所示。
现在，让我们深入了解实现上述图像中演示创建的关键组件。首先：新的 Llama 3.2 模型！
Meta Llama 3.2
Meta 最近宣布了首批轻量级量化 Llama 模型，这些模型旨在在流行的移动设备上运行。Meta 使用了两种技术来量化 Llama 3.2 1B 和 3B 模型：带有 LoRA 适配器的量化感知训练 (QAT) (QLoRA) 和最先进的后训练量化方法 SpinQuant。量化模型使用 PyTorch 的 ExecuTorch 框架作为推理引擎进行评估，并以 Arm CPU 作为后端。
这些指令微调模型保留了原始 1B 和 3B 模型的质量和安全性，同时实现了 2-4 倍的加速，并将模型大小平均减少了 56%，内存占用平均减少了 41%（与原始 BF16 格式相比）。
在本博客文章中，我们将演示我们在实验中观察到的性能改进。
ExecuTorch
ExecuTorch 是一个 PyTorch 原生框架，专门为在设备上部署 AI 模型而设计，旨在增强隐私并减少延迟。它支持部署尖端的开源 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 的诞生，KleidiAI 是一个开源库，为专为 Arm CPU 量身定制的人工智能 (AI) 工作负载提供优化的性能关键例程。您可以在 此处 了解有关 KleidiAI 的更多信息。
通过与 Meta 的 ExecuTorch 团队合作，Arm 为他们新型的每块量化方案的 4 位提供了软件优化，该方案用于加速 Transformer 层的 torch.nn.linear 运算符中 Llama 3.2 量化模型的矩阵乘法内核。ExecuTorch 提供的这种灵活的 4 位量化方案在模型精度和面向设备端 LLM 的低位矩阵乘法性能之间取得了平衡。
带有逐块量化的整数 4 位
在 KleidiAI 中，我们为这种新的 4 位整数量化方案引入了微内核 (matmul_clamp_f32_qai8dxp_qsi4c32p)
如下图所示，这种 4 位量化对权重（RHS 矩阵）量化使用逐块策略，对激活（LHS 矩阵）使用每行 8 位量化
正如您在上图中看到的，权重矩阵中的每个输出特征图 (OFM) 都被划分为大小相等的块（组大小），每个块都有一个以 BF16 格式存储的比例因子。BF16 的优势在于它保持了 32 位浮点 (FP32) 格式的动态范围，但位大小只有一半，并且使用简单的移位操作即可轻松地与 FP32 相互转换。这使得 BF16 非常适合节省模型空间、保持精度并确保与缺少 BF16 硬件加速的设备的向后兼容性。您可以在 这篇 Arm 社区博客文章中了解有关 BF16 格式的更多信息。
为了完整起见，这种 4 位量化方案以及我们在 KleidiAI 中的实现允许用户配置线性权重 (RHS) 的组大小，从而允许他们在模型大小、模型精度和模型性能之间进行权衡（如果模型由用户量化）。
此时，我们已准备好揭晓在使用 ExecuTorch 运行 Llama 3.2 1B 和 Llama 3.2 3B 时，Arm CPU 上记录的令人难以置信的性能。让我们首先回顾一下我们将用于评估 LLM 推理性能的指标。
LLM 推理的指标
通常，用于评估推理期间 LLM 性能的性能指标包括
- 首个 token 时间 (TTFT)：这衡量的是用户提供提示后生成第一个输出 token 所需的时间。此延迟或响应时间对于良好的用户体验非常重要，尤其是在手机上。TTFT 也与提示或提示 token 的长度有关。为了使此指标独立于提示长度，我们在此处使用预填充 token/秒作为代理。两者之间的关系是反向的：TTFT 越低，预填充 token/秒越高。
- 解码性能：这是每秒生成的平均输出 token 数，因此以 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 在配备 KleidiAI 的 S24+ 上实现了平均 >20% 的预填充性能提升，同时保持了相同的精度。这种性能优势不限于特定模型或设备，预计将使所有在 Arm CPU 上使用低位量化矩阵乘法的 ExecuTorch 模型受益。
为了评估 Kleidi 的影响，我们生成了两个以 Arm Cortex-A CPU 为目标的 ExecuTorch 运行时二进制文件，并比较了它们的性能。
- 第一个 ExecuTorch 运行时二进制文件通过 XNNPACK 库使用 Arm KleidiAI 库构建。
- 第二个二进制文件在没有 Arm KleidiAI 存储库的情况下构建，使用来自 XNNPACK 库的本机内核。
亲自尝试！
准备好亲身体验性能改进了吗？以下是如何在您的项目中使用 ExecuTorch 和 KleidiAI 提供的优化：这里有一个来自 Arm 的 学习路径链接，可开始使用 LLM、ExecuTorch 和 KleidiAI 开发您自己的应用程序。
我们期待收到您的反馈！
