| 390 | "sdpa": NDLSdpaAttention, #使用torch中的flash |
| 391 | } |
| 392 | class NDLDecoderlayer(nn.Module): |
| 393 | def __init__(self, config:ndlconfig, layer_idx: int,select_rmsnorm:'NDLRMSNorm'): |
| 394 | super().__init__() |
| 395 | self.hidden_size = config.hidden_size |
| 396 | config.attn_implementation = 'sdpa' |
| 397 | # print(config) |
| 398 | self.self_attn = NDL_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx) |
| 399 | |
| 400 | # print(NDL_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx)) |
| 401 | self.mlp = NDLFFN(config) #前馈神经网络层 |
| 402 | self.select_rmsnorm = select_rmsnorm #选择RMSNorm层是否使用flash_attn |
| 403 | if 'Flash' not in self.select_rmsnorm: |
| 404 | self.input_layernorm = NDLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 405 | self.post_attention_layernorm = NDLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 406 | else: |
| 407 | self.input_layernorm = NDLFlash_attnRMSNorm(config.hidden_size, eps=config.rms_norm_eps) #输入层的RMSNorm |
| 408 | self.post_attention_layernorm = NDLFlash_attnRMSNorm(config.hidden_size, eps=config.rms_norm_eps) #attention之后的RMSNorm |
| 409 | def forward( |
| 410 | self, |
| 411 | hidden_states: torch.Tensor, #输入：(batch, seq_len, embed_dim) |
| 412 | attention_mask: Optional[torch.Tensor] = None, |
| 413 | position_ids: Optional[torch.LongTensor] = None, |
| 414 | past_key_value: Optional[Tuple[torch.Tensor]] = None, |
| 415 | output_attentions: Optional[bool] = False, |
| 416 | use_cache: Optional[bool] = False, |
| 417 | # cache_position: Optional[torch.LongTensor] = None, |
| 418 | **kwargs, |
| 419 | ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: |
| 420 | residual = hidden_states #第一个残差块的原始x |
| 421 | hidden_states = self.input_layernorm(hidden_states) # 对输入张量进行RMSNorm |
| 422 | # 计算自注意力机制的输出 |
| 423 | hidden_states, self_attn_weights, present_key_value = self.self_attn( |
| 424 | hidden_states=hidden_states, |
| 425 | attention_mask=attention_mask, |
| 426 | position_ids=position_ids, |
| 427 | past_key_value=past_key_value, |
| 428 | output_attentions=output_attentions, |
| 429 | use_cache=use_cache, |
| 430 | # cache_position=cache_position, |
| 431 | **kwargs, |
| 432 | ) |
| 433 | hidden_states = residual + hidden_states #残差输出 |
| 434 | #进入全连接层 |
| 435 | residual = hidden_states #FFN的残差输入X |
| 436 | hidden_states = self.post_attention_layernorm(hidden_states) #进行RMSNorm归一化 |
| 437 | hidden_states = self.mlp(hidden_states) #进入FFN块输出结果 |
| 438 | hidden_states = residual + hidden_states #残差连接 |
| 439 | outputs = (hidden_states,) |
| 440 | if output_attentions: #如果要输出out_attention |
| 441 | outputs += (self_attn_weights,) |
| 442 | |
| 443 | if use_cache: #如果使用缓存机制 |
| 444 | outputs += (present_key_value,) |
| 445 | return outputs |
| 446 | |
| 447 | class NDLPreTrainedModel(PreTrainedModel): |
| 448 | |