基于MNN在個(gè)人設(shè)備上流暢運(yùn)行大語(yǔ)言模型該如何實(shí)現(xiàn)呢？

LLM（大語(yǔ)言模型）因其強(qiáng)大的語(yǔ)言理解能力贏得了眾多用戶的青睞，但LLM龐大規(guī)模的參數(shù)導(dǎo)致其部署條件苛刻；在網(wǎng)絡(luò)受限，計(jì)算資源有限的場(chǎng)景下無(wú)法使用大語(yǔ)言模型的能力；低算力，本地化部署的問(wèn)題亟待解決。ChatGLM-6B在60億參數(shù)的情況下做到了優(yōu)秀的中英文對(duì)話效果，且能夠支持在消費(fèi)級(jí)顯卡本地部署；因此在HuggingFace Trends上很快登頂。6B的參數(shù)量雖然能夠做到本地部署，但是目前的實(shí)現(xiàn)依賴庫(kù)較多，如Pytorch, transfomer；對(duì)于端側(cè)部署來(lái)說(shuō)要求仍然較高。因此我們嘗試將該模型轉(zhuǎn)換為MNN模型，極大降低了部署時(shí)的依賴項(xiàng)，能夠更方便的在各類端側(cè)設(shè)備上部署與測(cè)試；同時(shí)我們對(duì)MNN模型進(jìn)行了低bit量化，并實(shí)現(xiàn)了反量化與計(jì)算融合的計(jì)算kernel，大大降低了內(nèi)存需求。實(shí)測(cè)PC端小顯存顯卡能夠成流暢運(yùn)行浮點(diǎn)模型，在Android手機(jī)上能夠流暢運(yùn)行量化模型。

模型導(dǎo)出

模型導(dǎo)出采用了Pytorch到ONNX到MNN的轉(zhuǎn)換方式，并切對(duì)模型進(jìn)行了拆分導(dǎo)出，將embedding，28 x GLMBlock， lm分別導(dǎo)出；并且在導(dǎo)出時(shí)對(duì)詞表進(jìn)行了瘦身。模型導(dǎo)出代碼。

?導(dǎo)出方式

Pytorch實(shí)現(xiàn)的模型導(dǎo)出目前有2種主流方案：1. 導(dǎo)出為ONNX; 2. 導(dǎo)出為TorchScript。分析代碼后可知，ChatGLM的模型結(jié)構(gòu)比較簡(jiǎn)單，Embedding層，28層GLMBlock，線性層；其中GLMBlock結(jié)構(gòu)如為L(zhǎng)ayerNorm -> 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_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

因?yàn)樵撃Ｐ徒Y(jié)構(gòu)簡(jiǎn)單，使用的算子 ONNX全部支持；同時(shí)MNN對(duì)ONNX的支持完備性比較好；因此選擇使用ONNX導(dǎo)出模型。

?結(jié)構(gòu)拆分

在確定使用ONNX之后首先直接使用torch.onnx.export嘗試對(duì)模型進(jìn)行導(dǎo)出，導(dǎo)出過(guò)程非常緩慢，導(dǎo)出后模型的權(quán)重大小有28G。在將模型轉(zhuǎn)換到MNN時(shí)會(huì)執(zhí)行一些圖優(yōu)化Pass；因?yàn)槟Ｐ吞髮?dǎo)致占用內(nèi)存過(guò)高速度非常慢；因此考慮將模型進(jìn)行拆分優(yōu)化。拆分之后的優(yōu)化考慮如下：

Embedding層的參數(shù)大小為150528 * 4096, 單個(gè)權(quán)重使用內(nèi)存非常大；考慮到輸入文字?jǐn)?shù)量較?。ㄏ鄬?duì)于150528），使用Gather實(shí)現(xiàn)消耗大量?jī)?nèi)存/顯存，直接將參數(shù)存儲(chǔ)為二進(jìn)制文件，通過(guò)fseekfread實(shí)現(xiàn)Gather的操作能夠在稍微犧牲速度的情況下節(jié)約2.3G內(nèi)存；同時(shí)為了降低模型的文件大小，將embedding層的數(shù)據(jù)使用bf16格式存儲(chǔ)，能夠?qū)⑽募笮〗档鸵话?，?duì)精度和性能形象非常小。

