基于 MNN 在个人设备上流畅运行大语言模型

be5018a7557456e8fac51050c79ec509.gif

LLM(大语言模型)因其强大的语言理解能力赢得了众多用户的青睐,但LLM庞大规模的参数导致其部署条件苛刻;在网络受限,计算资源有限的场景下无法使用大语言模型的能力;低算力,本地化部署的问题亟待解决。ChatGLM-6B在60亿参数的情况下做到了优秀的中英文对话效果,且能够支持在消费级显卡本地部署;因此在HuggingFace Trends上很快登顶。6B的参数量虽然能够做到本地部署,但是目前的实现依赖库较多,如Pytorch, transfomer;对于端侧部署来说要求仍然较高。因此我们尝试将该模型转换为MNN模型,极大降低了部署时的依赖项,能够更方便的在各类端侧设备上部署与测试;同时我们对MNN模型进行了低bit量化,并实现了反量化与计算融合的计算kernel,大大降低了内存需求。实测PC端小显存显卡能够成流畅运行浮点模型,在Android手机上能够流畅运行量化模型。

代码实现:https://github.com/wangzhaode/ChatGLM-MNN

ab58627415816274b7363dcceb2fe836.png

模型导出

模型导出采用了Pytorch到ONNX到MNN的转换方式,并切对模型进行了拆分导出,将embedding,28 x GLMBlock, lm分别导出;并且在导出时对词表进行了瘦身。模型导出代码。

  导出方式

Pytorch实现的模型导出目前有2种主流方案:1. 导出为ONNX; 2. 导出为TorchScript。分析代码后可知,ChatGLM的模型结构比较简单,Embedding层,28层GLMBlock,线性层;其中GLMBlock结构如为LayerNorm -> SelfAttention -> LayerNorm -> MLP,代码如下:

attention_input = self.input_layernorm(hidden_states)# Self attention.
attention_outputs = self.attention(attention_input,position_ids,attention_mask=attention_mask,layer_id=layer_id,past_key_value=past_key_value,use_cache=use_cache,output_attentions=output_attentions
)attention_output = attention_outputs[0]outputs = attention_outputs[1:]# Residual connection.
alpha = (2 * self.num_layers) ** 0.5
hidden_states = attention_input * alpha + attention_outputmlp_input = self.post_attention_layernorm(hidden_states)# MLP.
mlp_output = self.mlp(mlp_input)# Second residual connection.
output = mlp_input * alpha + mlp_output

因为该模型结构简单,使用的算子 ONNX全部支持;同时MNN对ONNX的支持完备性比较好;因此选择使用ONNX导出模型。

  结构拆分

在确定使用ONNX之后首先直接使用torch.onnx.export尝试对模型进行导出,导出过程非常缓慢,导出后模型的权重大小有28G。在将模型转换到MNN时会执行一些图优化Pass;因为模型太大导致占用内存过高速度非常慢;因此考虑将模型进行拆分优化。拆分之后的优化考虑如下:

  1. Embedding层的参数大小为150528 * 4096, 单个权重使用内存非常大;考虑到输入文字数量较小(相对于150528),使用Gather实现消耗大量内存/显存,直接将参数存储为二进制文件,通过fseekfread实现Gather的操作能够在稍微牺牲速度的情况下节约2.3G内存;同时为了降低模型的文件大小,将embedding层的数据使用bf16格式存储,能够将文件大小降低一半,对精度和性能形象非常小。

  2. GLMBlock层的权重总大小为21G,仍然非常大,每个Block的大小为768M;考虑到要在端侧各类设备部署,可以将28层Block分别导出,对于浮点模型,这样的好处是能够在显存不足的情况下将部分Block放置在GPU,其余部分放置在CPU进行推理,这样能够充分利用设备算力;对与移动端设备,对这些block进行量化,分别获得int8/int4的权值量化模型,使用int4量化模型大小为2.6G,可以在端侧小内存设备部署。

  3. 线性层通过一个矩阵乘将hidden_state转换为词语的prob:[num, 4096] @ [4096, 150528];其实这里并不需要num全部参与运算,比如输入序列长度num = 10时,实际预测下一个词语时进需要使用最后一个[1, 4096]即可。因此可以先对输入变量做一个Gather然后执行矩阵乘:[1, 4096] @ [4096, 150528]即可。为了在端侧降低内存占用,这里同样使用int8/int4量化,量化后大小为256M。

  词表瘦身

