AI知识库

53AI知识库

学习大模型的前沿技术与行业应用场景


MNN在大模型端侧部署上的探索
发布日期:2024-11-20 18:42:25 浏览次数: 1555 来源:大淘宝技术






随着移动端(手机/平板等)算力、内存、磁盘空间的不断增长,在移动端部署大模型逐渐成为可能。在端侧运行大模型,可以有一系列好处:去除网络延迟,加快响应速度;降低算力成本,便于大规模应用;不需数据上传,保护用户稳私。


概述

为了在更广泛的设备上部署大模型,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.cpphttps://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 输入下对比图:


详细数据如下:

Qwen2-1.5B

prefill speed -- decode speed (tok/s)

prompt: 64 token

prompt: 256 token

prompt: 1024 token

MNN-LLM

CPU四线程

225.19 -- 53.65

298.89 -- 52.45

236.76 -- 45.70

OpenCL

153.42 -- 26.92

166.82 -- 26.33

237.31 -- 20.87

llama.cpp



CPU四线程

46.22 -- 30.05

42.06 -- 26.39

37.37 -- 22.07

OpenCL

6.05 -- 4.44

16.08 -- 4.26

23.09 -- 3.30

fastllm

CPU四线程

26.33 -- 9.29

25.91 -- 8.72

19.45 -- 6.47

mlc-llm

OpenCL

137.6 -- 25

146.3 -- 19.5

83.5 -- 11.9

Qwen2-7B

prefill speed -- decode speed (tok/s)

prompt: 64 token

prompt: 256 token

prompt: 1024 token

MNN-LLM

CPU四线程

61.24 -- 13.00

67.35 -- 11.77

59.27 -- 10.08

OpenCL

34.81 -- 8.91

37.42 -- 7.67

54.24 -- 6.98

llama.cpp



CPU四线程

7.78 -- 6.41

8.08 -- 5.90

6.86 -- 4.28

OpenCL

1.79 -- 1.65

4.45 -- 1.39

5.84 -- 1.23

fastllm

CPU四线程

4.32 -- 1.61

4.97 -- 1.80

3.62 -- 1.13

mlc-llm

OpenCL

33.8 -- 10.4

36.6 -- 9.0

卡死

Llama3-8B

prefill speed -- decode speed (tok/s)

prompt: 64 token

prompt: 256 token

prompt: 1024 token

MNN-LLM

CPU四线程

53.63 -- 11.23

65.25 -- 10.32

54.16 -- 8.97

OpenCL

32.25 -- 9.16

32.57 -- 9.36

56.78 -- 7.88

llama.cpp



CPU四线程

8.24 -- 6.03

7.65 -- 5.54

7.18 -- 5.01

OpenCL

1.64 -- 1.30

3.60 -- 1.31

4.15 -- 1.33

fastllm

CPU四线程

3.16 -- 1.28

3.18 -- 1.27

3.06 -- 1.08

mlc-llm

OpenCL

26.1 -- 9.60

33.8 -- 9.2

24.8 -- 7.0

  MNN-Diffusion

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

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


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

stable diffusion v1.5

512x512图

Android Mi 14

Apple Mac M3

CPU (int8)

GPU (float16)

CPU (int8)

GPU (float32)

MNN-diffusion

4.2 s/iter

2.0 s/iter

4.1 s/iter

1.1 s/iter

stable-diffusion.cpp

>1 min/iter

不支持

9.52s/iter

4.9 s/iter

Android OnnxRuntime

6.5 s/iter

不支持

部分示例


端侧大语言模型对话:


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

结语


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


团队介绍

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



本文作者:霞影




¤ 拓展阅读 ¤

3DXR技术 | 终端技术 | 音视频技术

服务端技术 | 技术质量 | 数据算法







53AI,企业落地应用大模型首选服务商

产品:大模型应用平台+智能体定制开发+落地咨询服务

承诺:先做场景POC验证,看到效果再签署服务协议。零风险落地应用大模型,已交付160+中大型企业

联系我们

售前咨询
186 6662 7370
预约演示
185 8882 0121

微信扫码

与创始人交个朋友

回到顶部

 
扫码咨询