ChatGLM2-6B模型推理流程和模型架构详解

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

文章目录

前言
1 ChatGLM是什么？
2 一代GLM
- 2.1 大模型架构
- 2.2 GLM特点
2 二代GLM：ChatGLM2-6B为例拆解
- 2.1 ChatGLM2-6B模型推理架构和流程
- 2.2 细节详解
- - 第一步：输入与分词、编码
  - 第二步：嵌入Embedding，可迁移重用
  - 第三步：GLMBlock*28
  - 第四步：生成next token
3 总结

前言

因为本人在做大模型优化方面的研究，之前拆了ChatGLM2的源代码，看看能从哪些地方深入。结果刚拆完没多久，昨天，也就是10 月 27 日，智谱 AI 在 2023 中国计算机大会(CNCC)上发布了自研第三代对话大模型 ChatGLM3，这是智谱 AI 在今年内第三次对 ChatGLM 基座模型进行了深度优化。目前还没去拆它的源代码，所以也不太清楚和2代之间有什么区别。但2代的结构我觉得可以先发以下。

1 ChatGLM是什么？

和ChatGPT类似，ChatGLM是基于GLM大模型的下游对话应用。GLM的全称是通用语言模型模型General Language model，是清华大学与智谱AI研发的中英双语大语言模型。官方API的ChatGLM是基于GLM-130B千亿基础模型，但官方也发布了GLM-6B小参数（62亿）版本，可在消费级显卡上部署。

2 一代GLM

一代发布的时间大概在一年多前，成果最早以一篇论文的形似和发布，有兴趣的可以去arxiv上看看原文（https://arxiv.org/pdf/2103.10360.pdf）。这篇文章的大概意思我稍稍总结了如下：

2.1 大模型架构

首先介绍一下目前基于transformer的大模型架构类别，主要有三类：

自编码：encoder架构，善于语言建模和理解，相当于特征更凝练的embedding，Bert系列之后好像就没有新模型了。
自回归：decoder架构，善于根据之前的信息生成，主流架构，如GPT等都在用，非常适合对话应用。
编码器-解码器：encoder-decoder架构，善于处理seq2seq任务如翻译，2017年标准transformer就是用翻译器作为示例的，比较有名的有T5。今天要讲的GLM在本质上可以说是对T5的优化改进

可以看到，自编码和自回归是两种不同的架构和自然语言处理应用思路。之前的语言模型各有优缺点，但没有一种框架能够在所有的自然语言处理任务中都表现出色。一些先前的工作尝试通过多任务学习的方式，将不同框架的目标结合起来，但由于自编码和自回归目标本质上的不同，简单的结合不能充分继承两者的优势。因此，清华大学提出了一种基于自回归空白填充的通用语言模型（GLM），来解决这个挑战。GLM 通过添加二维位置编码和允许任意顺序预测空白区域，改进了空白填充预训练，在自然语言理解任务上超越了 BERT 和 T5。

2.2 GLM特点

主要有三个特点

自编码，随机 MASK 输入中连续spans的 token
自回归，基于自回归空白填充的方法重新构建spans中的内容
2维的位置编码技术，来表示span间和span内位置信息

GLM 从输入文本中随机挖掉一些连续的词语（自编码思路），然后训练模型按照一定的顺序逐个恢复这些词语（自回归思路）。这种方法结合了自编码和自回归两种预训练方式的优点。
此外，GLM打乱了空白区域的预测顺序，并使用二维位置编码（第一个维度对span在原文本中的位置进行编码，第二个维度对token在span中的位置进行编码）。实验表明，GLM 在参数量和计算成本相同的情况下，能够在 SuperGLUE 基准测试中显著超越BERT，并且在使用相似规模的语料（158GB）预训练时，能够超越 RoBERTa 和 BART。GLM 还能够在自然语言理解和生成任务上显著超越 T5，而且使用的参数和数据更少。

2 二代GLM：ChatGLM2-6B为例拆解

GLM2和1在模型源代码方面没有什么区别，都是Prefix Decoder-only架构的，这种架构是在以GPT为代表的Causal Decoder-only架构上发展而来的，综合了单项注意力和双向注意力的优点（在前一部分采用双向注意力获得更加全面的信息，在后一部分采用单项注意力以适应生成式任务）。

2.1 ChatGLM2-6B模型推理架构和流程