词表大小为150528, 分析发现前20000个词语为<image>,在Chat中并没有使用,因此将可以将结构拆分后的Embedding层和最后的线性层进行删减。简单的方法是将2层的权重导出onnx模型,使用numpy.fromfile将onnx模型的权重加载,删除前[20000, 4096]的部分,在使用numpy.tofile保存即可。代码如下:

import numpy as np
embed = np.fromfile('transformer.word_embeddings.weight', dtype=np.float32, count=-1, offset=0)
embed = embed.reshape(-1, 4096) # shape is (150528, 4096)
embed = embed[20000:, :] # shape is (130528, 4096)
embed.tofile('slim_word_embeddings.bin')

对于删减后的词表,使用bf16格式存储可以降低一半的文件大小,使用C++代码将fp32转换为bf16,如下:

// read binary file
FILE* src_f = fopen("slim_word_embeddings.bin", "rb");
constexpr size_t num = 4096 * 130528;
std::vector<float> src_buffer(num);
fread(src_buffer.data(), 1, num * sizeof(float), src_f);
fclose(src_f);
// convert to bf16
std::vector<int16_t> dst_buffer(num);
for (int i = 0; i < num; i++) {dst_buffer[i] = reinterpret_cast<int16_t*>(src_buffer.data())[2 * i + 1];
}
// write to bianry file
FILE* dst_f = fopen("slim_word_embeddings_bf16.bin", "wb");
fwrite(dst_buffer.data(), 1, num * sizeof(int16_t), dst_f);
fclose(dst_f);

  动态形状

因为模型输入的形状是动态变化的,因此需要在导出时指定动态形状的维度,具体的导出方式如下:

def model_export(model,model_args: tuple,output_path: str,ordered_input_names,output_names,dynamic_axes,opset
):from torch.onnx import exportexport(model,model_args,f=output_path,input_names=ordered_input_names,output_names=output_names,dynamic_axes=dynamic_axes,do_constant_folding=True,opset_version=opset,verbose=False)model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, resume_download=True).float().cpu()
model_export(model,model_args=(torch.randn(4, 1, 4096),torch.tensor([[[[False, False, False,  True],[False, False, False,  True],[False, False, False,  True],[False, False, False, False]]]]),torch.tensor([[[0, 1, 2, 3], [0, 0, 0, 1]]]),torch.zeros(2, 0, 1, 32, 128)),output_path= "dyn_model/glm_block_{}.onnx".format(sys.argv[1]),ordered_input_names=["inputs_embeds", "attention_mask", "position_ids", "past_key_values"],output_names=["hidden_states", "presents"],dynamic_axes={"inputs_embeds" : { 0: "seq_len" },"attention_mask" : { 2: "seq_len", 3: "seq_len" },"position_ids" : { 2: "seq_len" },"past_key_values" : { 1: "history_len" }},opset= 14)

  其他问题

  • Tuple改为Tensor

原始实现中layer_past是Tuple,将其修改为Tensor方便模型导出后的模型输入。将代码中的Tuple操作替换为Tensor操作,如:

# 修改前
past_key, past_value = layer_past[0], layer_past[1]
key_layer = torch.cat((past_key, key_layer), dim=0)
value_layer = torch.cat((past_value, value_layer), dim=0)
present = (key_layer, value_layer)
# 修改后
key_layer = torch.cat((past_key_value[0], key_layer), dim=0)
value_layer = torch.cat((past_key_value[1], value_layer), dim=0)
present = torch.stack((key_layer, value_layer), dim=0)
  • view操作不支持动态形状

