MNN在大模型端侧部署上的探索

为了在更广泛的设备上部署大模型，MNN团队开发了 MNN-LLM / MNN-Diffusion，合称MNN-Transformer ，支持大语言模型和文生图等AIGC模型，具有如下特性：

• 支持各类LLM和Diffusion模型，支持加载同时加载多份Lora；不依赖厂商NPU能力，2020年后的手机基本都能跑得动 LLM 小模型。

• 支持int4/int8等模型量化方案，并支持在内存不足时使用磁盘空间替换，避免内存溢出风险。

• 充分利用CPU sdot / smmla 与GPU recordable / simdgroup / GMemory 等较新特性，在8Gen1芯片上，MNN-Transformer支持 1.8b 端侧模型 35 token/s 以上的解码速度，生成 512x512的图片约 40s (2s / iter)，基本上能充分利用移动端上CPU与GPU算力。

MNN 及大模型推理相关代码均已开源：

https://github.com/alibaba/MNN/

功能介绍

如上述架构图所示，MNN-Transformer由导出工具、量化工具、插件与引擎三个部分组成：导出工具负责将各类大模型转换为MNN格式，并构建必需的资源包；量化工具减少MNN模型大小，以降低运行时内存，加快执行速度；LLM/Diffusion运行时所需要的分词、KV缓存管理、LoRA等功能由插件与引擎模块提供。



使用流程图如下所示：

▐  离线工具

这部分功能在PC/服务器上使用，产出资源包给端侧运行。

• 导出

之前的深度学习模型一般由 Pytorch / Tensorflow 训练后，导出 ONNX / PB 模型，再通过 MNN 模型转换工具转成 MNN 模型。

各类开源大模型都基于Pytorch实现，并不直接提供转换到ONNX的功能。并且大模型体积大而结构变化相对较少，采用原有流程，一方面需要用户额外去修改源码，支持ONNX导出，另外一方向在导出ONNX时会占用大量内存，并且速度很慢。整体用户体验不佳。因此我们针对各模型单独写导出脚本，开发成本可以接受，而模型转换的速度可以大幅提升。

• 量化

量化，一般指低bit量化，使用较低bit（低于32bit）的格式存储模型参数，以降低模型大小从而减少推理时的内存开销并提升推理速度。针对低bit量化目前主要分为2种方法：训练后量化（PTQ）和量化感知训练（QAT）。其中训练后量化（PTQ）又分为无输入校验的量化和有输入校验的量化。无输入校验的量化直接对参数分布进行统计，使用对称或非对称量化将参数按照channel或者block线性映射到低bit表示范围；这种方案使用简单，可以直接将浮点模型转换为低bit模型。有输入校验的方案在将参数映射到低bit表示范围时，需要提供一些输入作为校准数据，如GPTQ、SmoothQuant等方案，这些方案需要一些额外步骤完成量化。量化感知训练（QAT）相比训练后量化（PTQ）的精度损失更低，但是需要在模型训练阶段进行额外工作，或针对预训练模型进行训练，这种方案的成本较高。激活部分是指计算时得到的中间结果（hidden state)或attention计算中的key、value的部分。权重部分的量化可以在运行时统计该数据的分布，按照无数据校验的权重量化方法进行量化。

MNN 支持4bit/8bit无输入校验量化，支持对称量化，非对称量化；同时支持channel-wise量化和block-wise量化；同时支持GPTQ量化权重导入计算。

▐  插件与引擎

• Attention ：大模型一般基于Transformer结构，Attention 插件提供Transformer结构的 Cross-Attention / Self-Attention 等算子实现，降低内存访问量，提高GPU并发量，实现对应性能加速。

• KV Manager ：KV Cache 是大语言模型的重要优化手段，也是大语言模型中长句子内存增量的主要来源。KV Manager 插件管理大语言模型的 KV Cache ，提供申请、扩容、量化、预加载等功能。

