作者：Scott Roy, Digant Desai, Kimish Patel

我们很高兴地宣布，已在 PyTorch 的原生低精度库 TorchAO 中为 Arm CPU 添加了具有低比特（1-8 比特）权重的嵌入操作，以及具有 8 比特动态量化激活和低比特（1-8 比特）权重的线性操作。这些操作可在所有 PyTorch 界面上无缝工作，包括 eager 模式、torch.compile、AOTI 和 ExecuTorch，并且可在 torchchat 中使用。

在开发这些线性操作时，我们的重点是在 PyTorch 和 ExecuTorch 之间共享代码，并在更高层级的操作和更底层级的内核之间建立清晰的界限。这种设计允许第三方供应商轻松替换自己的内核。我们还着手创建一个用于实验的场所和基础设施，以测试新的 CPU 量化想法并在 PyTorch 生态系统中进行验证。

通用低比特内核

目前没有硬件支持低比特算术。在我们称之为通用内核的设计中，我们以模块化的方式明确分离了将低比特值解压为 int8 值的逻辑以及 int8 GEMV 内核逻辑。我们从一个 8 比特内核开始，例如这个使用 Arm neondot 指令的1x8 8 比特 GEMV 内核。在 8 比特内核中，我们调用一个内联解压例程将低比特值转换为 int8 值。这个解压例程是强制内联的，并根据某些低比特值进行模板化。我们的实验表明，使用独立的强制内联解压例程与直接内嵌解压代码之间没有性能差异。

这种模块化设计的优点是提高了开发速度和代码可维护性。在编写了 8 比特内核后，通过编写简单的比特打包例程，我们很快实现了全面的低比特覆盖。实际上，负责比特打包例程的开发者无需成为 GEMV/GEMM 内核编写方面的专家。我们还将线性内核中的相同比特打包例程在嵌入内核中重复使用。未来，我们可以将相同的比特打包例程用于通用 GEMM 内核或基于 fma 或 i8mm 指令的内核。

PyTorch 与 ExecuTorch 之间的代码共享

为了实现 PyTorch 与 ExecuTorch 之间的代码共享，我们编写了使用原始指针而非 PyTorch 张量的内核。此外，我们将线性操作实现为一个头文件，该头文件包含在单独的 PyTorch 和 ExecuTorch 操作注册代码中。通过仅使用 ATen 和 ExecuTorch 张量共有的特性，我们确保了两个框架之间的兼容性。对于多线程计算，我们引入了torchao::parallel_1d，它根据编译时标志编译为at::parallel_for 或ExecuTorch 的线程池。

可插拔内核

我们对更高层级多线程线性操作的设计与更底层级单线程内核无关，允许第三方供应商替换自己的实现。操作与内核之间的接口由ukernel 配置定义，该配置指定了用于准备激活数据、准备权重数据和运行内核的函数指针。负责分块和调度的操作仅通过此配置与内核交互。

性能

下表展示了在配备 32GB RAM 的 M1 Macbook Pro 上使用 6 个 CPU 线程运行 Llama3.1 8B 模型的 token 生成性能。

位宽 x torch.compile（解码 tokens/秒） ExecuTorch（解码 tokens/秒） ExecuTorch PTE 大小 (GiB)
1 24.18 17.86 1.46
2 27.02 19.65 2.46
3 21.01 22.25 3.46
4 19.51 19.47 4.47
5 14.78 16.34 5.47
6 12.80 13.61 6.47
7 8.16 11.73 7.48

结果是在配备 32GB RAM 和 6 个线程的 M1 Macbook Pro（带有 8 个性能核心和 2 个能效核心）上使用 torchchat 运行的。在每个测试中，生成了最大序列长度为 128 个 tokens。对于每个位宽 x，嵌入层按组量化为 x 比特，组大小为 32。在线性层中，激活按 token 动态量化为 8 比特，权重按组量化为 x 比特，组大小为 256。我们在此关注的是性能，不报告准确度或困惑度数值。根据模型的不同，更低的位宽可能需要量化感知训练、对具有混合位宽的模型进行量化，或调整组大小以达到可接受的准确度。

Llama 3.1 chart

尝试并贡献！

如果您想亲身体验新的低比特内核，可以通过设置 torchchat 并使用内核在本地量化并运行 LLM 来尝试。

如果您想贡献力量，可以考虑在以下领域添加支持