指定了动态维度后,在实际测试中发现因为模型实现中有些view相关代码导出后会将形状固定为常量,导致导出后改变输入形状无法正确推理,因此需要对模型中非动态的实现进行修改,将attention_fn函数中所有view操作替换为squeezeunsqueeze操作,这样导出后与形状无关即可实现动态形状。

# 修改前
query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
# 修改后
query_layer = query_layer.squeeze(1)
key_layer = key_layer.squeeze(1)
  • squeeze简化

squeeze(1)在模型导出时会产生了额外If算子,apply_rotary_pos_emb_index函数中使用squeeze(1)会使模型中多出2个If,为了让模型更加简单,这里可以直接替换为squeeze可以。

# 修改前
cos, sin = F.embedding(position_id, cos.squeeze(1)).unsqueeze(2), \F.embedding(position_id, sin.squeeze(1)).unsqueeze(2)
# 修改后
cos = F.embedding(position_id, torch.squeeze(cos)).unsqueeze(2)
sin = F.embedding(position_id, torch.squeeze(sin)).unsqueeze(2)

21db23d62e60e3234036cc3eb78b22d5.png

推理实现

用户输入n个词语,在前处理转换为n个int后,通过对embedding数据的查询会生成一个[n, 4096]的向量,同时根据输入长度和结束符位置,生成position_ids和mask作为输入;对于非首次生成还会有一个历史信息的输入。其中position_ids和mask对于28个block完全一样,history每个block都使用上次生成时对应block的present输出;而输入的input_embedding则使用上一个block的输出hidden_state,具体结构如下:

e2256634f38b2f9051d2506bfdbffb58.png

  前后处理

官方提供的实现中使用了transformers库,该库提供了模型的前后处理的实现。其中前处理包括了分词,将词语转换为ids;后处理中包含了prob转换为词语,控制模型持续生成的逻辑。在转换到C++之后我们也需要实现相同的前后处理逻辑。

  • 前处理

前处理逻辑是将用户输入的句子进行分词,然后查询词表将词语转换为id;C++中实现如下:

  1. 分词:在C++上使用cppjieba进行分词;

  2. word2id: 将词表文件加载为map,通过查询map将词语转换为id;

std::vector<int> ids;
std::vector<std::string> words;
cppjieba::Jieba jieba(...);
jieba.Cut(input_str, words, true);
for (auto word : words) {const auto& iter = mWordEncode.find(word);if (iter != mWordEncode.end()) {ids.push_back(iter->second);}
}
ids.push_back(gMASK);
ids.push_back(BOS);
return ids;
  • 后处理

后处理部分是将prob转换为id,然后通过词表将id转换为词语,同时将一些特殊字符进行转义;C++中实现如下:

  1. prob2id:在lm层后接一个ArgMax即可将prob转换为id,实测效果与transformers中的实现结果一致;

  2. id2word: 次变文件加载为vector,直接读取即可获取word;

  3. 特殊词处理:针对一些特殊词语进行了替换;

auto word = mWordDecode[id];
if (word == "<n>") return "\n";
if (word == "<|tab|>") return "\t";
int pos = word.find("<|blank_");
if (pos != -1) {int space_num = atoi(word.substr(8, word.size() - 10).c_str());return std::string(space_num, ' ');
}
pos = word.find("▁");
if (pos != -1) {word.replace(pos, pos + 3, " ");
}
return word;

  模型推理

将28层block依次执行推理,将其中的hidden_state输出作为下一个block的输入;并将present输出保存起来作为下次推理的history即可,代码如下:

// embedding
FILE* file = fopen("slim_word_embeddings.bin", "rb");
auto input_embedding = _Input({static_cast<int>(seq_len), 1, HIDDEN_SIZE}, NCHW);
for (size_t i = 0; i < seq_len; i++) {fseek(file, input_ids[i] * size, SEEK_SET);fread(embedding_var->writeMap<char>() + i * size, 1, size, file);
}
fclose(file);
// glm_blocks
for (int i = 0; i < LAYER_SIZE; i++) {auto outputs = mModules[i]->onForward({hidden_states, attention_mask, position_ids, mHistoryVars[i]});hidden_states = outputs[0];mHistoryVars[i] = outputs[1];
}
// lm
auto outputs = mModules.back()->onForward({hidden_states});
int id = outputs[0]->readMap<int>()[0];

508983ca83201139650d0a512f4a1d07.png

推理优化

  MNN Module接口

MNN的Module接口相比于Session接口,能够支持控制流,输入输出不需要再考虑设备问题,用户可以透明的使用GPU,CPU等不同设备串联推理且不需要关心数据的设备问题。同时Module接口在模型加载时可以对权重进行预重排,在内存充足的情况下提升卷积,矩阵乘等算子的推理速度。在对模型逐层拆分后会考虑到不同设备加载,Module接口可以更加简洁的实现这种跨设备的推理。

  PC端低显存推理(浮点)

上述转换与推理使用的模型都是浮点模型,在实际推理中可以选择fp32或者fp16。在使用fp16推理时,显存要求在13G以上;目前主流的游戏显卡显存普遍达不到该要求,因此无法将全部模型加载到GPU中推理。考虑到我们对模型进行了分段划分,可以将一部分block放入显存使用GPU推理,剩余部分使用CPU推理。因此可以根据用户指定的显存大小动态的分配block到GPU中。分配规则为,fp16的情况下每个block占用显存大小为385M,推理过程中的特征向量大小预留2G的显存,因此可以加载到GPU中的层数为:(gpu_memory - 2) * 1024.0 / 385.0。代码实现如下:

void ChatGLM::loadModel(const char* fileName, bool cuda, int i) {Module::Config config;config.shapeMutable = true;config.rearrange = true;auto rtmgr = cuda ? mGPURtmgr : mCPURtmgr;std::shared_ptr<Module> net(Module::load({}, {}, fileName, rtmgr, &config));mModules[i] = std::move(net);
}// load model
int gpu_run_layers = (gpu_memory - 2) * 1024.0 / 385.0;
for (int i = 0; i < LAYER_SIZE; i++) {sprintf(buffer, "../resource/models/glm_block_%d.mnn", i);loadModel(buffer, i <= gpu_run_layers, i);
}

  移动端低内存推理(量化)

全部浮点模型加载使用CPU推理需要32G左右的内存大小,在移动设备上很难满足内存要求。使用模型量化的方法来降低内存占用是常用的方法;MNN支持模型权重量化,ChatGLM在int4量化后总权重大小在4G左右,理论上是可以在移动端全部加载运行的。但是由于MNN之前的设计偏向于中小模型的推理,在模型加载阶段就做了反量化的计算,导致实际推理的内存占用与浮点一致,没有降低推理时内存占用。这种模式针对传统的端侧模型是非常合适的,在降低了模型大小的同时不会有性能损失;但是遇到内存总大小成为瓶颈大模型时则不是非常合适;因此需要针对大模型的大内存需求进行优化。针对大模型内存瓶颈问题,MNN在运行时使用low_memory选项会将反量化过程放在矩阵乘中实现,以部分推理时的额外计算开销大幅降低内存占用与访存带宽占用。

对于权值量化模型的低内存实现,我们支持了int4和int8两种权值量化的模型的低内存模式。针对不同的硬件做了实现,针对X86 SSE, AVX2实现了int4@fp32, int8@fp32;针对ARM64实现了int4@fp32, int8@fp32和int4@fp16和int8@fp16。具体的是线上需要针对以上列举的情况分别实现对应的矩阵乘Kernel,并且在原来的浮点矩阵乘的输入里增加反量化需要的alphabias参数,在矩阵乘计算前需要先从内存中加载常量的int4/int8量化值,然后将其转换为浮点类型,之后再执行浮点矩阵乘操作,实际的矩阵乘基础操作如下公式:

979a5452d7d06a017b824d2835c8107b.png

以下为int4量化模型在ARMv8.2上的fp16矩阵乘实现:

mov x15, x1
ld1 {v12.8h, v13.8h}, [x14], #32 // alpha
mov w17, #0x0f
dup v3.16b, w17
mov w17, #7
dup v4.16b, w17
ld1 {v14.8h, v15.8h}, [x16], #32 // bias
subs x12, x9, #2
// load int4 weight
ld1 {v0.8h}, [x13], #16
// int4 to fp16
ushr v1.16b, v0.16b, #4
and v2.16b, v0.16b, v3.16b
sub v1.16b, v1.16b, v4.16b
sub v2.16b, v2.16b, v4.16b
zip1 v10.16b, v1.16b, v2.16b
zip2 v11.16b, v1.16b, v2.16b
sxtl v1.8h, v10.8b
sxtl2 v2.8h, v10.16b
scvtf v1.8h, v1.8h
scvtf v2.8h, v2.8h
mov v8.8h, v14.8h
mov v9.8h, v15.8h
// get fp16 in v8, v9
fmla v8.8h, v1.8h, v12.8h
fmla v9.8h, v2.8h, v13.8h// fp16 GEMM kernel
ld1 {v0.8h}, [x15], x11
fmul v16.8h, v8.8h, v0.h[0]
fmul v17.8h, v8.8h, v0.h[1]
fmul v18.8h, v8.8h, v0.h[2]
fmul v19.8h, v8.8h, v0.h[3]
...

52b88b76fd39fe8f576e892d92e64148.png

性能测试

  1. PC端测试:11G显存的2080Ti + AMD 3900X + 32G内存测;使用fp32精度模型(GPU显存不足情况下)GPU+CPU混合速度为3.5 tok/s; 仅使用CPU速度为 1.2 tok/s ;

  2. 移动端测试:Xiaomi12;使用int4模型精度,CPU速度为 1.5 tok/s,需要内存为 2.9 G。

9bc29f65edd0b5afc618d09d10d12dda.png

用户界面

  PC

在PC端提供了两种demo的用法,命令行与web的用户界面;

19bfef720536f93eeb192eab32f2f1db.png

  移动端

在移动端目前提供了Android的App;

1ac8cbfda63267e323f46b6ceb820bc6.png

9cdd72e463acff14647aeee4ceffc025.png

总结

ChatGLM-6B模型推理主要是内存瓶颈,针对这一问题本文提出了2中优化手段:分段加载;量化。使用了这两种优化方法后,我们能够保证模型精度完全无损(fp32)的情况下在消费级显卡的PC机器上部署运行;同时还支持略微损失模型精度(int4)的情况下在移动端设备上流畅运行。

a44b2fe6d149f8820d35accddeb8513b.png

团队介绍

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

本篇内容作者:王召德(雁行)

¤ 拓展阅读 ¤

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

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

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.hqwc.cn/news/21781.html

如若内容造成侵权/违法违规/事实不符,请联系编程知识网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

【测试开发】自动化测试 selenium 篇

目录 一. 什么是自动化测试 二. selenium 1. selenium的工作原理 2. seleniumJava的环境搭建(Chrome浏览器) 三. selenium中常用的API 1. 定位元素 findElement 1.1 css选择语法 1.2 xpath 2. 操作测试对象 2.1 sendKeys-在对象上模拟按键输入 2.2 click-点击对象…

IP 协议的相关特性

目录 IP协议有三大特点&#xff1a;无连接、无状态、不可靠。 四位版本号 四位头部长度 八位服务类型: 十六位总长度 16 位标识, 3 位标志, 13 位片偏移 八位生存时间 八位协议 十六位首部校验和 关于IP v4地址不够的问题 ip地址动态分配: ip地址转换(NAT) 数据传输…

