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

模型導(dǎo)出

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

?導(dǎo)出方式

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

因為該模型結(jié)構(gòu)簡單，使用的算子 ONNX全部支持；同時MNN對ONNX的支持完備性比較好；因此選擇使用ONNX導(dǎo)出模型。

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

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

Embedding層的參數(shù)大小為150528 * 4096, 單個權(quán)重使用內(nèi)存非常大；考慮到輸入文字?jǐn)?shù)量較小（相對于150528），使用Gather實現(xiàn)消耗大量內(nèi)存/顯存，直接將參數(shù)存儲為二進(jìn)制文件，通過fseekfread實現(xiàn)Gather的操作能夠在稍微犧牲速度的情況下節(jié)約2.3G內(nèi)存；同時為了降低模型的文件大小，將embedding層的數(shù)據(jù)使用bf16格式存儲，能夠?qū)⑽募笮〗档鸵话耄瑢群托阅苄蜗蠓浅Ｐ　?/p>

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

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

?詞表瘦身

詞表大小為150528, 分析發(fā)現(xiàn)前20000個詞語為，在Chat中并沒有使用，因此將可以將結(jié)構(gòu)拆分后的Embedding層和最后的線性層進(jì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')

對于刪減后的詞表，使用bf16格式存儲可以降低一半的文件大小，使用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);

?動態(tài)形狀

因為模型輸入的形狀是動態(tài)變化的，因此需要在導(dǎo)出時指定動態(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)

?其他問題

Tuple改為Tensor

原始實現(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操作不支持動態(tài)形狀

指定了動態(tài)維度后，在實際測試中發(fā)現(xiàn)因為模型實現(xiàn)中有些view相關(guān)代碼導(dǎo)出后會將形狀固定為常量，導(dǎo)致導(dǎo)出后改變輸入形狀無法正確推理，因此需要對模型中非動態(tài)的實現(xiàn)進(jìn)行修改，將attention_fn函數(shù)中所有view操作替換為squeeze和unsqueeze操作，這樣導(dǎo)出后與形狀無關(guān)即可實現(xiàn)動態(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簡化

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

?前后處理

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

前處理

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

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

word2id: 將詞表文件加載為map，通過查詢map將詞語轉(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，然后通過詞表將id轉(zhuǎn)換為詞語，同時將一些特殊字符進(jìn)行轉(zhuǎn)義；C++中實現(xiàn)如下：

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

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

特殊詞處理：針對一些特殊詞語進(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輸出作為下一個block的輸入；并將present輸出保存起來作為下次推理的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è)備問題，用戶可以透明的使用GPU，CPU等不同設(shè)備串聯(lián)推理且不需要關(guān)心數(shù)據(jù)的設(shè)備問題。同時Module接口在模型加載時可以對權(quán)重進(jìn)行預(yù)重排，在內(nèi)存充足的情況下提升卷積，矩陣乘等算子的推理速度。在對模型逐層拆分后會考慮到不同設(shè)備加載，Module接口可以更加簡潔的實現(xiàn)這種跨設(shè)備的推理。

?PC端低顯存推理（浮點）

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

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

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

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

?

以下為int4量化模型在ARMv8.2上的fp16矩陣乘實現(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端測試：11G顯存的2080Ti + AMD 3900X + 32G內(nèi)存測；使用fp32精度模型（GPU顯存不足情況下）GPU+CPU混合速度為3.5 tok/s; 僅使用CPU速度為 1.2 tok/s ；

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

用戶界面

?PC

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

?移動端

在移動端目前提供了Android的App；

總結(jié)

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

審核編輯：劉清