GLMBlock層的權(quán)重總大小為21G，仍然非常大，每個(gè)Block的大小為768M；考慮到要在端側(cè)各類設(shè)備部署，可以將28層Block分別導(dǎo)出，對(duì)于浮點(diǎn)模型，這樣的好處是能夠在顯存不足的情況下將部分Block放置在GPU，其余部分放置在CPU進(jìn)行推理，這樣能夠充分利用設(shè)備算力；對(duì)與移動(dòng)端設(shè)備，對(duì)這些block進(jìn)行量化，分別獲得int8/int4的權(quán)值量化模型，使用int4量化模型大小為2.6G，可以在端側(cè)小內(nèi)存設(shè)備部署。

線性層通過(guò)一個(gè)矩陣乘將hidden_state轉(zhuǎn)換為詞語(yǔ)的prob：[num, 4096] @ [4096, 150528]；其實(shí)這里并不需要num全部參與運(yùn)算，比如輸入序列長(zhǎng)度num = 10時(shí)，實(shí)際預(yù)測(cè)下一個(gè)詞語(yǔ)時(shí)進(jìn)需要使用最后一個(gè)[1, 4096]即可。因此可以先對(duì)輸入變量做一個(gè)Gather然后執(zhí)行矩陣乘：[1, 4096] @ [4096, 150528]即可。為了在端側(cè)降低內(nèi)存占用，這里同樣使用int8/int4量化，量化后大小為256M。

?詞表瘦身

詞表大小為150528, 分析發(fā)現(xiàn)前20000個(gè)詞語(yǔ)為，在Chat中并沒(méi)有使用，因此將可以將結(jié)構(gòu)拆分后的Embedding層和最后的線性層進(jìn)行刪減。簡(jiǎn)單的方法是將2層的權(quán)重導(dǎo)出onnx模型，使用numpy.fromfile將onnx模型的權(quán)重加載，刪除前[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')

對(duì)于刪減后的詞表，使用bf16格式存儲(chǔ)可以降低一半的文件大小，使用C++代碼將fp32轉(zhuǎn)換為bf16，如下：

// read binary file
FILE* 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 bf16
std::vector dst_buffer(num);
for (int i = 0; i < num; i++) {
    dst_buffer[i] = reinterpret_cast(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);

?動(dòng)態(tài)形狀

因?yàn)槟Ｐ洼斎氲男螤钍莿?dòng)態(tài)變化的，因此需要在導(dǎo)出時(shí)指定動(dòng)態(tài)形狀的維度，具體的導(dǎo)出方式如下：

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)

?其他問(wèn)題

Tuple改為Tensor

原始實(shí)現(xiàn)中l(wèi)ayer_past是Tuple，將其修改為Tensor方便模型導(dǎo)出后的模型輸入。將代碼中的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操作不支持動(dòng)態(tài)形狀

指定了動(dòng)態(tài)維度后，在實(shí)際測(cè)試中發(fā)現(xiàn)因?yàn)槟Ｐ蛯?shí)現(xiàn)中有些view相關(guān)代碼導(dǎo)出后會(huì)將形狀固定為常量，導(dǎo)致導(dǎo)出后改變輸入形狀無(wú)法正確推理，因此需要對(duì)模型中非動(dòng)態(tài)的實(shí)現(xiàn)進(jìn)行修改，將attention_fn函數(shù)中所有view操作替換為squeeze和unsqueeze操作，這樣導(dǎo)出后與形狀無(wú)關(guān)即可實(shí)現(xiàn)動(dòng)態(tài)形狀。

# 修改前
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簡(jiǎn)化