【xxl-job】分布式任务调度系统xxl-job搭建

XXL-JOB是一个轻量级分布式任务调度平台&#xff0c;其核心设计目标是开发迅速、学习简单、轻量级、易扩展、开箱即用。 更多介绍&#xff0c;请访问官网&#xff1a;分布式任务调度平台XXL-JOB 一、任务调度中心(基于docker)【Version 2.4.0】 前提条件&#xff1a;任务调度…

复习Javascript数组

JavaScript 数组 JS 数字属性JS 数组方法 JavaScript 数组用于在单一变量中存储多个值。 var cars ["Saab", "Volvo", "BMW"]; 什么是数组&#xff1f; 数组是一种特殊的变量&#xff0c;它能够一次存放一个以上的值。 如果您有一个项目清…

MySQL表单查询以及多表查询

1.单表查询 CREATE TABLE emp ( empno int(4) NOT NULL, ename varchar(255) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL, job varchar(255) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL, mgr int(4) NULL DEFAULT NULL, hi…

IOP JPCS独立出版 | 2023年第二届材料科学与工程国际会议(CoMSE 2023)

会议简介 Brief Introduction 2023年第二届材料科学与工程国际会议(CoMSE 2023) 会议时间&#xff1a;2023年7月21日-23日 召开地点&#xff1a;中国泰州 大会官网&#xff1a;www.icomse.org CoMSE 2023由四川大学、华南理工大学亚热带建筑科学国家重点实验室、国际电气电子与…

Raft算法之日志复制

Raft算法之日志复制 一、日志复制大致流程 在Leader选举过程中&#xff0c;集群最终会选举出一个Leader节点&#xff0c;而集群中剩余的其他节点将会成为Follower节点。Leader节点除了向Follower节点发送心跳消息&#xff0c;还会处理客户端的请求&#xff0c;并将客户端的更…

如何在 Windows 中免费合并 PDF 文件 [在线和离线]

PDF是一种广泛使用的文件格式&#xff0c;具有兼容性好、安全性高、易于打印、方便浏览等众多优点。在工作和学习过程中&#xff0c;经常需要将同一类型的PDF文件合并起来&#xff0c;以方便传输和查看&#xff0c;使得合并PDF文件成为一种重要的数据整合方法。 如果您想知道如…

android更换开机动画

android11 路径&#xff1a;device / {vendor-name} / {platform-name} / {device-name} / system / bootanimation.zip 例&#xff1a;android \ device \ softwinner \ ceres \ ceres-b6 \ system \ bootanimation.zip android13 路径&#xff1a;device / softwinner / {PRO…

基于单片机空气质量检测二氧化碳 一氧化碳温湿度PM2.5检测系统的设计与实现

功能介绍 以51单片机作为主控系统&#xff1b;对空气空气中有毒有害气体进行监测&#xff1b;使用LCD1602液晶显示&#xff0c;采集到的PM2.5值通过单片机串口传输&#xff1b;通过传感器对室内PM2.5粉尘进行检查&#xff1b;通过按键设置的上限值&#xff1b;当检测到有毒气体…

iOS开发 - NotificationService语音播报

iOS NotificationService语音播报 最近碰到个接收到推送要实现语音播报的需求&#xff0c;需要后台推送通知&#xff0c;APP客户端收到通知之后语音播放&#xff1a;“您的账户收到一笔巨款”的功能。 因为工程之前已经集成了极光推送服务。这里直接使用Notification Service…

Vue3+Vite+Pinia+Naive后台管理系统搭建之八:构建 login.vue 登录页

前言 如果对 vue3 的语法不熟悉的&#xff0c;可以移步Vue3.0 基础入门&#xff0c;快速入门。 项目所需要的图片&#xff0c;icon图标&#xff08;推荐&#xff1a;阿里巴巴矢量图标库&#xff09;自行获取&#xff0c;命名一致就行。 1. 构建 src/components/CopyRight.vu…