RM新时代|官方理财平台

算豐技術(shù)揭秘｜探索ChatGLM2-6B模型與TPU部署

概述

ChatGLM2-6B是由清華大學(xué)開(kāi)發(fā)的中英雙語(yǔ)對話(huà)模型，它能生成相當符合人類(lèi)偏好的回答，使每個(gè)人都能擁有自己的深度學(xué)習大模型，為了更好地支持開(kāi)源生態(tài)，ChatGLM2-6B對學(xué)術(shù)研究完全開(kāi)放，并允許免費商用。

算能BM1684X已經(jīng)實(shí)現ChatGLM2-6B的C++代碼部署，代碼實(shí)現鏈接：https://github.com/sophgo/ChatGLM2-TPU

本文總結了部署該模型過(guò)程中的一些技術(shù)點(diǎn)，希望通過(guò)詳細的教程幫助到更多的開(kāi)發(fā)者。文章內容包括ChatGLM2-6B的整體運行流程；如何將該動(dòng)態(tài)網(wǎng)路轉換成靜態(tài)網(wǎng)絡(luò )形式；如何將該模型導出成ONNX；如何將ONNX使用TPU-MLIR編譯器實(shí)現網(wǎng)絡(luò )的編譯；如何用C++代碼編寫(xiě)應用程序。

ChatGLM2-6b運行流程

?

