一张图说清楚：大模型“大”在哪？ChatGLM模型结构详解-阿里云开发者社区

写在前面

大型模型的神秘并不是不可透视的，今天我们以ChatGLM-6B为例，解析一下模型结构和代码。你会发现，大模型结构并没有那么神秘，相反还挺清晰的，就是Transformer的decoder改造而来的。我们还会看到模型中参数最密集的部分，这也是模型“大”的原因。

一、整体流程

ChatGLM和其他生成模型一样，都是迭代输出，每次生成一个token，然后把输出的token

拼到输入，进行下一轮迭代。我们以输入“你好吗”为例，假设输出最大长度512个token，简单梳理一下整个生成流程：

1.输入

上图中模拟了两轮迭代时模型中数据的形状，黑字是第一次迭代时模型中数据的形状，如果在第二次迭代中发生变化我用红字标出，在以后的轮次中都会遵循规律进行下去。可以看到，模型输入有三个，他们其实都要送入模型的核心模块—GLMBlock：

a.input_ids：输入的token序列，它会经过Embedding转换成hidden_state，每次迭代这个值都会拼上输出，进行下一次迭代。

b.attention_mask：根据输入的维度计算出的mask矩阵，每次迭代宽高都会加1，规律如下图，其中x是会被mask掉的token：

c.position_ids：2d旋转位置编码，详情请跳转这里。

2.流程

a.将输入做Embedding转换成hidden_state；根据输入的维度计算出的mask矩阵；根据input_ids和mask计算出position_ids。

b.将上面三个值送入28个GLMBlock，这28个GLMBlock序贯排列。每个GLMBlock的输入除了上述三个参数，还有一个layer_id，layer_id每轮+1，也就是从1-28。每个GLMBlock的输出是hidden_state，重新赋值hidden_state后进入下一轮迭代，总共28轮。

c.28个GLMBlock之后输出最后一个的hidden_state，更大规模的模型，GLMBlock就会更多。

d.hidden_state经过线性变换。也就是全连接转换成词典大小130528。

e.130528的维的结果执行argmax，得到输出的token

f.将新的token拼到input_ids后面

g.如果token序列达到预设的最大长度或最后的token=eos_token (30005)则迭代结束；否则继续迭代。

3.代码

代码可以从transformers包的transformers/generation/utils.py中的GenerationMixin.greedy_search()方法开始看，主要代码如下：

while True:
    ...
    # 1.准备输入，调用modeling_chatglm中的ChatGLMForConditionalGeneration.prepare_inputs_for_generation()方法
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
 
    # 2.执行模型推理，调用modeling_chatglm中的ChatGLMModel.forword()方法
    outputs = self(
                **model_inputs,
                return_dict=True,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
    )
 
    if synced_gpus and this_peer_finished:
        continue  # don't waste resources running the code we don't need
 
    # 3.输出后处理
    next_token_logits = outputs.logits[:, -1, :]
    next_tokens_scores = logits_processor(input_ids, next_token_logits)
 
    ...
 
    # 4.判断结束
    if eos_token_id_tensor is not None:
        unfinished_sequences = unfinished_sequences.mul(
                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
                )
 
    # stop when each sentence is finished
    if unfinished_sequences.max() == 0:
        this_peer_finished = True
 
    # stop if we exceed the maximum length
    if stopping_criteria(input_ids, scores):
        this_peer_finished = True
 
    # 5.修改input_ids 和一些状态数据，为下一和token的预测做准备
    input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
    ...

ChatGLMModel源码如下：

