手机流畅运行470亿大模型：上交大发布LLM手机推理框架PowerInfer-2，提速29倍

梦晨 2024-06-12 12:03:17 来源：量子位

梦晨 发自 凹非寺
量子位 | 公众号 QbitAI

苹果一出手，在手机等移动设备上部署大模型不可避免地成为行业关注焦点。

然而，目前在移动设备上运行的模型相对较小（苹果的是3B，谷歌的是2B），并且消耗大量内存，这在很大程度上限制了其应用场景。

即使是苹果，目前也需要与OpenAI合作，通过将云端GPT-4o大模型嵌入到操作系统中来提供能力更强的服务。

这样一来，苹果的混合方案引起了非常多关于数据隐私的讨论和争议，甚至马斯克都下场讨论。

如果苹果在操作系统层面集成OpenAI，那么苹果设备将被禁止在我的公司使用。这是不可接受的安全违规行为。

既然终端侧本地部署大模型的方案既让手机用户享受到AI强大的智能，又能保护好自己的隐私安全，为什么苹果还要冒着侵犯隐私的风险选择联手OpenAI采用云端大模型呢？主要挑战有两点：

• 手机内存不够大：按照大模型的Scaling Law法则，模型参数越大，能力对应的也就越强，这就意味着能力更强的模型对内存的要求越高。
• 手机算力不够强：即使勉强把通过量化等手段把模型塞进手机了，推理速度也慢，适合的应用场景也就非常有限了。

为了解决上述挑战，上海交大IPADS实验室推出了面向手机的大模型推理引擎（目前论文已在arxiv公开）：PowerInfer-2.0。

PowerInfer-2.0能够在内存有限的智能手机上实现快速推理，让Mixtral 47B模型在手机上达到11 tokens/s的速度。

与热门开源推理框架llama.cpp相比，PowerInfer-2.0的推理加速比平均达到25倍，最高达29倍。

为了充分释放出PowerInfer-2.0框架的最大潜力，上海交大团队还提出了配套的大模型优化技术Turbo Sparse，相关论文近期也上传了arxiv，并且已经在业内引起关注。

另外值得一提的是，去年底上海交大团队提出了针对PC场景的快速推理框架PowerInfer-1.0，在4090等消费级显卡的硬件上，实现了比llama.cpp高达11倍的推理加速，曾连续三天登顶GitHub趋势榜，5天获得了5k的GitHub star，目前已达到7.1k star。

相比PC，手机的内存和算力受到的约束更多，那么这次的PowerInfer-2.0是如何针对手机场景加速大模型推理呢？

动态神经元缓存

首先，针对手机运行内存（DRAM）不足的问题，PowerInfer-2.0利用了稀疏模型推理时的一个特点：每次只需要激活一小部分神经元，即“稀疏激活”。没有被激活的神经元即使不参与AI模型的推理计算，也不会对模型的输出质量造成影响。

稀疏激活为降低模型推理的内存使用创造了新的机会。为了充分利用稀疏激活的特性，PowerInfer-2.0把整个神经网络中的神经元分成了冷、热两种，并在内存中基于LRU策略维护了一个神经元缓存池。

近期频繁激活的”热神经元”被放置在运行内存中，而“冷神经元”只有在被预测激活的时候，才会被拉进内存，大幅降低了内存使用量。

其实冷热神经元分类，是继承自PowerInfer-1.0已有的做法。

而在去年12月，苹果在面向端侧的大语言模型推理方案“LLM in a Flash”中提出了和神经元缓存类似的“滑动窗口”技术。但这些工作主要针对的都是PC环境，直接迁移到手机环境，还会遇到新的难题。

首先手机平台的硬件条件远不及PC，无论是算力、内存总量还是存储带宽，都与PC存在较大差距。

其次，手机硬件平台存在CPU、GPU、NPU三种异构的计算单元，十分复杂。各大硬件平台宣发时都会强调一个总算力，实际上是把CPU、GPU、NPU提供的算力加起来。然而真正跑起大模型来，能不能高效利用各种异构算力还是个问题。

以神经元簇为粒度的异构计算

针对这一点，PowerInfer-2.0进一步把粗粒度的大矩阵计算分解成细粒度的“神经元簇”。

每个神经元簇可以包含若干个参与计算的神经元。对于不同的处理器，会根据处理器的特性来动态决定划分出来的神经元簇的大小。

例如，NPU擅长于做大矩阵的计算，那么可以把所有神经元合并成一个大的神经元簇，一起交给NPU计算，这样就可以充分利用NPU的计算能力。而在使用CPU时，可以拆出多个细粒度的神经元簇，分发给多个CPU核心一起计算。

具体而言，PowerInfer-2.0为模型推理的预填充阶段（Prefill）和解码阶段（Decoding）分别设计了两套神经元簇的划分方案：

预填充阶段会一次性输入很多token，基本上绝大部分神经元都会被激活，因此选择使用大神经元簇交给NPU计算。CPU此时也没有闲着，在后台为NPU执行反量化模型权重的操作。

