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

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_output
mlp_input = self.post_attention_layernorm(hidden_states)
# MLP.mlp_output = self.mlp(mlp_input)
# Second residual connection.output = mlp_input * alpha + mlp_output

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')

// 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 export export( 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)

# 修改前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)

# 修改前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)

# 修改前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)

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;

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];

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);}

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]...