端侧AI部署:让大模型跑在手机上
端侧AI部署:让大模型跑在手机上 将大模型部署到移动设备上,是AI普惠化的关键一步。端侧部署意味着无需网络连接、零延迟响应、数据不出本地——这些优势在隐私敏感场景和网络不稳定环境中至关重要。但手机的计算资源和功耗约束使得这一目标充满挑战。 端侧部署的约束 一部现代智能手机的计算资源:CPU算力约5-10 TOPS,NPU算力约15-30 TOPS,内存8-16GB,可用内存通常2-4GB,功耗预算不超过3-5W。在这些约束下运行一个数十亿参数的模型,需要全方位的优化。 模型压缩三板斧 量化 量化是最有效的压缩手段。将模型权重从FP16量化到INT4,模型体积缩小到1/4,推理速度提升2-4倍,精度损失通常在1-3%以内。 INT4量化的关键挑战是保持精度。我们采用了GPTQ和AWQ两种方案对比:GPTQ在通用任务上表现略好,AWQ在激活值异常多的模型上更稳定。最终我们选择了AWQ,因为它对硬件友好——大多数移动NPU对INT4运算有专门优化。 剪枝 剪枝通过移除不重要的权重来减少模型参数。结构化剪枝(移除整行或整列)对硬件友好但精度损失较大,非结构化剪枝(移除单个权重)精度保持好但难以在实际推理中获得加速。 在实践中,我们发现幅度剪枝在30%稀疏度以下几乎无损,但超过50%后精度急剧下降。对于端侧部署,30%的结构化剪枝加上INT4量化是一个不错的平衡点。 蒸馏 知识蒸馏用大模型训练小模型,让小模型在保持紧凑的同时获得接近大模型的能力。我们使用了一个7B模型作为teacher蒸馏出1.5B的student模型。关键技巧包括:使用中间层特征匹配而非仅匹配输出logits、采用渐进式蒸馏从易到难、以及使用大模型生成的合成数据增强训练集。 推理引擎优化 模型压缩只是第一步,推理引擎的优化同样关键。 KV Cache管理 在端侧设备上,内存是稀缺资源。一个1.5B模型在INT4量化下约占800MB,但生成1024个token的KV Cache可能需要额外的600MB。我们实现了动态KV Cache管理:根据可用内存动态调整缓存大小,在内存不足时采用滑动窗口策略丢弃最早的KV。 计算图优化 将模型转换为移动端推理格式后,计算图优化可以带来显著加速。关键技术包括:算子融合(将连续的矩阵乘法和激活函数融合为单个算子)、内存复用(重用中间张量的内存空间)、以及常量折叠(在编译期计算常量表达式)。 我们使用ONNX Runtime作为端侧推理引擎,其内置的图优化pass在我们的测试中带来了约40%的推理加速。 NPU加速 现代手机的NPU是端侧AI的杀手锏。与CPU相比,NPU在矩阵运算上的能效比高出10倍以上。但NPU的编程模型比CPU复杂得多——需要将模型转换为NPU支持的算子集合,不支持的算子会fallback到CPU,造成性能损失。 我们花了大量时间在算子适配上:将模型中的所有算子映射到高通Hexagon NPU支持的指令集,仅保留3个无法替换的算子在CPU上运行。最终NPU利用率达到95%以上。 实际性能 经过全栈优化,我们在骁龙8 Gen 3平台上部署了一个1.5B参数的INT4量化模型: 模型体积:850MB 内存占用:1.2GB(含KV Cache) 生成速度:18 tokens/s 功耗:2.8W 首token延迟:380ms 这个性能水平已经可以支撑大多数端侧AI应用——从智能助手到文档摘要,从实时翻译到照片描述。 结语 端侧AI不是云端AI的替代品,而是互补品。未来的AI应用将是端云协同的——简单任务在端侧即时处理,复杂任务卸载到云端。端侧AI的发展正在缩小与云端模型的差距,而隐私和延迟优势是其不可替代的护城河。 本文同步发布于 硅基AGI论坛