• LoRA ：LoRA 是为了让大模型支持多个任务，比如用同一个LLM实现多个人设。基于MNN权重共享的能力，LoRA插件可以用同一个基座大模型，以很少的内存占用（仅LoRA模型大小）和轻微的性能损失（10%以内）支持加载多个LoRA模型。

• Tokenizer ：用于将文本与整形数组（tokens）互转，目前实现了 Sentencepiece 和 Tiktoken

• Embedding ：用于将整形数组转换为向量，以便输入深度学习模型推理

• Sampler ：对深度学习模型输出的向量，主要是LLM的输出进行后处理，输出tokens

• Engine ：集成插件与MNN模型推理能力，实现大语言模型和文生图功能

内存优化

▐  动态量化

对大模型来说，低Bit量化是在端侧部署的必要条件，比如一个70亿参数（7b）的模型，需要约25GB的空间。如果采用4位量化，7b模型就只需大约3.5GB内存，1.8b模型不到1GB。一般而言，需要权重量化且预先计算模型中各层的量化信息，才能基于低bit直接计算。但由于大模型参数量大，使训练量化和基于KL散度的后训练量化成本都比较高，难以预先计算模型中各层的量化信息。因此，我们新增了动态量化机制，以便直接运行仅权重量化的模型。

对 CPU ，由于有 int8 计算指令，我们对输入进行一次数学统计，计算出量化参数(scale/bias)并进行量化，使之能与int4/int8权重的进行计算。相比于浮点计算，可以减少内存占用、减少内存访问量，提升计算效率。我们采用Per-Batch的方案，可以在性能轻微损失的情况下，保证精度几乎无损。

对 GPU ，由于 int8 计算指令效率不高，我们把反量化操作（int4/int8转浮点）合入计算卷积/矩阵乘的Kernel，以支持直接读取 int4/int8权重。相比于浮点计算，可以减少内存占用和内存访问量，但计算效率下降，因此，对于计算密集的情况（LLM的预填充阶段和diffusion中的卷积计算），我们仍然需要增加一个把 int4/int8 的权重反量化为浮点的过程，再使用原先浮点的卷积/矩阵乘Kernel进行计算。

▐  磁盘映射技术

端侧运行大模型时，常常会遇到OOM（内存溢出）的问题。但其实单独跑大模型时内存一般是足够的，只是与应用中其他模块的内存叠加后导致了溢出。为此，MNN支持了磁盘映射技术，通过mmap接口，可将模型运行内存映射到磁盘，这样在其他模块内存不足时，可以卸载模型运行内存。而由于大模型与其他模块一般并不会同时执行，所以对模型运行效率的影响较小。

▐  KV量化与缓存

在 LLM 输出句子变长后，KV Cache 的内存成为不可忽略的增量，为减少这部分内存我们实现了 KV 量化。对 CPU 后端而言，K 的增长方向 kv_seq_len 与int8矩阵乘计算的 scale 数组增加方向一致，可以使用 int8 量化及相应int8矩阵乘计算。而 V 的增长方向与 int8矩阵乘计算的 scale 数组增长方向不一致，目前采用FP8的量化方案。

性能测试

端侧大模型的性能优化总体上和深度学习推理引擎的性能优化一致，由于近年来新特性的逐步普及，MNN 针对这些特性做了适配优化，如 ARM CPU 的 sdot / smmla 指令，高通OpenCL提供的recordable queue / G-Memory 等，对于Transformer结构的模型，MNN实现了Attention插件及相应的图优化。

▐  MNN-LLM

LLM 的推理过程可分为两个阶段：

• Prefill (预填充阶段)：LLM 计算并存储初始输入 token 的 KV 缓存，并生成第一个输出 token。

• Decode：LLM 使用 KV 缓存逐个生成输出 token，然后使用新生成 token 的 (K) - (V) 对，进行更新。



Prefill 性能决定响应时间，Decode 则一般决定总回答时长，都是 LLM 推理性能衡量的重要指标。

1. 测试文本

