写点什么

大模型端侧 CPU 部署最高提效 6 倍!微软亚研院新开源项目 T-MAC 技术解析来了

  • 2024-08-12
    北京
  • 本文字数:2670 字

    阅读完需:约 9 分钟

大模型端侧 CPU 部署最高提效 6 倍!微软亚研院新开源项目 T-MAC 技术解析来了

为增强设备上的智能性,在边缘设备部署大型语言模型(LLMs)成为了一个趋势,比如微软的 Windows 11 AI + PC。目前部署的大语言模型多会量化到低比特。然而,低比特 LLMs 在推理过程中需要进行低精度权重和高精度激活向量的混合精度矩阵乘法(mpGEMM)。现有的系统由于硬件缺乏对 mpGEMM 的原生支持,不得不将权重反量化以进行高精度计算。这种间接的方式导致了显著的推理开销,并且无法随着比特数进一步降低而获得加速。


为此,微软亚洲研究院的研究员们开发了 T-MAC。T-MAC 采用基于查找表(LUT)的计算范式,无需反量化,直接支持混合精度矩阵乘,其高效的推理性能以及其统一且可扩展的特性为在资源受限的边缘设备上实际部署低比特 LLMs 铺平了道路。


此外,当前大模型的部署普遍依赖于专用加速器,如 NPU 和 GPU 等,而 T-MAC 可以摆脱专用加速器的依赖,仅利用 CPU 部署 LLMs,推理速度甚至能够超过同一片上的专用加速器,使 LLMs 可以部署在各类包括 PC、手机、树莓派等边缘端设备。T-MAC 现已开源。

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


T-MAC 的关键创新在于采用基于查找表(LUT)的计算范式,而非传统的乘累加(MAC)计算范式。T-MAC 利用查找表直接支持低比特计算,从而消除了其他系统中必须的反量化 (dequantization) 操作,并且显著减少了乘法和加法操作的数量。


经过实验,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 倍。即使在较低端的设备如 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 的性能。



图 1:BitNet on T-MAC vs llama.cpp on Apple M2



图 2:在不同端侧设备 CPU(Surface Laptop 7, NVIDIA AGX Orin, Apple M2-Ultra)的各核数下 T-MAC 和 llama.cpp 的 token 生成速度可达 llama.cpp 的 4-5 倍。达到相同的生成速率,T-MAC 所需的核心数仅为原始 llama.cpp 的 1/4 至 1/6。

矩阵乘不需乘,只需查表 (LUT)


对于低比特参数 (weights),T-MAC 将每一个比特单独进行分组(例如,一组 4 个比特),这些比特与激活向量相乘,预先计算所有可能的部分和,然后使用 LUT 进行存储。之后,T-MAC 采用移位和累加操作来支持从 1 到 4 的可扩展位数。通过这种方法,T-MAC 抛弃了 CPU 上效率不高的 FMA (乘加)指令,转而使用功耗更低效率也更高的 TBL/PSHUF(查表)指令。



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

以比特为核心的计算,取代以数据类型为核心的计算


传统的基于反量化的计算,实际上是以数据类型为核心的计算,这种方式需要对每一种不同的数据类型单独定制。每种激活和权重的位宽组合,如 W4A16(权重 int4 激活 float16) 和 W2A8,都需要特定的权重布局和计算内核。例如,W3 的布局需要将 2 位和另外 1 位分开打包,并利用不同的交错或混洗方法进行内存对齐或快速解码。然后,相应的计算内核需要将这种特定布局解包到硬件支持的数据类型进行执行。


而 T-MAC 通过从比特的视角观察低比特矩阵乘计算,只需为单独的一个比特设计最优的数据结构,然后通过堆叠的方式扩展到更高的 2/3/4 比特。同时,对于不同精度的激活向量(float16/float32/int8),仅有构建表的过程需要发生变化,在查表的时候不再需要考虑不同的数据结构。



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


同时,传统基于反量化的方法,从 4- 比特降低到 3/2/1- 比特时,尽管内存占用更少,但是计算量并未减小,而且由于反量化的开销不减反增,性能反而可能会更差。但 T-MAC 的计算量随着比特数降低能够线性减少,从而在更低比特带来更好加速,为最新的工作 BitNet, EfficientQAT 等发布的 1- 比特 /2- 比特模型提供了高效率的部署方案。



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

高度优化的算子实现


基于比特为核心的计算具有许多优势,但将其实现在 CPU 上仍具有不小的挑战:(i) 与激活和权重的连续数据访问相比,表的访问是随机的。表在快速片上内存中的驻留对于最终的推理性能尤为重要,(ii) 然而,片上内存是有限的,查找表(LUT)方法相比传统的 mpGEMV 增大了片上内存的使用。这是因为查找表需要保存激活向量与所有可能的位模式相乘的结果。这比激活本身要多得多。



图 6:T-MAC 与 llama.cpp 在计算数据流上的不同


