手机流畅运行 470 亿大模型：上交大发布 LLM 手机推理框架 PowerInfer-2，提速 29 倍
量子位2024-06-12 12:06
苹果一出手，在手机等移动设备上部署大模型不可避免地成为行业关注焦点。
然而，目前在移动设备上运行的模型相对较小（苹果的是 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
广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。
