ChatGLM2-6B解析与TPU部署
概述#
算能bm1684x芯片已经实现chatglm2-6b的c++代码部署，代码实现链接：https://github.com/sophgo/chatglm2-tpu。
本文总结部署该模型过程中的一些技术点。首先是chatglm2-6b的整体运行流程，然后介绍如何将该动态网路转换成静态网络形式，接着介绍如何将该模型导出成onnx。
最后如何将onnx使用tpu-mlir编译器实现网络的编译以及用c++代码编写应用程序，可以直接看源码就可以理解，这里不做介绍。
chatglm2-6b运行流程#
如图该网络基本上可以分为5个阶段：
将句子通过分词器(使用google的sentencepiece）转换成tokens，如图中的的数据。注意tokens里面64790, 64792是起始符号。
通过wordembedding将tokens转换成词向量，得到的结果是的数据。
通过tranformer进行神经网络推理，推理结果是，答案在最后的词向量上，所以做一个额外的slice操作，得到。这里transformer网络是由28个block组成，每个block的核心是一个attention运算，并输出kv cache给下一轮transform做为输入。
经过lmhead操作生成的结果，也就是输出的token。lmhead的组成如图所示。
token经过分词器转成词语，且传递给下一轮推理，进入第一阶段。直到token == eos_id结束。
转成静态网络#
chatglm2-6b从前面的描述中，可以看到有两处是动态的，一是因句子的长短不同，transformer的输入shape有所有不同；二是每一轮transformer生成的kv cache会逐步增长。为了方便部署，根据网络特点转换成静态网络。转换后的运行流程如下：
从图中可以看到句子不论长短，转换后的tokens始终是，kv cache的数据也始终是。
这里介绍最关键的几点：
将原始tokens输入尾部补0，从转换成。
将position_ids从glmblock中提取出来，并固定长度为，也是尾部补0，本例中的数值为[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,0,0,0,0,...0]，用于位置编码。因为在原始网络中这部分是变长，提取出来后就可以做成定长。
将attention_mask从glmblock中提取出来，并固定长度为，注意无效部分全部补1，因为它之后会有masked_fill操作将mask为1的部分全部配置为-inf。然后经过softmax使超出部分全部清0，保留有效部分，从而保证最终结果与原始结果一致。如下图，为说明问题，attention做了简化。
第一轮transformer后，kv_chache有效部分是[0:17]，我们将该部分移到末尾[512-17:]，并头部清0。因为kv cache的累加发生在尾部。从第二轮开始累加后做slice操作去掉头部1个单位，取[1:]，这样就保证了kv cache始终保持在512。同时attention mask也要是从尾部实际token len长度的0，头部全部置1。
导出onnx#
将该网络分为4块：workembedding，glmblock，glmblockcache，lmhead。这里分别介绍这四块是如何导出的。
导出前，先要指定python路径，如下：
1 export pythonpath=/workspace/chatglm2-6b:$pythonpath
需要先加载原始chatglm2-6b，如下代码：
1234567891011 chatglm2_path = /workspace/chatglm2-6borigin_model = automodel.from_pretrained(chatglm2_path, trust_remote_code=true).float()origin_model.eval()transformer = origin_model.transformermax_len = transformer.seq_lengthfor param in origin_model.parameters(): param.requires_grad = falsenum_layers = transformer.encoder.num_layerslayers = transformer.encoder.layers
workembedding#
直接使用原模型中的word_embeddings，构建成独立网络，导出即可
123456789101112131415161718 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组合构建独立网路，导出。因为有28个block，所以需要导出28个onnx模型。这里的position_ids和attention_mask作为外部输入，前面有介绍。其实这28个block是可以组合成一个模型，但是这样导致onnx权重过大(f16约12gb)，导出麻烦，部署也麻烦，所以单个导出。
1234567891011121314151617181920212223 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_kvdef 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是类似的，但是需要额外的kv cache参数。注意这里 最后会把头部1个单位去除掉。
123456789101112131415161718192021222324 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，其实也是可以用argmax，看芯片实现哪一种效率高。
12345678910111213141516 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 tokendef 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模型后都已经是静态网络，通过tpu-mlir，可以很容易的转换成f16的模型。但是特别要注意的是rmsnorm需要用f32。之后就可以按照执行逻辑编写c++代码。演示效果如下：
双通道数字调节可变增益放大器MAX2063
abb机器人将向客户免费开放关键的软件服务
中美5G形成打劫局面 技术“硬核”即是先手
搭载骁龙835的小米6到底有多强？GPU频率峰值超过800Mhz官方都没摸透？
立景创新科技联合华为打造高品质制造园区，柔性生产、高效运营
ChatGLM2-6B解析与TPU部署
隆利科技荣获“Mini/Micro LED产业推动先锋奖”奖项
石英射灯电路的检修
日本补贴4760亿日元，台积电只需为熊本县工厂出资60%
微软官方对新主机命名作出解释
Nature:同行评审屡遭毒喷,中国欲启用AI系统解决问题
LLC串联谐振拓扑谐振腔电压概述
新兴存储器铁电RAM嵌入式应用的优势是什么
AR“小巨人”｜亮风台继专精特新“小巨人”后，再获上海市“科技小巨人”
高通可向荣耀供货5G芯片，无需审批
动力电池新蓝海 装备巨头入局叠片领域
华为首款全面屏新机9月22日发布,不是华为Mate10也不是荣耀Note9!而是麦芒6,惊不惊喜意不意外
如何快速成为点灯大师？
自恢复保险丝(PPTC)规避的使用四要素
科里奥利流量计的工作原理及设计
算能bm1684x芯片已经实现chatglm2-6b的c++代码部署，代码实现链接：https://github.com/sophgo/chatglm2-tpu。
本文总结部署该模型过程中的一些技术点。首先是chatglm2-6b的整体运行流程，然后介绍如何将该动态网路转换成静态网络形式，接着介绍如何将该模型导出成onnx。
最后如何将onnx使用tpu-mlir编译器实现网络的编译以及用c++代码编写应用程序，可以直接看源码就可以理解，这里不做介绍。
chatglm2-6b运行流程#
如图该网络基本上可以分为5个阶段：
将句子通过分词器(使用google的sentencepiece）转换成tokens，如图中的的数据。注意tokens里面64790, 64792是起始符号。
通过wordembedding将tokens转换成词向量，得到的结果是的数据。
通过tranformer进行神经网络推理，推理结果是，答案在最后的词向量上，所以做一个额外的slice操作，得到。这里transformer网络是由28个block组成，每个block的核心是一个attention运算，并输出kv cache给下一轮transform做为输入。
经过lmhead操作生成的结果，也就是输出的token。lmhead的组成如图所示。
token经过分词器转成词语，且传递给下一轮推理，进入第一阶段。直到token == eos_id结束。
转成静态网络#
chatglm2-6b从前面的描述中，可以看到有两处是动态的，一是因句子的长短不同，transformer的输入shape有所有不同；二是每一轮transformer生成的kv cache会逐步增长。为了方便部署，根据网络特点转换成静态网络。转换后的运行流程如下：
从图中可以看到句子不论长短，转换后的tokens始终是，kv cache的数据也始终是。
这里介绍最关键的几点：
将原始tokens输入尾部补0，从转换成。
将position_ids从glmblock中提取出来，并固定长度为，也是尾部补0，本例中的数值为[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,0,0,0,0,...0]，用于位置编码。因为在原始网络中这部分是变长，提取出来后就可以做成定长。
将attention_mask从glmblock中提取出来，并固定长度为，注意无效部分全部补1，因为它之后会有masked_fill操作将mask为1的部分全部配置为-inf。然后经过softmax使超出部分全部清0，保留有效部分，从而保证最终结果与原始结果一致。如下图，为说明问题，attention做了简化。
第一轮transformer后，kv_chache有效部分是[0:17]，我们将该部分移到末尾[512-17:]，并头部清0。因为kv cache的累加发生在尾部。从第二轮开始累加后做slice操作去掉头部1个单位，取[1:]，这样就保证了kv cache始终保持在512。同时attention mask也要是从尾部实际token len长度的0，头部全部置1。
导出onnx#
将该网络分为4块：workembedding，glmblock，glmblockcache，lmhead。这里分别介绍这四块是如何导出的。
导出前，先要指定python路径，如下：
1 export pythonpath=/workspace/chatglm2-6b:$pythonpath
需要先加载原始chatglm2-6b，如下代码：
1234567891011 chatglm2_path = /workspace/chatglm2-6borigin_model = automodel.from_pretrained(chatglm2_path, trust_remote_code=true).float()origin_model.eval()transformer = origin_model.transformermax_len = transformer.seq_lengthfor param in origin_model.parameters(): param.requires_grad = falsenum_layers = transformer.encoder.num_layerslayers = transformer.encoder.layers
workembedding#
直接使用原模型中的word_embeddings，构建成独立网络，导出即可
123456789101112131415161718 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组合构建独立网路，导出。因为有28个block，所以需要导出28个onnx模型。这里的position_ids和attention_mask作为外部输入，前面有介绍。其实这28个block是可以组合成一个模型，但是这样导致onnx权重过大(f16约12gb)，导出麻烦，部署也麻烦，所以单个导出。
1234567891011121314151617181920212223 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_kvdef 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是类似的，但是需要额外的kv cache参数。注意这里 最后会把头部1个单位去除掉。
123456789101112131415161718192021222324 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，其实也是可以用argmax，看芯片实现哪一种效率高。
12345678910111213141516 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 tokendef 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模型后都已经是静态网络，通过tpu-mlir，可以很容易的转换成f16的模型。但是特别要注意的是rmsnorm需要用f32。之后就可以按照执行逻辑编写c++代码。演示效果如下：
双通道数字调节可变增益放大器MAX2063
abb机器人将向客户免费开放关键的软件服务
中美5G形成打劫局面 技术“硬核”即是先手
搭载骁龙835的小米6到底有多强？GPU频率峰值超过800Mhz官方都没摸透？
立景创新科技联合华为打造高品质制造园区，柔性生产、高效运营
ChatGLM2-6B解析与TPU部署
隆利科技荣获“Mini/Micro LED产业推动先锋奖”奖项
石英射灯电路的检修
日本补贴4760亿日元，台积电只需为熊本县工厂出资60%
微软官方对新主机命名作出解释
Nature:同行评审屡遭毒喷,中国欲启用AI系统解决问题
LLC串联谐振拓扑谐振腔电压概述
新兴存储器铁电RAM嵌入式应用的优势是什么
AR“小巨人”｜亮风台继专精特新“小巨人”后，再获上海市“科技小巨人”
高通可向荣耀供货5G芯片，无需审批
动力电池新蓝海 装备巨头入局叠片领域
华为首款全面屏新机9月22日发布,不是华为Mate10也不是荣耀Note9!而是麦芒6,惊不惊喜意不意外
如何快速成为点灯大师？
自恢复保险丝(PPTC)规避的使用四要素
科里奥利流量计的工作原理及设计