基于 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 src_buffer(num);fread(src_buffer.data(), 1, num * sizeof(float), src_f);fclose(src_f);// convert to bf16std::vector dst_buffer(num);for (int i = 0; i     dst_buffer[i] = reinterpret_cast(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 ids;std::vector<: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 == "") return "n";if (word == "") return "t";int pos = word.find("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(seq_len), 1, HIDDEN_SIZE}, NCHW);for (size_t i = 0; i     fseek(file, input_ids[i] * size, SEEK_SET);    fread(embedding_var->writeMap() + i * size, 1, size, file);}fclose(file);// glm_blocksfor (int i = 0; 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()[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 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     sprintf(buffer, "../resource/models/glm_block_%d.mnn", i);    loadModel(buffer, 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]...