在实验室服务器上本地部署后（如果想知道怎么部署的之后可以单写一篇），用"你好"这个最简单的输入进行测试，得到整体流程如下（这里更正一下，在输入shape的问题上，在使用kv_cache的情况下，只有第一轮迭代的输入序列长度是大于1的，之后序列长度都等于1）

可以看到，这个模型在推理阶段主要由两层循环组成。

第一层是while true循环，每循环一次生成一个next token，退出条件是模型生成了这个token，就是结束符。
第二层循环是固定28次的for循环，对GLMBlock顺序运行28次，根据的attention scores得到最有可能的token id。

2.2 细节详解

第一步：输入与分词、编码

输入“你好”
（1）先被自动填充嵌入一个简单的prompt变成“[Round 1] 问：你好答：”
（2）根据分词器的词典进行分词，模式采用的是wordpiece分词法，基本原理就是一个预先给定的词表中去套用输入文本，根据概率将给定文本分割成基本单元词片，然后用Int32的下标作为词片的id。
从上图可以看到，“[Round 1] 问：你好答：”字符串在经过这一步时被处理成了一个长度为17的整数数组。其中前两个数字64790、64792是固定的开头标识。那么这个给定的分词表之后还会加上一些特殊的Token，最后形成一个长度为65024的词表。

第二步：嵌入Embedding，可迁移重用

Embedding层的参数是可训练，在你下载到本地的model文件夹中有7个二进制参数文件和一个映射表，表中可以查看ChatGLM2-6B的所有层的参数存在哪个二进制文件中。模型的Embedding层参数以及预训练好了，从图中可以看到，Embedding层的形状为65024*4096，即对第一步中提到的size=65024的词表中的每个词都可以映射为长度为4096的特征向量。
的所以只需要把第一步生成的形状为17*1的整数数组输入，即可得到一个17*4096的嵌入。由于ChatGLM2-6B支持同时输入多句话，所以真实的输入维度为[17,1,4096]分别对应序列长度（多句输入时会对所有句子padding统一到最长序列）、批数、嵌入的特征空间维度。

第三步：GLMBlock*28

GLMBlock结构图如下。可以看到也是对Transformer进行了一个魔改，但主体还是一个注意力模块和一个MLP全连接模块。下面还是从输入流的角度看看block的结构，。
（1）首先一上来就是一个RMS归一化层。
（2）进入注意力模块，对输入数据进行QKV映射，得到输入数据的Query、Key、Value值，形状分别为[17,1,32,128]、[17,1,2,128]、[17,1,2,128]，可以看到Key-value的形状保持一致。
（3）然后经过核心注意力运算，也就是缩放点积注意力层，Query和Key点积后消除量纲，再点积Value，随后reshape回[17,1,4096]的形式。
（4）离开attention模块、进入MLP模块前，要完成三个操作：Dropout、残差连接、后归一化，如图。这里的残差add的值是还没经过前归一化的原始输入。
（5）进入MLP模块：要经过两次变换，中间的激活函数是SwiGLU。在这一层，虽然输入、输出的维度还是4096，但在中间过程涨到了27392，极大丰富了表示能力。
（6）离开MLP后，也要经过Dropout和残差连接，这里加的是attention结束后、后归一化前的值。最后输出GLMBlock的是一个和输入Block的shape一样的矩阵。
（7）在for循环控制下，这个输出矩阵被当作下一轮的输入，一共要走28次。这28次的Block的参数都不一样，具体见参数映射表。

第四步：生成next token

在第三步的28轮GLMBlcok循环后，最后一层block输出的attention scores会被用于输出处理。输出处理主要有三个组件：
第一个是RMS归一化，相当于弥补了Block中MLP层出来后没有进行的归一化。
第二个是嵌入的逆操作，将输出的[17,1,4096]维的嵌入向量还原为id值，变成一组65024长度的输出logits。这一步实际上是根据attention scores对65025的词表进行了一个概率估计的操作。
第三个是Softmax操作，将logits变为概率（和为1），并选择其中的最大值，输出其id。这个id就是这一轮28次循环生成的token的id。在我的实例中，这个id是36474，代表“你”。
在最外层的while true循环下，只要这个id不是2（这个Token的id），就会一直生成next token。本次测试最后生成的结果如图1最下部。