01.概述
随着移动端（手机/平板等）算力、内存、磁盘空间的不断增长，在移动端部署大模型逐渐成为可能。在端侧运行大模型，可以有一系列好处：去除网络延迟，加快响应速度；降低算力成本，便于大规模应用；不需数据上传，保护用户稳私。
为了在更广泛的设备上部署大模型，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推理端侧模型均已在魔搭社区开源：
https://modelscope.cn/organization/MNN
效果体验视频（LLM）：
https://developer.aliyun.com/live/254741?spm=a2c6h.26396819.creator-center.8.404f3e184lgDmA
效果体验视频（多模态大模型）：
https://developer.aliyun.com/live/254740?spm=a2c6h.26396819.creator-center.10.404f3e18rb0Xb3
02.MNN系统架构
如上述架构图所示，MNN-Transformer由导出工具、量化工具、插件与引擎三个部分组成：导出工具负责将各类大模型转换为MNN格式，并构建必需的资源包；量化工具减少MNN模型大小，以降低运行时内存，加快执行速度；LLM/Diffusion运行时所需要的分词、KV缓存管理、LoRA等功能由插件与引擎模块提供。
模型导出
在将深度学习模型从研究原型转换为实际可部署的产品时，模型导出阶段的顺畅与否对于整个工作流程至关重要。通常，这个过程涉及将模型从训练框架中导出到一个中间表示，如 ONNX（开放神经网络交换格式），然后再转换为目标部署框架——在本例中为 MNN格式。为了简化并标准化这一过程，研究团队开发了一个名为 llm-export 的工具。
llm-export 工具的核心思想在于对大型语言模型（LLM）进行了高度抽象，建立了一个统一化的导出框架。这个项目的目标是消除将各种 LLM 模型导出到 ONNX 格式的障碍，确保无论是何种架构的 LLM 都能通过一个清晰定义的接口进行处理。在 llm-export 中，研究团队定义了一套公用的导出逻辑，这意味着对于任何特定的 LLM，开发者只需实现模型的加载逻辑。这极大地减少了从多样化的训练环境向 ONNX 模型迁移的复杂性，并显著提高了整个导出过程的易用性。模型一旦被成功导出至 ONNX，即可利用现有的mnnconver工具转换到 MNN 格式，从而使用MNN完成llm模型的推理。
llm-export中将llm模型抽象为4部分：tokenizer, embedding, blocks, lm；主要代码如下：
class LLM(torch.nn.Module): def __init__(self, args): super().__init__() # load tokenizer, embed, blocks, lm self.load_model(args.path) def forward(self, input_ids, attention_mask, position_ids, past_key_values): hidden_states = self.embed(input_ids) presents = [] for i in range(self.block_nums): hidden_states, kv = self.blocks[i](hidden_states, attention_mask, position_ids, past_key_values[i]) presents.append(kv) token_id = self.lm(hidden_states).view(1) presents = torch.stack(presents) return token_id, presents def export(self): # export llm to onnx and mnn ... class Chatglm2_6b(LLM): def load_model(self, model_path: str): # chatglm2 load impl ... class Qwen_7b(LLM): def load_model(self, model_path: str): # qwen load impl ...
模型部署
在部署大型语言模型（LLM）时，兼容性和易用性是关键因素。为了解决这一挑战，MNN团队开发了一个名为 mnn-llm 的项目。考虑到MNN在跨平台上支持上的优秀表现，该项目基于 MNN 构建，旨在为各种平台提供一个统一的 LLM 模型部署解决方案。mnn-llm 项目使得从 llm-export 导出的模型能够无缝进行端到端的推理，并为开发者提供了一个简易的文本到文本（txt2txt）调用接口。
在mnn-llm中移植实现了目前主流的tokenizer工具：Sentencepiece 和 Tiktoken。这些 tokenizer 组件是处理自然语言输入的关键部分，它们能够将原始文本转换为模型能理解的格式。同时为了轻量化，两种模型都使用文本的方式存储，移除了Sentencepiece中迪对protobuf的依赖。
此外，考虑到内存占用在移动设备上尤为宝贵，同时还在 mnn-llm 中引入了 disk embedding 功能。这意味着用户可以根据需要选择：在模型推理过程中使用 embedding 模型在内存计算，或者直接从磁盘加载 embedding 值。这种灵活性不仅降低了运行时的内存需求，也为用户提供了更多的选择来平衡推理性能和资源使用。为了确保 mnn-llm 的通用性和扩展性，研究团队设计了一种易于扩展的架构。开发者可以通过继承基类并实现特定的配置来支持不同的 LLM 模型。这种架构设计使得整合新的 LLM 模型变得简单快捷，大大降低了将最新的研究成果应用到实际产品中的门槛。
性能优化
内存优化
DRAM-Flash混合存储
在移动设备上部署大型语言模型的主要瓶颈在于DRAM的限制。MNN-LLM采用混合存储策略，以缓解内存使用情况，并确保在受限内存条件下保持大型语言模型推理的可用性。尽管闪存存储容量比DRAM大得多，但其读取速度要慢得多；MNN-LLM的混合存储策略针对模型的操作特性进行了定制：对于参数存储，它评估利用率并将低利用率参数分配给闪存以最小化速度影响。对于KV数据，预取技术被用于降低闪存读取延迟，从而减轻对性能的影响。
大型语言模型 (LLM) 模型的大量参数是其高内存消耗的主要原因。结构上，这些参数可以分为三个类别：embedding、layers和Lm head。embedding和Lm head参数的大小通常计算为词汇表大小×隐藏单元大小，由于词汇表大小通常很大，因此嵌入参数不会像其他参数一样参与计算。层是指每个解码器层中的参数，包括注意力和MLP线性层，通常大小为隐藏单元大小×隐藏单元大小或中间大小×隐藏单元大小，在各层中具有相同的参数规模。在Qwen2 7B 模型中，不参与计算的嵌入参数约占总参数量的15％。
利用Flash来存储Embedding层允许在不显著影响推理性能的情况下减少DRAM的使用。例如，Qwen-7B通过使用bfloat16存储可以减少大约2.18GB的DRAM使用，大大提高了对内存受限的移动设备模型推理的可行性。
在输入文本较长或生成长度较广的场景中，连续增长的KV缓存可能导致显著内存使用。MNN-LLM通过采用混合存储策略来解决这一挑战，利用闪存来保存部分KV缓存，从而确保在长上下文条件下仍能进行LLM推理。最初，所有KV缓存值都存储在DRAM中，但随着上下文的扩展和KV缓存大小的增长，任何超过一定阈值的部分都会转移到Flash上。由于每个计算仅产生一组新的KV值，Qwen2 7B模型的总KV值约为1KB，最大限度地减少了存储开销。
随着存储在闪存中的KV缓存值数量的增加，从闪存中加载它们所需的时间将逐渐增加，并且可以减慢推理速度，如图所示。为了减轻从闪存中加载KV缓存对推理时间的影响，MNN实现了预取：在当前层的MLP阶段和下一层的qkv投影阶段，在内存中预取闪存中的KV缓存值。当预取时间小于或等于计算时间时，LLM推理速度不受影响。
组合量化
大型语言模型（LLM）的参数量大是其高内存消耗的主要原因，量化可以显著减少参数量，从而降低内存使用。然而，量化会影响模型的推理精度；一般来说，较低位数会导致信息损失更大，并对准确性产生更大的影响。有各种方法、数据类型和位数可用于量化，因此选择适当的量化方法以平衡内存使用、运行时性能和模型精度至关重要。
对于Embedding、layer和Lm head的参数，MNN-LLM采用组合量化策略来平衡精度和计算开销。embedding的权重约占总模型重量的15％。由于这些权重在每个解码步骤中只使用一小部分，在闪存中存储它们不会占用DRAM空间。这允许使用bfloat16存储，确保计算准确性。非嵌入参数包括layer和LM head的权重，必须完全在每个计算中加载，使得它们的大小对推理性能产生显著影响。特别是，在解码阶段，由于内存受限，推断时间直接与这些参数的大小成正比。因此，对于这些权重使用低比特量化至关重要。考虑到精度和硬件计算指令——边缘CPU特别适合int8运算——这些参数被量化为int4或int8。非对称量化如下：
在处理长上下文时，KV缓存的内存使用量会继续增长，并且量化策略可以有效地减少这种内存消耗。MNN-LLM根据其计算角色为键和值提供不同的量化方法。在注意力计算中，查询、键和值的形状是 [头数，序列长度，头维度] 。当对键和查询进行矩阵乘法运算时，被缩减的维度是头维度，这是一个固定值。
因此，可以将int4/int8量化应用于键上，允许新键值直接量化并存储。相反，在与值进行注意力得分矩阵乘法运算时，被缩减的维度是序列长度。对于值采用int4/int8量化会影响现有值的数据分布，需要更新它们的量化值并且产生额外开销。为此，MNN-LLM采用了fp8量化用于值，允许新值直接量化而不会影响现有的值。
计算优化
硬件驱动的数据重排序
对大型语言模型（LLMs）的推理过程进行分析，发现主要耗时的操作是线性运算和注意力机制，这两种操作本质上都依赖于矩阵乘法。因此，优化这些操作中的矩阵乘法对于提高LLM性能至关重要。循环块化是一种常见的优化技术，它增强了内存访问的局部性，显著影响了性能。循环块化的最佳块大小极大地影响最终的矩阵乘法性能，并受到设备内存、缓存和计算硬件的影响。因此，在硬件和数据规模的基础上选择最合适的数据重组织和计算方法以实现峰值性能非常重要。MNN-LLM采用一种针对这两种操作类型的计算特征而定制的硬件驱动的数据重新排序策略来确定块化方法和大小，从而优化LLM推理性能。
对于Attention操作，使用与Linear相同的重新排列策略。键和值直接存储在重新排列的数据布局中，确保无需在每次计算期间重新排列历史KV。
多核工作负载平衡
现代CPU通常具有多个内核，因此有效地利用多核心计算能力对于优化性能至关重要。MNN-LLM利用CPU的多线程并行性将操作平行化沿seqlen和-dimensions方向。MNN-LLM在启动时根据其实际计算能力为不同内核指定计算负载。在并行计算期间，MNN-LLM根据内核的负载率分配计算工作量。这种平衡的工作负载分布策略与均匀工作负载策略相比可以提高多线程计算性能。
主流移动SoC通常具有一个主核心和三个性能核心，例如Snapdragon 8 Gen 3。高负载计算一般使用主核心和性能核心。当线程数超过一个时，在主核心与性能核心之间发生并行计算，如图4所示。在这种情况下，工作量平衡显著提高了多线程加速比，相比于均匀的工作负载分布。
混合浮点精度
在之前的矩阵操作讨论中，使用了低精度量化方法来加速计算。对于非矩阵乘法运算，MNN-LLM也支持混合精度结果，确保准确性的同时提升推理性能。ARMv8.2和更新的CPU支持float16计算，相比float32可以节省一半内存，并且float16 NOEN指令的吞吐量是float32的两倍。然而，float16有一些精度限制；对于需要更宽精度范围的计算，可能会出现显著误差，尤其是在值超过65,504时尤其明显。为了应对这一问题，MNN-LLM采用混合精度策略以保持推理准确度。在LLM推理过程中，注意力中的Softmax计算对数据精度非常敏感，因此MNN-LLM保证Softmax使用float32。在查询与键的矩阵相乘中，查询值可能很大，在累积后可能导致溢出。为了解决这个问题，可以直接将查询除以dk从而缩小其值范围并防止最终结果发生溢出。这种方法优化了整体内存使用和推理性能，同时保持了准确性。
几何计算
LLMs的计算图还包括诸如转置、聚集和连接等长尾操作。尽管这些操作可能不会显著影响整体执行时间，但在数据规模较大时它们会导致大量内存访问。为了应对这些长尾操作，MNN-LLM采用了几何计算方法， 抽象所有数据重新排列操作为地址线性映射。
对于计算图中的连续数据重新排序操作，此抽象产生大量连续的区域。MNN-LLM 实现了一个基于循环展开、循环交换、循环分区和循环融合规则的自动区域融合算法。该算法可以自动合并兼容的区域，从而减少数据重新排序操作的数据读写操作次数，并提高性能。通过利用几何运算对 LLM 模型进行推理，可以降低长尾操作的开销，提升约 3% 的性能。
LoRA优化
在移动设备上，不同的任务可能需要不同类型的LLM模型。由于模型参数数量庞大，直接使用多个模型会导致带宽和存储资源的过度消耗。因此，在多任务处理中，结合基线模型与多个LoRA模型是一种更高效的解决方案。
MNN-LLM支持合并的LoRA模型部署和多个LoRA模型的在线加载。在使用多个LoRA模型时，首先加载基模型，然后加载LoRA模型的计算图和权重，并且LoRA模型共享基模型的权重。由于LoRA权重通常较小，因此内存开销很小。与预合并方法相比，在线加载LoRA模型更加灵活，适用于多任务场景，虽然会增加额外的计算成本。
评估
实验设计基于量化模型，即Qwen2 1.5B、Qwen2 7B和Llama3 8B。使用小米14作为测试设备，并在CPU（利用4个线程）和GPU（通过OpenCL）上进行推理效果的比较评估。
使用诸如 llama.cpp、MLC-LLM 和 fastllm 等推理引擎的架构。由于 MLC-LLM 不支持基于 CPU 的推理，而 fastllm 缺乏对 GPU 的兼容性，因此这些引擎的相关实验被排除在外。在长度不同的提示（64、256 和 1024 个令牌）下进行了广泛的试验，并且在解码阶段施加了最多 16 个令牌的限制。
由于MLC-LLM在处理非对称量化模型时表现不佳，因此报告的MLC-LLM结果基于对称量化模型，但竞争引擎明确参与了使用非对称模型进行推理的任务。性能结果以预填充和解码速度的形式反映，并图形化地表示在下图中。
在CPU基准测试中，MNN-LLM表现出色，在预填充速度方面比llama.cpp快了8.6倍，并且比fastllm快了20.5倍。同时，解码速度也分别提高了2.3倍和8.9倍。而在基于GPU的评估中，与MLC-LLM相比，MNN-LLM的表现略有下降，特别是在使用Qwen2-7B时，由于MLC-LLM采用了优势对称量化技术，导致性能有所降低。然而，MNN-LLM仍然表现卓越，在预填充速度上比llama.cpp快了最多25.3倍，并且在解码速度上比llama.cpp快了7.1倍；此外，相对于MLC-LLM，其分别实现了2.8倍和1.7倍的改进。
02.总结与展望
在大模型端侧部署上，基于 MNN 实现的 mnn-llm 项目已经展现出业界领先的性能，特别是在 ARM 架构的 CPU 上。目前利用 mnn-llm 的推理能力，qwen-1.8b在mnn-llm的驱动下能够在移动端达到端侧实时会话的能力，能够在较低内存（<2G）的情况下，做到快速响应，可以下载 qwen-1.8b-apk来体验。然而，端侧部署 LLM 仍然面临着一系列挑战。尽管 qwen-1.8b 模型能在端侧设备上达到实用性能，但更大规模的模型如 6b 或 7b 仍然存在较高的部署门槛。这些模型往往要求更多的内存（~3GB）并且在处理较长文本（~ 100 token）时，很难在 1 秒内给出响应，这限制了其在实时应用场景下的可用性。
MNN后续会在如下方面进一步研发：
● 进一步优化高通设备上低精度计算的GPU运行性能
● 研究更有效的量化方案，并支持更低 Bit 的模型运行及中低端机型性能优化
● 提供输出限制条件情况下的优化方案
● 支持更多大语言模型和扩散模型
● 支持CPU/GPU/NPU混用计算
点击链接阅读全文：ModelScope 魔搭社区