如圖該網(wǎng)絡(luò )基本上可以分為5個(gè)階段：

  1. 將句子通過(guò)分詞器(使用google的sentencepiece）轉換成tokens，如圖中的<1x17 xi32>的數據。注意tokens里面64790, 64792是起始符號。
  2. 通過(guò)WordEmbedding將tokens轉換成詞向量，得到的結果是<1x17x4096 xf32>的數據。
  3. 通過(guò)Tranformer進(jìn)行神經(jīng)網(wǎng)絡(luò )推理，推理結果是<1x17x4096 xf32>，答案在最后的詞向量上，所以做一個(gè)額外的slice操作，得到<1x1x4096 xf32>。這里Transformer網(wǎng)絡(luò )是由28個(gè)Block組成，每個(gè)Block的核心是一個(gè)Attention運算，并輸出kv cache給下一輪Transform做為輸入。
  4. 經(jīng)過(guò)LmHead操作生成<1x1 xi32>的結果，也就是輸出的Token。LmHead的組成如圖所示。
  5. Token經(jīng)過(guò)分詞器轉成詞語(yǔ)，且傳遞給下一輪推理，進(jìn)入第一階段。直到token == EOS_ID結束。

轉成靜態(tài)網(wǎng)絡(luò )

ChatGLM2-6B從前面的描述中，可以看到有兩處是動(dòng)態(tài)的，一是因句子的長(cháng)短不同，Transformer的輸入Shape有所有不同；二是每一輪Transformer生成的kv cache會(huì )逐步增長(cháng)。為了方便部署，根據網(wǎng)絡(luò )特點(diǎn)轉換成靜態(tài)網(wǎng)絡(luò )。轉換后的運行流程如下：

從圖中可以看到句子不論長(cháng)短，轉換后的tokens始終是<1x512x i32>，kv cache的數據也始終是<512x1x2x128x f32>。

這里介紹最關(guān)鍵的幾點(diǎn)：

  1. 將原始tokens輸入尾部補0，從<1x17x i32>轉換成<1x512x i32>。
  2. 將position_ids從GlmBlock中提取出來(lái)，并固定長(cháng)度為<1x512x i32>，也是尾部補0，本例中的數值為[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,0,0,0,0,...0]，用于位置編碼。因為在原始網(wǎng)絡(luò )中這部分是變長(cháng)，提取出來(lái)后就可以做成定長(cháng)。
  3. 將attention_mask從GlmBlock中提取出來(lái)，并固定長(cháng)度為<1x512x512x f32>，注意無(wú)效部分全部補1，因為它之后會(huì )有masked_fill操作將mask為1的部分全部配置為-inf。然后經(jīng)過(guò)Softmax使超出部分全部清0，保留有效部分，從而保證最終結果與原始結果一致。如下圖，為說(shuō)明問(wèn)題，Attention做了簡(jiǎn)化。
  4. 第一輪Transformer后，kv_chache有效部分是[0:17]，我們將該部分移到末尾[512-17:]，并頭部清0。因為kv cache的累加發(fā)生在尾部。從第二輪開(kāi)始累加后做Slice操作去掉頭部1個(gè)單位，取[1:]，這樣就保證了kv cache始終保持在512。同時(shí)attention mask也要是從尾部實(shí)際token len長(cháng)度的0，頭部全部置1。

導出ONNX

將該網(wǎng)絡(luò )分為4塊：WorkEmbedding，GlmBlock，GlmBlockCache，LmHead。這里分別介紹這四塊是如何導出的。

導出前，先要指定python路徑，如下：

export?PYTHONPATH=/workspace/chatglm2-6b:$PYTHONPATH

需要先加載原始ChatGLM2-6B，如下代碼：

CHATGLM2_PATH?=?"/workspace/chatglm2-6b"
origin_model?=?AutoModel.from_pretrained(CHATGLM2_PATH,
?????????????????????????????????????????trust_remote_code=True).float()
origin_model.eval()
transformer?=?origin_model.transformer
MAX_LEN?=?transformer.seq_length
for?param?in?origin_model.parameters():
????param.requires_grad?=?False
num_layers?=?transformer.encoder.num_layers
layers?=?transformer.encoder.layers

WordEmbedding

直接使用原模型中的word_embeddings，構建成獨立網(wǎng)絡(luò )，導出即可

class?Embedding(torch.nn.Module):
????def?__init__(self):
????????super().__init__()
????def?forward(self,?input_ids):
????????return?transformer.embedding.word_embeddings(input_ids)
def?convert_embedding():
????model?=?Embedding()
????torch.onnx.export(model,?(torch.tensor([0,?1,?2,?3])),
??????????????????????f'./tmp/embedding.onnx',
??????????????????????verbose=False,
??????????????????????input_names=['input_ids'],
??????????????????????output_names=['input_embed'],
??????????????????????dynamic_axes={"input_ids":?{0:?"length"}},
??????????????????????do_constant_folding=True,
??????????????????????opset_version=15)

GlmBlock

需要將transformer.rotary_pos_emb和transformer.encoder.layers組合構建獨立網(wǎng)路，導出。因為有28個(gè)block，所以需要導出28個(gè)ONNX模型。這里的position_ids和attention_mask作為外部輸入，前面有介紹。其實(shí)這28個(gè)Block是可以組合成一個(gè)模型，但是這樣導致onnx權重過(guò)大(F16約12GB)，導出麻煩，部署也麻煩，所以單個(gè)導出。

class?GlmBlock(torch.nn.Module):
????def?__init__(self,?layer_id):
????????super().__init__()
????????self.layer?=?layers[layer_id]
????def?forward(self,?hidden_states,?position_ids,?attention_mask):
????????rotary_pos_emb?=?transformer.rotary_pos_emb(MAX_LEN)[position_ids]
????????rotary_pos_emb?=?rotary_pos_emb.transpose(0,?1).contiguous()
????????hidden_states,?past_kv?=?self.layer(hidden_states,?attention_mask,
????????????????????????????????????????????rotary_pos_emb=rotary_pos_emb)
????????return?hidden_states,?past_kv
def?convert_glm_block(layer_id):
????model?=?GlmBlock(layer_id)
????torch.onnx.export(
????????model,?(hidden_states,?position_ids,?attention_mask),
????????f'./tmp/glm_block_{layer_id}.onnx',
????????verbose=False,
????????input_names=['input_states',?'position_ids',?'attention_mask'],
????????output_names=['hidden_states',?'past_k',?'past_v'],
????????do_constant_folding=True,
????????opset_version=15)

GlmBlockCache

與`GlmBlock是類(lèi)似的，但是需要額外的kv cache參數。注意這里 最后會(huì )把頭部1個(gè)單位去除掉。

class?GlmBlockCache(torch.nn.Module):
????def?__init__(self,?layer_id):
????????super().__init__()
????????self.layer?=?layers[layer_id]
????def?forward(self,?hidden_states,?position_ids,?attention_mask,?past_k,?past_v):
????????rotary_pos_emb?=?transformer.rotary_pos_emb(MAX_LEN)[position_ids]
????????rotary_pos_emb?=?rotary_pos_emb.transpose(0,?1).contiguous()
????????hidden_states,?past_kv?=?self.layer(hidden_states,?attention_mask,
????????????????????????????????????????????kv_cache=(past_k,?past_v),
????????????????????????????????????????????rotary_pos_emb=rotary_pos_emb)
????????past_k,?past_v?=?past_kv
????????return?hidden_states,?past_k[1:],?past_v[1:]
def?convert_glm_block_cache(layer_id):
????model?=?GlmBlockCache(layer_id)
????torch.onnx.export(
????????model,?(hidden_states,?position_ids,?attention_mask,?past_k,?past_v),
????????f'./tmp/glm_block_cache_{layer_id}.onnx',
????????verbose=False,
????????input_names=['input_states',?'position_ids',?'attention_mask',?'history_k',?'history_v'],
????????output_names=['hidden_states',?'past_k',?'past_v'],
????????do_constant_folding=True,
????????opset_version=15)

LmHead

這里取m_logits后使用topk，其實(shí)也是可以用argmax，看處理器實(shí)現哪一種效率高。

class?LmHead(torch.nn.Module):
????def?__init__(self):
????????super().__init__()
????def?forward(self,?hidden_states):
????????hidden_states?=?transformer.encoder.final_layernorm(hidden_states)
????????m_logits?=?transformer.output_layer(hidden_states)
????????_,?token?=?torch.topk(m_logits,?1)
????????return?token
def?convert_lm_head():
????model?=?LmHead()
????input?=?torch.randn(1,?4096)
????torch.onnx.export(model,?(input),?f'./tmp/lm_head.onnx',?verbose=False,
??????????????????????input_names=['hidden_states'],
??????????????????????output_names=['token'],
??????????????????????do_constant_folding=True,
??????????????????????opset_version=15)

部署

上述轉完ONNX模型后都已經(jīng)是靜態(tài)網(wǎng)絡(luò )，通過(guò)TPU-MLIR，可以很容易的轉換成F16的模型。但是特別要注意的是RmsNorm需要用F32。之后就可以按照執行邏輯編寫(xiě)C++代碼。演示效果如下：

?

分享到 ：

最新資訊

RM新时代|官方理财平台