聊聊ChatGLM-6B的源碼分析

基於ChatGLM-6B第一版，要注意還有ChatGLM2-6B以及ChatGLM3-6B

轉載請備註出處：https://www.cnblogs.com/zhiyong-ITNote/

PrefixEncoder

作用：在微調時（以P-Tuning V2爲例），方法訓練時凍結模型的全部參數，只激活PrefixEncoder的參數。
其源碼如下，整體來看是比較簡單的。

class PrefixEncoder(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.prefix_projection = config.prefix_projection
        if self.prefix_projection:
            # 使用一個兩層(線性層)的MLP編碼prefix
            self.embedding = torch.nn.Embedding(config.pre_seq_len, config.hidden_size)
            self.trans = torch.nn.Sequential(
                torch.nn.Linear(config.hidden_size, config.hidden_size),
                torch.nn.Tanh(),
                torch.nn.Linear(config.hidden_size, config.num_layers * config.hidden_size * 2)
            )
        else:
            self.embedding = torch.nn.Embedding(config.pre_seq_len, config.num_layers * config.hidden_size * 2)

    def forward(self, prefix: torch.Tensor):
        if self.prefix_projection:
            prefix_tokens = self.embedding(prefix)
            past_key_values = self.trans(prefix_tokens)
        else:
            past_key_values = self.embedding(prefix)
        return past_key_values

爲什麼源碼註釋中會說到MLP？定位追溯：

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 default_init(cls, *args, **kwargs):
    return cls(*args, **kwargs)

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

# 轉載請備註出處：https://www.cnblogs.com/zhiyong-ITNote/

init_method對應到default_init，這個函數的作用與直接調用類構造函數相同，但它提供了一種更靈活的方式來創建類的實例，因爲它可以接受任意數量的位置參數和關鍵字參數。在Pytorch中，用於模塊化的構造函數。從源碼分析來看，GLU/MLP類就是構造了兩個線性層與gelu激活函數，其結構可簡化如下：

從PrefixEncoder類的初始化方法來看，其就是embedding層與MLP的組合。其結構可簡化如下：

詳細解讀可參考 ChatGLM的模型架構

Q：在這裏還有一個問題，從哪裏可以定位溯源到微調時禁用了全部的參數，只激活PrefixEncoder的參數並調用了該類？

轉載請備註出處：https://www.cnblogs.com/zhiyong-ITNote/

激活函數與位置編碼

代碼簡單明瞭，RoPE的理論知識可以多瞭解。

attention_fn

僞代碼表示爲：

def attention_fn(
        self,
        query_layer,
        key_layer,
        value_layer,
        attention_mask,
        hidden_size_per_partition,
        layer_id,
        layer_past=None,
        scaling_attention_score=True,
        use_cache=False,
):
    xxxx

標準的注意力機制計算公式如下：




多頭注意力就是將多個單頭注意力的結果拼接起來，再點乘一個新的權重參數。


attention_fn函數實現了注意力的核心計算過程(即上述數學表達式)，包括計算注意力分數、注意力概率和上下文層。這些計算對於實現許多自然語言處理任務，如語言建模、命名實體識別等，都是非常重要的。

SelfAttention

僞代碼表示爲：

class SelfAttention(torch.nn.Module):
    xxxx

attention_mask_func將注意力掩碼應用於Transformer模型中的注意力得分中。

@staticmethod
def attention_mask_func(attention_scores, attention_mask):
    attention_scores.masked_fill_(attention_mask, -10000.0)
    return attention_scores

apply_rotary_pos_emb_index函數爲注入了RoPE位置信息，然後調用attention_fn計算注意力概率、上下文層表示，並得到返回值。這些都是在forward函數中調用處理的。

最後還調用了dense對上下文表示做線性計算，返回輸出。

轉載請備註出處：https://www.cnblogs.com/zhiyong-ITNote/

GLU

GLU也可以理解爲是MLP，在後面版本的ChatGLM中，去掉了GLU類的定義聲明，直接換成了MLP。在上面已經寫過不再贅述。

GLMBlock

一般都會把GLMBlock對應爲transformer結構的實現。從其構造函數來看，主要是拼接各個層到一起。

從代碼來看，中間有兩次的殘差連接，如下所示

# 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

ChatGLMPreTrainedModel

TODO....

ChatGLMModel

TODO....

轉載請備註出處：https://www.cnblogs.com/zhiyong-ITNote/