为此,微软亚洲研究院的研究员们深入探究了基于查表的计算数据流,为这种计算范式设计了高效的数据结构和计算流程,其中包括:


  1. 将 LUT 存入片上内存,以利用 CPU 上的查表向量指令 (TBL/PSHUF) 提升随机访存性能。

  2. 改变矩阵 axis 计算顺序,以尽可能提升放入片上内存的有限 LUT 的数据重用率。

  3. 为查表单独设计最优矩阵分块 (Tiling) 方式,结合 autotvm 搜索最优分块参数

  4. 参数 weights 的布局优化

  5. weights 重排,以尽可能连续访问并提升缓存命中率

  6. weights 交错,以提升解码效率

  7. 对 Intel/ARM CPU 做针对性优化,包括

  8. 寄存器重排以快速建立查找表

  9. 通过取平均数指令做快速 8- 比特累加


研究员们在一个基础实现上,一步步应用各种优化,最终相对于 SOTA 低比特算子获得显著加速:



图 7:在实现各种优化后,T-MAC 4- 比特算子最终相对于 llama.cpp 获得显著加速

开源易用的工具


T-MAC 现已开源 https://github.com/microsoft/T-MAC,简单输入几条命令即可在自己的笔记本电脑上高效运行 Llama-3-8B-instruct 模型。


代码:https://github.com/microsoft/T-MAC

论文:https://www.arxiv.org/pdf/2407.00088

2024-08-12 10:339846
用户头像
李冬梅 加V:busulishang4668

发布了 1111 篇内容, 共 722.8 次阅读, 收获喜欢 1253 次。

关注

评论

发布
暂无评论
发现更多内容

软件测试 | 测试开发 | 接口测试之HTTP、HTTPS 抓包分析

测吧(北京)科技有限公司

HTTP

呆猫云工作站助力Omniverse云上部署试水元宇宙

神奇视野

教育信息化迈入2.0时代,呆猫云工作站破除技术壁垒

神奇视野

看完这篇你将get VR/AR沉浸式技术的“创作密码”,速来!

神奇视野

技术分享 | 软件项目管理与跨部门沟通协作

霍格沃兹测试开发学社

软件测试 | 测试开发 | 接口自动化测试中如何对xml 格式做断言验证?

测吧(北京)科技有限公司

接口自动化测试

【Django | 开发】面试招聘网站(增加csv,excel导出&企业域账号集成&日志管理功能)

计算机魔术师

8月月更

【Django | 开发】面试招聘信息网站(用户登录注册&投在线递简历)

计算机魔术师

8月月更

【Django | allauth】重写allauth重置密码方法

计算机魔术师

8月月更

汽车制造企业如何最大化数据资产价值?

Kyligence

数据分析 智能多维数据库

NFT系统:数字藏品交易平台app开发功能

开源直播系统源码

NFT 数字藏品 数字藏品软件开发 数字藏品开发 数字藏品系统

leetcode 28. Implement strStr() 实现 strStr()(简单)

okokabcd

LeetCode 数据结构与算法

“云”上交通,“翼”路畅行

天翼云开发者社区

加快云网融合发展,打通算力传输大动脉!

天翼云开发者社区

【小程序项目开发-- 京东商城】uni-app之自定义搜索组件(下) -- 搜索历史

计算机魔术师

8月月更

内网穿透是什么意思?有什么用?用什么软件好?

行云管家

运维 内网穿透 内网

主机监控是什么意思?用什么软件好?咨询电话多少?

行云管家

运维 主机 主机监控

恭喜天翼云“翼起飞”战队在CCF国际AIOps挑战赛中夺得亚军!

天翼云开发者社区

多线程原理和常用方法以及Thread和Runnable的区别

共饮一杯无

多线程 8月月更

【Django | 开发】面试招聘信息网站(快速搭建核心需求)

计算机魔术师

8月月更

华为云联合信通院发布业界首个《云原生数据库白皮书》

IT资讯搬运工

【Django | 开发】面试招聘信息网站(划分面试官权限&集成钉钉消息)

计算机魔术师

8月月更

云聚华为伙伴暨开发者大会GaussDB专场,与客户伙伴共话金融科技新发

IT资讯搬运工

华为云GaussDB践行云原生 带来极致弹性、全域可用和极简体验

IT资讯搬运工

内容小程序化,是在线教育服务推广的最佳格式

Speedoooo

小程序 在线教育 移动开发 小程序容器

【Django | allauth】useprofile 用户模型扩展

计算机魔术师

8月月更

【Django | 开发】面试招聘信息网站(处理产品细节和权限&美化页面样式)

计算机魔术师

8月月更

helm实战之开发Chart

程序员欣宸

Kubernetes Helm 8月月更

OpenHarmony技术挑战课题征集

OpenHarmony开发者

OpenHarmony

软件测试 | 测试开发 | 接口测试中如何使用Json 来进行数据交互 ?

测吧(北京)科技有限公司

json

华为云GaussDB深耕数据库根技术,助力能源行业数字化转型

IT资讯搬运工

大模型端侧 CPU 部署最高提效 6 倍!微软亚研院新开源项目 T-MAC 技术解析来了_微软_微软亚洲研究院_InfoQ精选文章