class ChatGLMModel(ChatGLMPreTrainedModel):

    def __init__(self, config: ChatGLMConfig, empty_init=True):
        super().__init__(config)
        if empty_init:
            init_method = skip_init
        else:
            init_method = default_init
        # recording parameters
        self.max_sequence_length = config.max_sequence_length
        self.hidden_size = config.hidden_size
        self.params_dtype = torch.half
        self.num_attention_heads = config.num_attention_heads
        self.vocab_size = config.vocab_size
        self.num_layers = config.num_layers
        self.layernorm_epsilon = config.layernorm_epsilon
        self.inner_hidden_size = config.inner_hidden_size
        self.hidden_size_per_attention_head = self.hidden_size // self.num_attention_heads
        self.position_encoding_2d = config.position_encoding_2d
        self.pre_seq_len = config.pre_seq_len
        self.prefix_projection = config.prefix_projection

        self.word_embeddings = init_method(
            torch.nn.Embedding,
            num_embeddings=self.vocab_size, embedding_dim=self.hidden_size,
            dtype=self.params_dtype
        )
        self.gradient_checkpointing = False

        def get_layer(layer_id):
            return GLMBlock(
                self.hidden_size,
                self.num_attention_heads,
                self.layernorm_epsilon,
                layer_id,
                inner_hidden_size=self.inner_hidden_size,
                hidden_size_per_attention_head=self.hidden_size_per_attention_head,
                layernorm=LayerNorm,
                use_bias=True,
                params_dtype=self.params_dtype,
                position_encoding_2d=self.position_encoding_2d,
                empty_init=empty_init
            )

        self.layers = torch.nn.ModuleList(
            [get_layer(layer_id) for layer_id in range(self.num_layers)]
        )

        # Final layer norm before output.
        self.final_layernorm = LayerNorm(self.hidden_size, eps=self.layernorm_epsilon)

        if self.pre_seq_len is not None:
            for param in self.parameters():
                param.requires_grad = False
            self.prefix_tokens = torch.arange(self.pre_seq_len).long()
            self.prefix_encoder = PrefixEncoder(config)
            self.dropout = torch.nn.Dropout(0.1)

            # total_params = sum(p.numel() for p in self.parameters())
            # trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
            # print("Using p-tuning v2: # trainable_params = {} / {}".format(trainable_params, total_params))


    def forward(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
            inputs_embeds: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    )
 -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPast]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape[:2]
        elif inputs_embeds is not None:
            batch_size, seq_length = inputs_embeds.shape[:2]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        # [seq_len, batch, hidden_size]
        hidden_states = inputs_embeds.transpose(0, 1)

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None

        if attention_mask is None:
            attention_mask = torch.zeros(1, 1, device=input_ids.device).bool()
        else:
            attention_mask = attention_mask.to(hidden_states.device)

        for i, layer in enumerate(self.layers):

            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            layer_past = past_key_values[i]

            layer_ret = layer(
              hidden_states,
              position_ids=position_ids,
              attention_mask=attention_mask,
              layer_id=torch.tensor(i),
              layer_past=layer_past,
              use_cache=use_cache,
              output_attentions=output_attentions
            )

            hidden_states = layer_ret[0]

            if use_cache:
                presents = presents + (layer_ret[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_ret[2 if use_cache else 1],)

        # Final layer norm.
        hidden_states = self.final_layernorm(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )

二、GLMBlock

GLMBlock是核心模块，整体流程如下：

1.整体流程

a.LayerNorm:hidden_states经过LayerNorm，然后送入SelfAttention

b.SelfAttention:

输入经过全连接升维至12288，然后拆出QKV；

QK结合position_ids和block_position_ids，分别经过2d旋转位置编码重新编排；

新的Q和K做内积并进行Attention缩放，得到attention_scores；

使用mask将部分attention_scores的值去掉（设置成一个很小的值），再经过softmax处理，attention_scores是一个概率图，标志每个token和其它token的相关性；

将attention_scores与V做内积，即对V施加注意力机制，得到attention_output；

attention_input与attention_output做残差相加，attention_input的系数是，其中num_layers=28。

c.LayerNorm：再LayerNorm一遍

d.GLU:Gated Linear Unit（门控线性单元），也可以叫做FFN(Feedforward Neural Network)或者MLP（Multilayer Perceptron），由两个全连接夹着一个GELU激活函数组成。GLU之后接一个带权重的残差输出。

GLMBlock代码：

class GLMBlock(torch.nn.Module):
    def __init__(
            self,
            hidden_size,
            num_attention_heads,
            layernorm_epsilon,
            layer_id,
            inner_hidden_size=None,
            hidden_size_per_attention_head=None,
            layernorm=LayerNorm,
            use_bias=True,
            params_dtype=torch.float,
            num_layers=28,
            position_encoding_2d=True,
            empty_init=True
    ):
        super(GLMBlock, self).__init__()
        # Set output layer initialization if not provided.
 
        self.layer_id = layer_id
 
        # Layernorm on the input data.
        self.input_layernorm = layernorm(hidden_size, eps=layernorm_epsilon)
 
        self.position_encoding_2d = position_encoding_2d
 
        # Self attention.
        self.attention = SelfAttention(
            hidden_size,
            num_attention_heads,
            layer_id,
            hidden_size_per_attention_head=hidden_size_per_attention_head,
            bias=use_bias,
            params_dtype=params_dtype,
            position_encoding_2d=self.position_encoding_2d,
            empty_init=empty_init
        )
 
        # Layernorm on the input data.
        self.post_attention_layernorm = layernorm(hidden_size, eps=layernorm_epsilon)
 
        self.num_layers = num_layers
 
        # GLU
        self.mlp = GLU(
            hidden_size,
            inner_hidden_size=inner_hidden_size,
            bias=use_bias,
            layer_id=layer_id,
            params_dtype=params_dtype,
            empty_init=empty_init
        )
 
    def forward(
            self,
            hidden_states: torch.Tensor,
            position_ids,
            attention_mask: torch.Tensor,
            layer_id,
            layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
            use_cache: bool = False,
            output_attentions: bool = False,
    ):
        """
        hidden_states: [seq_len, batch, hidden_size]
        attention_mask: [(1, 1), seq_len, seq_len]
        """
 
        # Layer norm at the begining of the transformer layer.
        # [seq_len, batch, hidden_size]
        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,
            layer_past=layer_past,
            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
 
        if use_cache:
            outputs = (output,) + outputs
        else:
            outputs = (output,) + outputs[1:]
 
        return outputs  # hidden_states, present, attentions

2.旋转位置编码

上面提到的旋转位置编码是一种用绝对位置编码的方式实现相对位置编码的编码方式，详情请看这里。

3.Attention缩放

上面是Attention公式，其中d是查询向量和键向量的维度，这里是128，Attention缩放指的是。那为什么QK要除以呢，原因如下：

a.如果softmax的输入值过大，那它的梯度会趋近于0，导致梯度消失，所以要保证softmax的输入大小合适，最好是一个有效的概率分布。

b.在注意力机制中，通常会对注意力分数进行归一化处理，即N(0, 1)。在经过layernom之后，QK能符合这个标准，他俩的内积均值还是0 ，但是标准差变成了d，推导过程如下：

由于每个元素和都是独立的，并且符合 N(0,1) 分布，所以矩阵 Q 和 K 的内积的方差等于所有元素的方差之和。

由于和都是来自 N(0,1)的独立随机变量，它们的方差都是 1，所以矩阵 Q 和 K 的内积的方差为所有元素的方差之和。

因此，在这种情况下，矩阵 Q 和 K 内积的标准差是它们的维度 nn。所以要让std(QK)=1，除以就可以了。

4.mask的作用

我们观察mask矩阵发现，在生成每个token时，通常会“mask掉”刚刚生成的token，这主要有下面几个原因：

a.避免重复：如果模型在生成下一个token时还看到刚刚生成的token，它可能会过分关注这个token，并倾向于重复它，这会使得生成的文本显得不自然和重复。

b.防止错误累积：在生成较长的文本序列时，如果模型在某个步骤生成了一个不合适的token，而这个token又被用于生成后续的token，可能会导致错误在序列中累积，影响整个生成文本的质量。通过mask掉刚刚生成的token，可以减少这种错误累积的可能性。

c.缺乏远见：当模型在生成文本时，它应该考虑整个上下文的历史，并预测一个能够推动故事或对话向前发展的token。如果模型过分关注刚刚生成的token，它可能会忽视更广泛的上下文，导致生成的文本缺乏远见和创造性。

5.残差系数

在两个残差处，都使用了缩放系数，其中num_layers=28。

因为28个GLMBlock是串联的，每个GLMBlock的输出都要与输入做融合。那么有一个缩放系数是必要的，从这个系数可以看出输入与输出的权重比例大约为7.4:1，随着层数的增加，每层输出的影响力是一点一点增加的，这很合理。

那缩放系数为什么是这个样子呢，我能力有限没有找到原因，欢迎知道的大佬指点一下。但是可以简单的看一下这个系数的意义：

a.模型的规模变大，num_layers变大，组成结果的层数变大，每层（即输出）的权重变小是合理的；反之亦然。所以这个系数与层数成正比是合理的。

b.开根号可以使这个数的分布更平滑

c.2可以调节输入的权重规模，起到一定的人为调控的作用。

6.FFN （又称MLP）

FFN由两个全连接夹着一个GELU组成，两个全连接的中间维度一般设置为hidden_state的整数倍，这里是4（即4096*4=16384）

值得一提的是这里是模型参数最密集的地方，有两个全连接。而且从Transformer诞生到现在，有很多部分已经得到优化，唯独没有这块的优化。实验表明，如果FFN变小，模型整体的性能也会随之下降，基本可以认定，模型的“知识”是存在于FFN中的，所以这里需要更多的参数。

代码如下：

class GLU(torch.nn.Module):
    def __init__(self, hidden_size, inner_hidden_size=None,
                 layer_id=None, bias=True, activation_func=gelu, params_dtype=torch.float, empty_init=True):
        super(GLU, self).__init__()
        if empty_init:
            init_method = skip_init
        else:
            init_method = default_init
        self.layer_id = layer_id
        self.activation_func = activation_func
 
        # Project to 4h.
        self.hidden_size = hidden_size
        if inner_hidden_size is None:
            inner_hidden_size = 4 * hidden_size
        self.inner_hidden_size = inner_hidden_size
        self.dense_h_to_4h = init_method(
            torch.nn.Linear,
            self.hidden_size,
            self.inner_hidden_size,
            bias=bias,
            dtype=params_dtype,
        )
        # Project back to h.
        self.dense_4h_to_h = init_method(
            torch.nn.Linear,
            self.inner_hidden_size,
            self.hidden_size,
            bias=bias,
            dtype=params_dtype,
        )
 
    def forward(self, hidden_states):
        """
        hidden_states: [seq_len, batch, hidden_size]
        """
 
        # [seq_len, batch, inner_hidden_size]
        intermediate_parallel = self.dense_h_to_4h(hidden_states)
 
        intermediate_parallel = self.activation_func(intermediate_parallel)
 
        output = self.dense_4h_to_h(intermediate_parallel)
 
        return output

三、总结

最后总结一下：

1.ChatGLM是根据Transformer的decoder改造的；

2.位置编码方式是旋转位置编码；

3.残差缩放系数；

4.GLMBlock的结构是LayerNorm+SelfAttention+LayerNorm+GLU（FFN）；

5.ChatGLM的核心是GLMBlock，不通大小的模型是通过调整GLMBlock的个数来实现，比如6B有28个，10B有48个；

6.模型的“知识”是存在于FFN中的，这也是大模型最“大”的地方。

7.从模型结构来看，理论上输入和输出的长度是没有限制的，但实际使用中一定要设置一个最大长度（从ChatGLM的原来来看，这个长度值得是输入和输出的总长度），这是因为算力限制，而且训练集的长度不是无限的，一般情况下训练集限制在什么长度决定了推理时能支持的最大长度。

ChatGLM的结构整体看下来中规中矩，但是效果很不错，除了模型本身，训练方式也很重要，接下来开始训练模型，敬请期待。

ChatGLM模型结构就介绍到这里，关注不迷路(#^.^#)

一张图说清楚：大模型“大”在哪？ChatGLM模型结构详解