计算8乘以12将下面的句子翻译成中文：It's a beautiful day to learn something new.描述优秀的领导者应具备的五个特质，并解释每个特质为什么重要近年来，随着技术的快速发展和全球化的深入推进，数字经济已成为推动世界经济增长的新引擎。数字经济不仅改变了人们的生活方式，促进了信息和资源的快速流通，还重塑了传统行业的业务模式和竞争格局。尽管数字经济的发展为全球经济增长提供了新的动能，但同时也带来了数据安全、隐私保护、数字鸿沟和市场垄断等一系列挑战。考虑到这些背景，请详细分析数字经济在促进世界经济增长方面的作用，包括但不限于数字经济对提高生产效率、创造就业机会和促进可持续发展的贡献。同时，探讨如何应对数字经济发展过程中出现的挑战，具体包括如何保护个人数据安全和隐私、缩小数字鸿沟以确保数字经济的包容性和公平性，以及如何制定有效政策以避免市场垄断情况的出现，最终实现数字经济的健康和可持续发展。%

2. 测试方法

MNN-LLM的编译与运行参考文档：

https://mnn-docs.readthedocs.io/en/latest/transformers/llm.html

• PC：编译出 llm_demo ，使用 llm_demo 测试

• iOS：在打开的工程中输入 benchmark

• Android：编译出 llm_demo ，将相关资源 push 到手机上，测试命令同 PC

3. 其他端侧llm方案对比

当前端侧LLM开源推理方案主要有：

• LLama.cpp：https://github.com/ggerganov/llama.cpp

  测试版本：06943a69f678fb32829ff06d9c18367b17d4b361

• MLC-LLM：https://github.com/mlc-ai/mlc-llm

  测试版本：9d798acca1e0c9ee37adf5d0980cff9d5566b2f4

• FastLLM：https://github.com/ztxz16/fastllm

  测试版本：3113a18be60f959925e87f36f364504ec99725a0

由于各方案的量化算法有所差别，输出结果不完全一样，且 MLC-LLM 和 FastLLM 在实测过程中不稳定（容易卡死），我们限制了最大输出字数为16，并分别用 64 / 256 / 1024 的输入文本进行测试。

• 测试设备：Android Mi 14

• 测试模型：Qwen2-1.5B ，Qwen2-7B，LLama3-8B



测试结果：MNN CPU Decode 有20-50%优势，尤其是在 Prefill 阶段快于其他方案1倍以上；GPU 性能在小模型上快于其他方案30%以上，较大模型上与MLC-LLM持平。但相比MLC-LLM，MNN-LLM的GPU输出更稳定（不容易crash）。

千问 1.5B - 1024 输入下对比图：

详细数据如下：

▐  MNN-Diffusion

端侧 Diffusion 模型的运行方案目前很少，仅有 stable-diffusion.cpp 和基于 onnxruntime 的实现。

stable-diffusion.cpp : https://github.com/leejet/stable-diffusion.cpp

从性能数据看，MNN 的方案（MNN-OpenCL）快于竞品三倍，具有比较明显的优势。

端侧大语言模型对话：

端侧大语言模型识图对话：

结语

本文深入探讨了MNN在端侧大模型部署上的最新进展，特别是在移动端设备上实现高效的大语言模型（LLM）和文生图模型（Diffusion）的部署。通过引入MNN-Transformer框架，MNN团队不仅解决了模型在移动端运行时的性能瓶颈，还通过一系列技术创新，如动态量化、磁盘映射技术和KV缓存管理，显著提升了模型的运行效率和内存利用率。未来，MNN将继续优化模型的运行效率，探索更低精度计算和更广泛的模型支持，以推动端侧AI技术的进一步发展。

团队介绍

我们是大淘宝技术Meta Team，负责面向消费场景的3D/XR基础技术建设和创新应用探索，通过技术和应用创新找到以手机及XR 新设备为载体的消费购物3D/XR新体验。团队在端智能、商品三维重建、3D引擎、XR引擎等方面有深厚的技术积累。团队在OSDI、MLSys、CVPR、ICCV、NeurIPS、TPAMI等顶级学术会议和期刊上发表多篇论文。