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

模型導出

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

?導出方式

Pytorch實現的模型導出目前有2種主流方案：1. 導出為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.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

因為該模型結構簡單，使用的算子 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。

?詞表瘦身

詞表大小為150528, 分析發現前20000個詞語為，在Chat中并沒有使用，因此將可以將結構拆分后的Embedding層和最后的線性層進行刪減。簡單的方法是將2層的權重導出onnx模型，使用numpy.fromfile將onnx模型的權重加載，刪除前[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轉換為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);

?動態形狀

因為模型輸入的形狀是動態變化的，因此需要在導出時指定動態形狀的維度，具體的導出方式如下：

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

原始實現中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簡化

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

?前后處理

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

前處理

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

分詞：在C++上使用cppjieba進行分詞；

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

prob2id：在lm層后接一個ArgMax即可將prob轉換為id，實測效果與transformers中的實現結果一致；

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

特殊詞處理：針對一些特殊詞語進行了替換；

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依次執行推理，將其中的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];

推理優化

?MNN Module接口

MNN的Module接口相比于Session接口，能夠支持控制流，輸入輸出不需要再考慮設備問題，用戶可以透明的使用GPU，CPU等不同設備串聯推理且不需要關心數據的設備問題。同時Module接口在模型加載時可以對權重進行預重排，在內存充足的情況下提升卷積，矩陣乘等算子的推理速度。在對模型逐層拆分后會考慮到不同設備加載，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 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);
}

?移動端低內存推理（量化）

全部浮點模型加載使用CPU推理需要32G左右的內存大小，在移動設備上很難滿足內存要求。使用模型量化的方法來降低內存占用是常用的方法；MNN支持模型權重量化，ChatGLM在int4量化后總權重大小在4G左右，理論上是可以在移動端全部加載運行的。但是由于MNN之前的設計偏向于中小模型的推理，在模型加載階段就做了反量化的計算，導致實際推理的內存占用與浮點一致，沒有降低推理時內存占用。這種模式針對傳統的端側模型是非常合適的，在降低了模型大小的同時不會有性能損失；但是遇到內存總大小成為瓶頸大模型時則不是非常合適；因此需要針對大模型的大內存需求進行優化。針對大模型內存瓶頸問題，MNN在運行時使用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量化值，然后將其轉換為浮點類型，之后再執行浮點矩陣乘操作，實際的矩陣乘基礎操作如下公式：

?

以下為int4量化模型在ARMv8.2上的fp16矩陣乘實現：

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內存測；使用fp32精度模型（GPU顯存不足情況下）GPU+CPU混合速度為3.5 tok/s; 僅使用CPU速度為 1.2 tok/s ；

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

用戶界面

?PC

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

?移動端

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

總結

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

審核編輯：劉清