(self, config:ndlconfig, layer_idx: int,select_rmsnorm:'NDLRMSNorm')
| 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) |
nothing calls this directly
no test coverage detected