squeeze(1)在模型導(dǎo)出時(shí)會(huì)產(chǎn)生了額外If算子，apply_rotary_pos_emb_index函數(shù)中使用squeeze(1)會(huì)使模型中多出2個(gè)If，為了讓模型更加簡(jiǎn)單，這里可以直接替換為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個(gè)詞語(yǔ)，在前處理轉(zhuǎn)換為n個(gè)int后，通過(guò)對(duì)embedding數(shù)據(jù)的查詢會(huì)生成一個(gè)[n, 4096]的向量，同時(shí)根據(jù)輸入長(zhǎng)度和結(jié)束符位置，生成position_ids和mask作為輸入；對(duì)于非首次生成還會(huì)有一個(gè)歷史信息的輸入。其中position_ids和mask對(duì)于28個(gè)block完全一樣，history每個(gè)block都使用上次生成時(shí)對(duì)應(yīng)block的present輸出；而輸入的input_embedding則使用上一個(gè)block的輸出hidden_state，具體結(jié)構(gòu)如下：

?前后處理

官方提供的實(shí)現(xiàn)中使用了transformers庫(kù)，該庫(kù)提供了模型的前后處理的實(shí)現(xiàn)。其中前處理包括了分詞，將詞語(yǔ)轉(zhuǎn)換為ids；后處理中包含了prob轉(zhuǎn)換為詞語(yǔ)，控制模型持續(xù)生成的邏輯。在轉(zhuǎn)換到C++之后我們也需要實(shí)現(xiàn)相同的前后處理邏輯。

前處理

前處理邏輯是將用戶輸入的句子進(jìn)行分詞，然后查詢?cè)~表將詞語(yǔ)轉(zhuǎn)換為id；C++中實(shí)現(xiàn)如下：

分詞：在C++上使用cppjieba進(jìn)行分詞；

word2id: 將詞表文件加載為map，通過(guò)查詢map將詞語(yǔ)轉(zhuǎn)換為id;

std::vector ids;
std::vector 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轉(zhuǎn)換為id，然后通過(guò)詞表將id轉(zhuǎn)換為詞語(yǔ)，同時(shí)將一些特殊字符進(jìn)行轉(zhuǎn)義；C++中實(shí)現(xiàn)如下：

prob2id：在lm層后接一個(gè)ArgMax即可將prob轉(zhuǎn)換為id，實(shí)測(cè)效果與transformers中的實(shí)現(xiàn)結(jié)果一致；

id2word: 次變文件加載為vector，直接讀取即可獲取word；

特殊詞處理：針對(duì)一些特殊詞語(yǔ)進(jìn)行了替換；

auto word = mWordDecode[id];
if (word == "") return "
";
if (word == "<|tab|>") return "	";
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依次執(zhí)行推理，將其中的hidden_state輸出作為下一個(gè)block的輸入；并將present輸出保存起來(lái)作為下次推理的history即可，代碼如下：

