# GLM 块包括注意力层、FFN层和之间的残差
class GLMBlock(torch.nn.Module):
    """A single transformer layer.

    Transformer layer takes input with size [s, b, h] and returns an
    output of the same size.
    """

    def __init__(self, config: ChatGLMConfig, layer_number, device=None):
        super(GLMBlock, self).__init__()
        self.layer_number = layer_number

        self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm

        self.fp32_residual_connection = config.fp32_residual_connection
        # 判断使用 RMS 还是 LN
        LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
        # LN1
        self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
                                             dtype=config.torch_dtype)

        # 注意力层
        self.self_attention = SelfAttention(config, layer_number, device=device)
        # Dropout
        self.hidden_dropout = config.hidden_dropout

        # LLN2
        self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
                                                      dtype=config.torch_dtype)

        # FFN
        self.mlp = MLP(config, device=device)

    def forward(
            self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
    ):
        # hidden_states: [s, b, h]

        # 输入 -> LN1 -> 注意力层 -> ...
        layernorm_output = self.input_layernorm(hidden_states)
        attention_output, kv_cache = self.self_attention(
            layernorm_output,
            attention_mask,
            rotary_pos_emb,
            kv_cache=kv_cache,
            use_cache=use_cache
        )

        # 判断残差是否在LN1后面
        # 如果为真，那么：
        # 输入 -> LN1 -> 注意力 -> Dropout -> ⊕ -> ...
        #  |                                  ↑
        #  +----------------------------------+
        # 否则：
        # 输入 -> LN1 -> 注意力 -> Dropout -> ⊕ -> ...
        #          |                          ↑
        #          +--------------------------+
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states
        
        layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
        layernorm_input = residual + layernorm_input

        # ... -> LN2 -> FFN -> ...
        layernorm_output = self.post_attention_layernorm(layernorm_input)
        mlp_output = self.mlp(layernorm_output)

        # 判断残差是否在LN1后面
        # 如果为真，那么：
        # ... -> LN2 -> FFN -> Dropout -> ⊕ -> 输出
        #  |                               ↑
        #  +-------------------------------+
        # 否则：
        # ... -> LN2 -> FFN -> Dropout -> ⊕ -> 输出
        #         |                        ↑
        #         +------------------------+
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = layernorm_input

        output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
        output = residual + output

        return output, kv_cache