一、前言 本文将介绍ChatGLM-6B的模型结构组件源码。
代练链接：https:///THUDM/chatglm-6b/blob/main/modeling_chatglm.py
二、激活函数 @torch.jit.script def gelu_impl (x) :'''OpenAI's gelu implementation.''' return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x *1.0 + 0.044715 * x * x)))def gelu (x) :return gelu_impl(x)
三、位置编码 ChatGLM-6B的位置编码采用的旋转位置编码（详细推导过程见原文：ROPE），简单来说其目的就是构建一个包含相对位置信息的Attention矩阵，其公式如下：
式中，
原作者提到，由于
这里直接上代码阅读
class RotaryEmbedding (torch.nn.Module) :def __init__ (self, dim, base=10000 , precision=torch.half, learnable=False) :1. / (base ** (torch.arange(0 , dim, 2 ).float() / dim))if learnable:None else :'inv_freq' , inv_freq)None None None def _load_from_state_dict (self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, :pass def forward (self, x, seq_dim=1 , seq_len=None) :if seq_len is None :if self.max_seq_len_cached is None or (seq_len > self.max_seq_len_cached):None if self.learnable else seq_len# 在计算旋转嵌入之前，根据当前的嵌入维度和基数计算频率因子 inv_freq。将其转换为半精度数据类型（如果指定的 precision 为 bfloat16，则转换为单精度）。 # 使用 爱因斯坦求和函数 einsum 将 t 和 inv_freq 相乘，得到频率矩阵 freqs。 'i,j->ij' , t, self.inv_freq)# 通过在频率矩阵 freqs 中进行重复和拼接操作，生成旋转嵌入矩阵 emb，其维度为 [seq_len, 2 * dim]。 -1 ).to(x.device)if self.precision == torch.bfloat16:# 将旋转嵌入矩阵 emb 分别进行余弦和正弦运算。 None , :]None , :]if self.precision == torch.bfloat16:if self.learnable:return cos_cached, sin_cached# 按照序列长度截取 return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]四、注意力层 ChatGLM-6B代码中这一层采用的位置编码是GLM的中提出的2D位置编码，详细原理见原文：GLM: General Language Model Pretraining with Autoregressive Blank Infilling，其原理图如下图：
输入的序列是
ChatGLM-6B相比标准的自注意力机制在Q和K中注入了RoPE位置信息。
def attention_fn ( :# 考虑过去的信息 if layer_past is not None :0 )0 )# seqlen, batch, num_attention_heads, hidden_size_per_attention_head if use_cache:else :None # 对查询层进行缩放操作，即将其除以（隐藏层大小的平方根乘以查询层的缩放系数）。这是为了控制注意力得分的尺度。 1 )if scaling_attention_score:# =================================== # Raw attention scores. [b, np, s, s] # =================================== # # 注意力分数的输出形状: [batch_size, num_heads, seq_length, seq_length] 1 ), query_layer.size(2 ), query_layer.size(0 ), key_layer.size(0 ))# 形状重塑：[seq_length, batch_size, num_heads, head_dim] -> [seq_length, batch_size*num_heads, head_dim] # [sq, b, np, hn] -> [sq, b * np, hn] 2 ], output_size[0 ] * output_size[1 ], -1 )# [sk, b, np, hn] -> [sk, b * np, hn] 3 ], output_size[0 ] * output_size[1 ], -1 )0 ] * output_size[1 ],2 ],3 ],# 计算原始的注意力得分，通过转置和重塑操作，将查询、键和值的张量形状调整为合适的形状。 0 , 1 ), # [b * np, sq, hn] 0 , 1 ).transpose(1 , 2 ), # [b * np, hn, sk] 0.0 ,1.0 ,# 重塑形状为：[batch_size,num_head,seq_length,seq_length] # 如果指定了缩放的掩码 softmax（scale_mask_softmax），则将注意力得分传递给缩放的掩码 softmax 函数进行处理，以获得归一化的注意力概率。 # 否则，将应用 softmax 操作，并根据需要填充一个较大的负数值（-10000.0）来屏蔽无效位置。 if self.scale_mask_softmax:else :# 对注意力分数进行mask if not (attention_mask == 0 ).all():# if auto-regressive, skip -10000.0 )-1 )# ========================= # Context layer. [sq, b, hp] # ========================= # value_layer -> context layer. # [sk, b, np, hn] --> [b, np, sq, hn] # context layer shape: [b, np, sq, hn] 1 ), value_layer.size(2 ), query_layer.size(0 ), value_layer.size(3 ))# change view [sk, b * np, hn] 0 ), output_size[0 ] * output_size[1 ], -1 )# 对注意力分数进行mask # change view [b * np, sq, sk] 0 ] * output_size[1 ], output_size[2 ], -1 )# matmul: [b * np, sq, hn] 0 , 1 ))# change view [b, np, sq, hn] # [b, np, sq, hn] --> [sq, b, np, hn] 2 , 0 , 1 , 3 ).contiguous()# [sq, b, np, hn] --> [sq, b, hp] -2 ] + (hidden_size_per_partition,)# 重塑上下文层 return outputs
SelfAttention的目的是为了捕捉序列中的位置信息，应用RoPE将位置信息注入Q和K。
class SelfAttention (torch.nn.Module) :def __init__ (self, hidden_size, num_attention_heads, :2 )if position_encoding_2delse self.hidden_size // self.num_attention_heads,10000 ,False ,None if hidden_size_per_attention_head is None :else :# Strided linear layer. 3 * self.inner_hidden_size, @staticmethod def attention_mask_func (attention_scores, attention_mask) :-10000.0 )return attention_scoresdef split_tensor_along_last_dim (self, tensor, num_partitions, :'''Split a tensor along its last dimension. # Get the size and dimension. 1 # Split. # Note: torch.split does not create contiguous tensors by default. if contiguous_split_chunks:return tuple(chunk.contiguous() for chunk in tensor_list)return tensor_listdef forward ( :''' # [seq_len, batch, 3 * hidden_size] # [seq_len, batch, 3 * hidden_size] --> [seq_len, batch, num_attention_heads, 3 * hidden_size_per_attention_head] -1 ] + (3 * self.hidden_size_per_attention_head,# [seq_len, batch, num_attention_heads, hidden_size_per_attention_head] 3 )# 根据是否使用二维位置编码，对查询和键应用旋转嵌入，并根据位置信息进行索引操作。 if self.position_encoding_2d:2 , dim=(query_layer.ndim - 1 ))2 , dim=(key_layer.ndim - 1 ))1 )0 , :].transpose(0 , 1 ).contiguous(), \1 , :].transpose(0 , 1 ).contiguous()# 拼接嵌入不同位置信息的query和key，这样query和key中包含了两种位置信息 1 ))1 ))else :# RoPE 0 , 1 )1 )# [seq_len, batch, num_attention_heads, hidden_size_per_attention_head] # 调用 attention_fn 方法计算注意力得分和上下文层，其中使用了注意力函数的代码块 # [seq_len, batch, hidden_size] if output_attentions:return outputs # output, present, attention_probs 五、GLU层 根据代码，GLU形式化表示为：
class GEGLU (torch.nn.Module) :def __init__ (self) :def forward (self, x) :# dim=-1 breaks in jit for pt<1.10 2 , dim=(x.ndim - 1 ))return x1 * self.activation_fn(x2)class GLU (torch.nn.Module) :def __init__ (self, hidden_size, inner_hidden_size=None, :# Project to 4h. if inner_hidden_size is None :4 * hidden_size# Project back to h. def forward (self, hidden_states) :''' # [seq_len, batch, inner_hidden_size] return output
六、GLMBlock 根据代码，GLMBlock由Layer Norm、Self Attention、Layer Norm和GLU模块构成。
class GLMBlock (torch.nn.Module) :def __init__ (28 , :# Set output layer initialization if not provided. # LayerNorm层 # 是否使用2维位置编码 # 自注意力层 # LayerNorm层 # GLU层 def forward ( :''' # 输入进行Layer Norm # [seq_len, batch, hidden_size] # 自注意力 0 ]1 :]# Residual connection. 2 * self.num_layers) ** 0.5 # 执行注意力残差连接 # 对注意力残差连接后的输出进行层归一化 # 使用GLU层对归一化后的输出进行非线性变换 # 执行GLU残差连接 if use_cache:else :1 :]return outputs # hidden_states, present, attentions 七、ChatGLMPreTrainedModel 这一块主要看看其中的MASK 和Position_ids
ChatGLM-6B采用prefix-LM的Mask，其对于输入的前缀使用双向注意力 ，对于后续的生成部分则是Causal Mask 。
def get_masks (self, input_ids, device) :# context_lengths记录了batch中每个样本的真实长度 for seq in input_ids]# 生成causal mask，即下三角以及对角线为1，上三角为0 # 将前缀部分的注意力改为双向注意力 for i, context_length in enumerate(context_lengths):1 1 )0.5 ).bool()return attention_mask
def get_position_ids (self, input_ids, mask_positions, device, use_gmasks=None) :''' if use_gmasks is None :False ] * batch_size# context_lengths：未被padding前，batch中各个样本的长度 for seq in input_ids]# 2维位置编码 if self.position_encoding_2d:# [0,1,2,...,seq_length-1] 0 ).repeat(batch_size, 1 )# 将原始输入后所有位置的postion id都设置为[Mask]或者[gMask]的位置id for i, context_length in enumerate(context_lengths):# 原始输入的位置编码全部设置为0，待生成的位置添加顺序的位置id # 例如：[0,0,0,0,1,2,3,4,5] 1 for context_length in context_lengths]0 )# 将postion_ids和block_position_ids堆叠在一起，用于后续的参数传入； # 在注意力层中，还有将这个position_ids拆分为两部分 1 )else :0 ).repeat(batch_size, 1 )for i, context_length in enumerate(context_lengths):if not use_gmasks[i]:return position_ids八、ChatGLMModel 这一块主要是模型的各部件的组合结构，直接看源码：
class ChatGLMModel (ChatGLMPreTrainedModel) :def __init__ (self, config: ChatGLMConfig, empty_init=True) :if empty_init:else :# recording parameters False def get_layer (layer_id) :return GLMBlock(True ,for layer_id in range(self.num_layers)]# Final layer norm before output. ''' if self.pre_seq_len is not None :for param in self.parameters():False 0.1 )def get_prompt (self, batch_size, device, dtype=torch.half) :''' 0 ).expand(batch_size, -1 ).to(device)2 ,# seq_len, b, nh, hidden_size 2 , 1 , 0 , 3 , 4 ]).split(2 )# past_key_values = [(v[0], v[1]) for v in past_key_values] return past_key_valuesdef forward ( -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPast]:''' if past_key_values is None :if self.pre_seq_len is not None :0 ], device=input_ids.device,else :None ] * len(self.layers))if attention_mask is None :if position_ids is None :''' for seq in seqs:if gMASK in seq else MASKif self.pre_seq_len is not None and attention_mask is not None :1 , input_ids.size(-1 ), self.pre_seq_len).to(0.5 ).bool()3 )''' # [seq_len, batch, hidden_size] 0 , 1 )if use_cache else None if output_attentions else None if output_hidden_states else None if attention_mask is None :1 , 1 , device=input_ids.device).bool()else :for i, layer in enumerate(self.layers):if output_hidden_states:if self.gradient_checkpointing and self.training:else :0 ]if use_cache:1 ],)if output_attentions:2 if use_cache else 1 ],)# Final layer norm. if output_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(
如有不详细之处，还望一起交流学习。
参考文献 GLM: General Language Model Pretraining with Autoregressive Blank Infilling) modeling_chatglm.py · THUDM/chatglm-6b at main () Transformer升级之路：2、博采众长的旋转式位置编码 - 科学空间|Scientific Spaces (kexue.fm)
