手机运行大型模型速度提升4-5倍:微软亚洲研究院开源新技术T-MAC,只需CPU即可。
编辑日期:2024年08月10日
确实,为了优化模型的端侧部署,微软亚洲研究院提出了一项新技术——T-MAC。
这项技术主要强调性价比,不仅能加速端侧模型的运行速度,还能减少资源消耗。
是如何做到的呢?
通常来说,若要在如手机、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,但在T-MAC的支持下,CPU仅利用两核心就能达到每秒12.6个token的处理速度,最高甚至可提升至每秒22个token。
这些都远超过了人类的平均阅读速度,相比原始的llama.cpp框架,速度提高了4到5倍。
即使在像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个令牌,四核每秒28个令牌,从而大幅度超越NPU的性能。
好的,讲完了效果之后,我们继续来探讨T-MAC的技术细节。
对于低比特权重(parameters/weights),T-MAC将每个比特独立分组(例如,4个比特为一组),这些比特与激活向量相乘,并预先计算出所有可能的部分和,然后利用查找表(LUT)进行存储。
然后,T-MAC 使用移位和累加操作来支持位数的可扩展性,范围可以从 1 到 4。
通过这种方式,T-MAC放弃了CPU上效率较低的FMA(乘加)指令,转而采用功耗更低、效率更高的TBL/PSHUF(查表)指令。
传统的基于反量化计算的方法,实际上是围绕数据类型为核心进行计算的,这种方式要求对每种不同的数据类型进行单独定制。
每种激活和权重的位宽组合,例如 W4A16(权重 int4 激活 float16)和 W2A8,都需要特定的权重布局和计算内核。
例如,W3的布局需要将2位和另外1位分开进行打包,并采用不同的交织或混洗方法以实现内存对齐或快速解码。
然后,相应的计算内核需要将这种特定的布局解包到硬件支持的数据类型以进行执行。
而 T-MAC 通过从比特的角度来审视低比特矩阵乘法计算,只需要为单一比特设计最优的数据结构,随后可通过堆叠方式扩展至更高的2/3/4比特。
同时,对于不同精度的激活向量(如 float16 / float32 / int8),仅在构建表的过程中需要做出调整,在进行查表操作时则无需考虑不同的数据结构。
同时,传统的基于反量化的方案在从4位减少到3/2/1位时,虽然占用的内存减少了,但计算量并未减少。由于反量化带来的开销不降反增,性能可能反而会变得更差。
但是,T-MAC的计算量能随着位数的减少而线性降低,因此在低位数时可以实现更好的加速效果。这一特性为最近推出的如BitNet、EfficientQAT等工作中的2位模型提供了高效的部署方案。
如下图所示:
总的来说,以比特为基础的计算方式有许多优点,但在CPU上的实现仍然存在不少挑战。
与对激活和权重的连续数据访问相比,对表格的访问是随机的。
表驻留在快速芯片的内存中对最终的推理性能尤其重要。然而,芯片上的内存是有限的。与传统的mpGEMV方法相比,查表(LUT)方法增加了芯片上内存的使用量。
这是因为查找表需要保存激活向量与所有可能的位模式相乘的结果,这比激活本身的数量要多得多。
为了实现这一目标,微软亚洲研究院的研究人员对基于查表的计算数据流进行了深入研究,并为这种计算模式设计了高效的数据结构和计算流程。这包括:
1、将 LUT 存储在片上内存中,以便利用 CPU 上的查表向量指令 (TBL/PSHUF) 来提高随机访问内存的性能。
2、更改矩阵的轴计算顺序,以尽可能提高放入片上内存的有限查找表(LUT)的数据重用率。
3、为查表单独设计最优化的矩阵分块(Tiling)方法,并结合autotvm搜索最佳的分块参数。
4、对参数 weights 的布局进行优化:
a、对权重进行重新排列,以便尽可能实现连续访问并提高缓存命中率。
b、权重交错,以提高解码效率。
5、针对 Intel / ARM CPU 进行专门的优化,包括
a、通过寄存器重排以快速构建查找表
b、通过取平均数指令进行快速的8位累加
研究人员在基本实现的基础上逐步应用各种优化措施,最终相较于最先进的低比特运算器实现了显著的速度提升。
例如,在实施了各种优化之后,T-MAC 4-比特运算符最终相比 llama.cpp 实现了显著的速度提升:
最后,T-MAC 已经开源,相关的论文也已在 arXiv 上公开,有兴趣的话可以进一步了解。
开源地址(包含代码):
https://github.com/microsoft/T-MAC
(这段文本是一个网址,不需要重写中文,但是如果需要描述它的话,可以这样写:)
这是微软公司在GitHub上发布的T-MAC项目的链接。
论文:
(注:您提供的文本内容不全,如果需要具体重写的内容,请提供完整的句子或段落。)
您的链接是指向一篇在arXiv上的论文。但由于这是一个未来的日期(截至2023年),我们无法确认该链接的内容。对于链接的描述您没有提供具体的信息,根据链接格式,我猜测这可能是对一篇学术论文的引用。如果是需要转换成中文形式,并非是内容的翻译,可以这样写:
"https://www.arxiv.org/pdf/2407.00088"
如果有关于这篇论文的具体信息或者内容需要翻译,请提供更多的上下文。
本文来源于微信公众号:量子位(ID:QbitAI),作者为微软亚洲研究院。原标题为《手机运行大型模型速度提升4-5倍!微软亚洲研究院开源新技术,只需CPU即可》。