基于 MNN 在个人设备上流畅运行大语言模型
LLM（大语言模型）因其强大的语言理解能力赢得了众多用户的青睐，但LLM庞大规模的参数导致其部署条件苛刻；在网络受限，计算资源有限的场景下无法使用大语言模型的能力；低算力，本地化部署的问题亟待解决。ChatGLM-6B在60亿参数的情况下做到了优秀的中英文对话效果，且能够支持在消费级显卡本地部署；因此在HuggingFace Trends上很快登顶。6B的参数量虽然能够做到本地部署，但是目前的实现依赖库较多，如Pytorch, transfomer；对于端侧部署来说要求仍然较高。因此我们尝试将该模型转换为MNN模型，极大降低了部署时的依赖项，能够更方便的在各类端侧设备上部署与测试；同时我们对MNN模型进行了低bit量化，并实现了反量化与计算融合的计算kernel，大大降低了内存需求。实测PC端小显存显卡能够成流畅运行浮点模型，在Android手机上能够流畅运行量化模型。
代码实现：https://github.com/wangzhaode/ChatGLM-MNN
模型导出采用了Pytorch到ONNX到MNN的转换方式，并切对模型进行了拆分导出，将embedding，28 x GLMBlock， lm分别导出；并且在导出时对词表进行了瘦身。模型导出代码。
▐ 导出方式
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.5hidden_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；因为模型太大导致占用内存过高速度非常慢；因此考虑将模型进行拆分优化。拆分之后的优化考虑如下：
-
Embedding层的参数大小为
150528 * 4096, 单个权重使用内存非常大；考虑到输入文字数量较小（相对于150528），使用Gather实现消耗大量内存/显存，直接将参数存储为二进制文件，通过fseekfread实现Gather的操作能够在稍微牺牲速度的情况下节约2.3G内存；同时为了降低模型的文件大小，将embedding层的数据使用bf16格式存储，能够将文件大小降低一半，对精度和性能形象非常小。 -
GLMBlock层的权重总大小为21G，仍然非常大，每个Block的大小为768M；考虑到要在端侧各类设备部署，可以将28层Block分别导出，对于浮点模型，这样的好处是能够在显存不足的情况下将部分Block放置在GPU，其余部分放置在CPU进行推理，这样能够充分利用设备算力；对与移动端设备，对这些block进行量化，分别获得int8/int4的权值量化模型，使用int4量化模型大小为2.6G，可以在端侧小内存设备部署。
-
线性层通过一个矩阵乘将hidden_state转换为词语的prob：
[num, 4096] @ [4096, 150528]；其实这里并不需要num全部参与运算，比如输入序列长度num = 10时，实际预测下一个词语时进需要使用最后一个[1, 4096]即可。因此可以先对输入变量做一个Gather然后执行矩阵乘：[1, 4096] @ [4096, 150528]即可。为了在端侧降低内存占用，这里同样使用int8/int4量化，量化后大小为256M。
▐ 词表瘦身
numpy.fromfile 将onnx模型的权重加载，删除前 [20000， 4096] 的部分，在使用 numpy.tofile 保存即可。代码如下：
import numpy as npembed = 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 fileFILE* 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 bf16std::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 fileFILE* 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操作替换为squeeze和unsqueeze操作，这样导出后与形状无关即可实现动态形状。
# 修改前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简化
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)
[n, 4096] 的向量，同时根据输入长度和结束符位置，生成position_ids和mask作为输入；对于非首次生成还会有一个历史信息的输入。其中position_ids和mask对于28个block完全一样，history每个block都使用上次生成时对应block的present输出；而输入的input_embedding则使用上一个block的输出hidden_state，具体结构如下：
▐ 前后处理
官方提供的实现中使用了transformers库，该库提供了模型的前后处理的实现。其中前处理包括了分词，将词语转换为ids；后处理中包含了prob转换为词语，控制模型持续生成的逻辑。在转换到C++之后我们也需要实现相同的前后处理逻辑。
-
前处理
-
分词：在C++上使用 cppjieba进行分词； -
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;
-
后处理
-
prob2id：在lm层后接一个ArgMax即可将prob转换为id，实测效果与 transformers中的实现结果一致； -
id2word: 次变文件加载为vector，直接读取即可获取word； -
特殊词处理：针对一些特殊词语进行了替换；
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;
▐ 模型推理
// embeddingFILE* 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_blocksfor (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];}// lmauto outputs = mModules.back()->onForward({hidden_states});int id = outputs[0]->readMap<int>()[0];
推理优化
▐ MNN 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 modelint 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);}
▐ 移动端低内存推理（量化）
low_memory 选项会将反量化过程放在矩阵乘中实现，以部分推理时的额外计算开销大幅降低内存占用与访存带宽占用。
对于权值量化模型的低内存实现，我们支持了int4和int8两种权值量化的模型的低内存模式。针对不同的硬件做了实现，针对X86 SSE, AVX2实现了int4@fp32, int8@fp32；针对ARM64实现了int4@fp32, int8@fp32和int4@fp16和int8@fp16。具体的是线上需要针对以上列举的情况分别实现对应的矩阵乘Kernel，并且在原来的浮点矩阵乘的输入里增加反量化需要的alpha和bias参数，在矩阵乘计算前需要先从内存中加载常量的int4/int8量化值，然后将其转换为浮点类型，之后再执行浮点矩阵乘操作，实际的矩阵乘基础操作如下公式：
mov x15, x1ld1 {v12.8h, v13.8h}, [x14], #32 // alphamov w17, #0x0fdup v3.16b, w17mov w17, #7dup v4.16b, w17ld1 {v14.8h, v15.8h}, [x16], #32 // biassubs x12, x9, #2// load int4 weightld1 {v0.8h}, [x13], #16// int4 to fp16ushr v1.16b, v0.16b, #4and v2.16b, v0.16b, v3.16bsub v1.16b, v1.16b, v4.16bsub v2.16b, v2.16b, v4.16bzip1 v10.16b, v1.16b, v2.16bzip2 v11.16b, v1.16b, v2.16bsxtl v1.8h, v10.8bsxtl2 v2.8h, v10.16bscvtf v1.8h, v1.8hscvtf v2.8h, v2.8hmov v8.8h, v14.8hmov v9.8h, v15.8h// get fp16 in v8, v9fmla v8.8h, v1.8h, v12.8hfmla v9.8h, v2.8h, v13.8h// fp16 GEMM kernelld1 {v0.8h}, [x15], x11fmul 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]...
-
PC端测试：11G显存的2080Ti + AMD 3900X + 32G内存测；使用fp32精度模型（GPU显存不足情况下）GPU+CPU混合速度为3.5 tok/s; 仅使用CPU速度为 1.2 tok/s ； -
移动端测试：Xiaomi12；使用int4模型精度，CPU速度为 1.5 tok/s，需要内存为 2.9 G。
▐ PC
▐ 移动端
总结
大淘宝技术Meta Team，负责面向消费场景的3D/XR基础技术建设和创新应用探索，通过技术和应用创新找到以手机及XR 新设备为载体的消费购物3D/XR新体验。团队在端智能、商品三维重建、3D引擎、XR引擎等方面有深厚的技术积累。先后发布端侧推理引擎MNN，端侧实时视觉算法库PixelAI，商品三维重建工具Object Drawer等技术。团队在OSDI、MLSys、CVPR、ICCV、NeurIPS、TPAMI等顶级学术会议和期刊上发表多篇论文。
本篇内容作者：王召德（雁行）
本文分享自微信公众号 - 大淘宝技术（AlibabaMTT）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。
