手机跑大模型提速 4-5 倍：微软亚研院开源新技术 T-MAC，有 CPU 就行

量子位
2024-08-09 12:58

有 CPU 就能跑大模型，性能甚至超过 NPU / GPU！

没错，为了优化模型端侧部署，微软亚洲研究院提出了一种新技术 —— T-MAC。

这项技术主打性价比，不仅能让端侧模型跑得更快，而且资源消耗量更少。

咋做到的？？

在 CPU 上高效部署低比特大语言模型

一般来说，要想在手机、PC、树莓派等端侧设备上使用大语言模型，我们需要解决存储和计算问题。

常见的方法是模型量化，即将模型的参数量化到较低的比特数，比如 4 比特、3 比特甚至更低，这样模型所需的存储空间和计算资源就会减少。

不过这也意味着，在执行推理时，需要进行混合精度的矩阵乘法运算（mpGEMM），即用低精度的权重和高精度的激活向量进行计算。

然而，现有的系统和硬件并不原生支持这种混合精度的矩阵乘法，因此它们通常需要将低精度的权重转换回高精度，这个过程叫做反量化 (dequantization)。

但这种方法不仅效率低，而且当比特数进一步降低时，并不能带来性能上的提升。

对此，新技术 T-MAC 采用基于查找表（LUT）的计算范式，无需反量化，直接支持混合精度矩阵乘。

这样，T-MAC 不仅提高了推理性能，还使得模型更加统一和可扩展，尤其适合在资源受限的端侧设备部署。

此外，T-MAC 不依赖于专用的硬件加速器 NPU 或 GPU，能够仅利用 CPU 部署模型。甚至在某些情况下，它的推理速度可以超过专用加速器。

T-MAC 的关键创新在于采用基于查找表（LUT）的计算范式，而非传统的乘累加（MAC）计算范式。

T-MAC 利用查找表直接支持低比特计算，从而消除了其他系统中必须的反量化操作，并且显著减少了乘法和加法操作的数量。

经过实验，T-MAC 展现出了卓越的性能：

在配备了最新高通 Snapdragon X Elite 芯片组的 Surface AI PC 上，3B BitNet-b1.58 模型的生成速率可达每秒 48 个 token，2bit 7B llama 模型的生成速率可达每秒 30 个 token，4bit 7B llama 模型的生成速率可达每秒 20 个 token。

这甚至超越了 NPU 的性能！

当部署 llama-2-7B-4bit 模型时，尽管使用 NPU 可以生成每秒 10.4 个 token，但 CPU 在 T-MAC 的助力下，仅使用两核便能达到每秒 12.6 个 token，最高甚至可以飙升至每秒 22 个 token。

这些都远超人类的平均阅读速度，相比于原始的 llama.cpp 框架提升了 4~5 倍。

▲ BitNet on T-MAC（基于 LUT）vs llama.cpp（基于反量化）

即使在较低端的设备如 Raspberry Pi 5 上，T-MAC 针对 3B BitNet-b1.58 也能达到每秒 11 个 token 的生成速率。

同时，T-MAC 也具有显著的功耗优势：

达到相同的生成速率，T-MAC 所需的核心数仅为原始 llama.cpp 的 1/4 至 1/6，降低能耗的同时也为其它应用留下计算资源。

值得注意的是，T-MAC 的计算性能会随着比特数的降低而线性提高，这一现象在基于反量化去实现的 GPU 和 NPU 中是难以观察到的。

这进一步使得 T-MAC 能够在 2 比特下实现单核每秒 10 个 token，四核每秒 28 个 token，大大超越了 NPU 的性能。

采用新的计算范式

好了，说完了效果，咱们接着展开 T-MAC 的技术细节。

1、矩阵乘不需乘，只需查表（LUT）

对于低比特参数 (weights)，T-MAC 将每一个比特单独进行分组（例如，一组 4 个比特），这些比特与激活向量相乘，预先计算所有可能的部分和，然后使用 LUT 进行存储。

之后，T-MAC 采用移位和累加操作来支持从 1 到 4 的可扩展位数。

通过这种方法，T-MAC 抛弃了 CPU 上效率不高的 FMA（乘加）指令，转而使用功耗更低、效率也更高的 TBL / PSHUF（查表）指令。

▲ 混合精度 GEMV 基于现有反量化的实现范式 vs T-MAC 基于查找表的新范式

2、以比特为核心的计算，取代以数据类型为核心的计算

传统的基于反量化的计算，实际上是以数据类型为核心的计算，这种方式需要对每一种不同的数据类型单独定制。

每种激活和权重的位宽组合，如 W4A16（权重 int4 激活 float16）和 W2A8，都需要特定的权重布局和计算内核。

例如，W3 的布局需要将 2 位和另外 1 位分开打包，并利用不同的交错或混洗方法进行内存对齐或快速解码。

然后，相应的计算内核需要将这种特定布局解包到硬件支持的数据类型进行执行。

而 T-MAC 通过从比特的视角观察低比特矩阵乘计算，只需为单独的一个比特设计最优的数据结构，然后通过堆叠的方式扩展到更高的 2/3/4 比特。

同时，对于不同精度的激活向量（float16 / float32 / int8），仅有构建表的过程需要发生变化，在查表的时候不再需要考虑不同的数据结构。

▲ 以比特为核心的查表计算混合精度 GEMV

同时，传统基于反量化的方法，从 4-比特降低到 3/2/1-比特时，尽管内存占用更少，但是计算量并未减小，而且由于反量化的开销不减反增，性能反而可能会更差。

但 T-MAC 的计算量随着比特数降低能够线性减少，从而在更低比特带来更好加速，为最新的工作 BitNet, EfficientQAT 等发布的 2-比特模型提供了高效率的部署方案。

比如下图展示了：

（1）使用不同端侧设备 CPU 的单核，T-MAC 在 4 到 1 比特的混合精度 GEMV 算子相较 llama.cpp加速 3-11 倍。
（2）T-MAC 的 GEMM 耗时能随着比特数减少线性减少，而基于反量化的 llama.cpp 无法做到（1 比特 llama.cpp 的算子性能由其 2 比特实现推算得到）。