解码阶段每次只有一个token，具有较高的稀疏性，因此更加适合划分成若干细粒度的神经元簇，交给CPU灵活调度和执行计算。

神经元簇这一概念除了能够更好的适应手机的异构计算环境，还能天然地支持计算与存储I/O的流水线并行执行。

PowerInfer-2.0提出了分段神经元缓存和神经元簇级的流水线技术，在一个神经元簇等待I/O的同时，可以及时地把另一个已经准备好的神经元簇调度到处理器上进行计算，从而充分隐藏了I/O的延迟。

同时，这种基于神经元簇的流水线打破了传统推理引擎中逐矩阵计算的方式，可以允许来自不同参数矩阵的神经元簇交错执行，达到最高的并行效率。

I/O加载神经元的速度对于模型推理也至关重要。

分段缓存会针对不同的权重类型采取不同策略（如注意力权重、预测器权重、前馈网络权重）采取不同的缓存策略，提高缓存命中率，减少不必要的磁盘 I/O。

缓存还会使用LRU替换算法动态更新每个神经元的实际冷热情况，确保缓存中放着的都是最热的神经元。此外PowerInfer-2.0还针对手机UFS 4.0存储的性能特点，设计了专门的模型存储格式，提高读取性能。

最后再来看一下实测成绩，使用一加12和一加Ace 2两款测试手机，在内存受限的情况下，PowerInfer-2.0的预填充速度都显著高于llama.cpp与LLM in a Flash（简称“LLMFlash”)：

解码阶段同样是PowerInfer-2.0占据很大优势。特别是对于Mixtral 47B这样的大模型，也能在手机上跑出11.68 tokens/s的速度：

而对于Mistral 7B这种可以放进手机运行内存的模型，PowerInfer-2.0可以节约40%内存的情况下，达到与llama.cpp和MLC-LLM同水平甚至更快的解码速度：

PowerInfer-2.0是一个模型-系统协同设计的方案，也就是需要模型中可预测稀疏性的配合。

如何以低成本的形式调整模型以适配PowerInfer-2.0框架，也是一个重大挑战。

低成本高质量地大幅提升模型稀疏性

传统简单的ReLU稀疏化会给模型原本的能力造成不小的影响。

为了克服这个问题，上海交大IPADS联合清华和上海人工智能实验室提出一个低成本地稀疏化方法，不仅大幅提升模型的稀疏性，还能保持住模型原本的能力！

首先，论文深入分析了模型稀疏化中的问题：

• 在类LLaMA模型中中简单引入ReLU，虽然能引入一定程度的稀疏性，但稀疏度仍然有限。
• 稀疏化过程由于训练语料的不足和训练token的不足导致模型精度下降的问题。

为了提升模型的稀疏度，论文在ReLU基础上提出dReLU激活函数，采用替换原有激活函数后继续预训练的方式增加模型稀疏性。

将SwiGLU替换为dReLU一方面直观地提高了输出值中的零元素比例，另一方面能更有效地在稀疏化的过程中复用原本模型训练完成的gate和up矩阵权重。

为了克服模型能力下降的问题，团队收集了包括网页、代码和数学数据集在内的多样化继续训练语料库。高质量、多样化的训练数据有助于模型在稀疏化后更好地保持和提升性能。

最后，团队训练了2个TurboSparse大模型进行验证，分别是8x7B和7B的大模型。得益于高质量的继续训练语料，TurboSparse系列模型模型的精度甚至还能反超原版模型（具体见表6）。

而在稀疏度方面效果也非常显著。相比于原本的Mixtral模型需要激活13B参数量，TurboSparse-Mixtral只需要激活4.3B的参数量，激活的参数量是原本模型的三分之一。

而关于稀疏化过程的成本问题，TurboSparse论文中介绍，改造过程中模型需要继续训练150B tokens，相比于预训练（假设3T tokens）还不到5%，说明其成本是很低的。

让技术加速走出实验室

从推理框架和改造模型两个角度出发，上海交大团队的成果实现了大语言模型在手机等资源受限场景下的快速推理。

而且这套方案的潜力不止于手机，未来在车载设备、智能家居等方向还有更多应用前景。

最后再正式介绍一下团队。上海交通大学并行与分布式系统研究所（简称IPADS），由陈海波教授领导，现有13名教师，100多名学生。

IPADS长期从事计算机系统的研究，近10年在权威榜单CSRankings的Operating Systems领域排名全球前二，仅次于MIT；上海交大也是排名前十中唯一上榜的亚洲高校。

目前，上海交大IPADS已经在Huggingface上开放了稀疏化的模型权重。在未来，如果PowerInfer-2.0能够与手机厂商进一步紧密合作，相信可以加速相关技术走出实验室，落地到各种真实场景。

PowerInfer-2论文：https://arxiv.org/abs/2406.06282
TurboSparse论文：https://arxiv.org/abs/2406.05955
模型权重：https://huggingface.co/PowerInfer/TurboSparse-Mixtral