//embedding
FILE* 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 < seq_len; i++) {
    fseek(file, input_ids[i] * size, SEEK_SET);
    fread(embedding_var->writeMap() + 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()[0];

推理優(yōu)化

?MNN Module接口

MNN的Module接口相比于Session接口，能夠支持控制流，輸入輸出不需要再考慮設(shè)備問(wèn)題，用戶可以透明的使用GPU，CPU等不同設(shè)備串聯(lián)推理且不需要關(guān)心數(shù)據(jù)的設(shè)備問(wèn)題。同時(shí)Module接口在模型加載時(shí)可以對(duì)權(quán)重進(jìn)行預(yù)重排，在內(nèi)存充足的情況下提升卷積，矩陣乘等算子的推理速度。在對(duì)模型逐層拆分后會(huì)考慮到不同設(shè)備加載，Module接口可以更加簡(jiǎn)潔的實(shí)現(xiàn)這種跨設(shè)備的推理。

?PC端低顯存推理（浮點(diǎn)）

上述轉(zhuǎn)換與推理使用的模型都是浮點(diǎn)模型，在實(shí)際推理中可以選擇fp32或者fp16。在使用fp16推理時(shí)，顯存要求在13G以上；目前主流的游戲顯卡顯存普遍達(dá)不到該要求，因此無(wú)法將全部模型加載到GPU中推理?？紤]到我們對(duì)模型進(jìn)行了分段劃分，可以將一部分block放入顯存使用GPU推理，剩余部分使用CPU推理。因此可以根據(jù)用戶指定的顯存大小動(dòng)態(tài)的分配block到GPU中。分配規(guī)則為，fp16的情況下每個(gè)block占用顯存大小為385M，推理過(guò)程中的特征向量大小預(yù)留2G的顯存，因此可以加載到GPU中的層數(shù)為：(gpu_memory - 2) * 1024.0 / 385.0。代碼實(shí)現(xiàn)如下：

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

?移動(dòng)端低內(nèi)存推理（量化）

全部浮點(diǎn)模型加載使用CPU推理需要32G左右的內(nèi)存大小，在移動(dòng)設(shè)備上很難滿足內(nèi)存要求。使用模型量化的方法來(lái)降低內(nèi)存占用是常用的方法；MNN支持模型權(quán)重量化，ChatGLM在int4量化后總權(quán)重大小在4G左右，理論上是可以在移動(dòng)端全部加載運(yùn)行的。但是由于MNN之前的設(shè)計(jì)偏向于中小模型的推理，在模型加載階段就做了反量化的計(jì)算，導(dǎo)致實(shí)際推理的內(nèi)存占用與浮點(diǎn)一致，沒(méi)有降低推理時(shí)內(nèi)存占用。這種模式針對(duì)傳統(tǒng)的端側(cè)模型是非常合適的，在降低了模型大小的同時(shí)不會(huì)有性能損失；但是遇到內(nèi)存總大小成為瓶頸大模型時(shí)則不是非常合適；因此需要針對(duì)大模型的大內(nèi)存需求進(jìn)行優(yōu)化。針對(duì)大模型內(nèi)存瓶頸問(wèn)題，MNN在運(yùn)行時(shí)使用low_memory選項(xiàng)會(huì)將反量化過(guò)程放在矩陣乘中實(shí)現(xiàn)，以部分推理時(shí)的額外計(jì)算開銷大幅降低內(nèi)存占用與訪存帶寬占用。

對(duì)于權(quán)值量化模型的低內(nèi)存實(shí)現(xiàn)，我們支持了int4和int8兩種權(quán)值量化的模型的低內(nèi)存模式。針對(duì)不同的硬件做了實(shí)現(xiàn)，針對(duì)X86 SSE, AVX2實(shí)現(xiàn)了int4@fp32, int8@fp32；針對(duì)ARM64實(shí)現(xiàn)了int4@fp32, int8@fp32和int4@fp16和int8@fp16。具體的是線上需要針對(duì)以上列舉的情況分別實(shí)現(xiàn)對(duì)應(yīng)的矩陣乘Kernel，并且在原來(lái)的浮點(diǎn)矩陣乘的輸入里增加反量化需要的alpha和bias參數(shù)，在矩陣乘計(jì)算前需要先從內(nèi)存中加載常量的int4/int8量化值，然后將其轉(zhuǎn)換為浮點(diǎn)類型，之后再執(zhí)行浮點(diǎn)矩陣乘操作，實(shí)際的矩陣乘基礎(chǔ)操作如下公式：

?

以下為int4量化模型在ARMv8.2上的fp16矩陣乘實(shí)現(xiàn)：

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

PC端測(cè)試：11G顯存的2080Ti + AMD 3900X + 32G內(nèi)存測(cè)；使用fp32精度模型（GPU顯存不足情況下）GPU+CPU混合速度為3.5 tok/s; 僅使用CPU速度為 1.2 tok/s ；

移動(dòng)端測(cè)試：Xiaomi12；使用int4模型精度，CPU速度為 1.5 tok/s，需要內(nèi)存為 2.9 G。

用戶界面

?PC

在PC端提供了兩種demo的用法，命令行與web的用戶界面；

?移動(dòng)端

在移動(dòng)端目前提供了Android的App；

總結(jié)

ChatGLM-6B模型推理主要是內(nèi)存瓶頸，針對(duì)這一問(wèn)題本文提出了2中優(yōu)化手段：分段加載；量化。使用了這兩種優(yōu)化方法后，我們能夠保證模型精度完全無(wú)損（fp32）的情況下在消費(fèi)級(jí)顯卡的PC機(jī)器上部署運(yùn)行；同時(shí)還支持略微損失模型精度（int4）的情況下在移動(dòng)端設(shè)備上流暢運(yùn)行。

審核編輯：劉清