Method forward

model/modeling_ndl.py:409–445

(
        self,
        hidden_states: torch.Tensor, #输入：(batch, seq_len, embed_dim)
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        # cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    )

            self.input_layernorm = NDLFlash_attnRMSNorm(config.hidden_size, eps=config.rms_norm_eps) #输入层的RMSNorm
            self.post_attention_layernorm = NDLFlash_attnRMSNorm(config.hidden_size, eps=config.rms_norm_eps) #attention之后的RMSNorm
    def forward(
        self,
        hidden_states: torch.Tensor, #输入：(batch, seq_len, embed_dim)
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        # cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        residual = hidden_states #第一个残差块的原始x
        hidden_states = self.input_layernorm(hidden_states) # 对输入张量进行RMSNorm
        # 计算自注意力机制的输出
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            # cache_position=cache_position,
            **kwargs,
        )
        hidden_states = residual + hidden_states #残差输出
        #进入全连接层
        residual = hidden_states #FFN的残差输入X
        hidden_states = self.post_attention_layernorm(hidden_states) #进行RMSNorm归一化
        hidden_states = self.mlp(hidden_states) #进入FFN块输出结果
        hidden_states = residual + hidden_states #残差连接
        outputs = (hidden_states,)
        if output_attentions: #如果要输出out_attention
            outputs += (self_attn_weights,)
        
        if use_cache: #如果使用缓存机制
            outputs += (present_key_value,)
        return outputs

class NDLPreTrainedModel(PreTrainedModel):