(self, config: MoleConfig, layer_idx: int)
| 605 | |
| 606 | class MoleDecoderLayer(nn.Module): |
| 607 | def __init__(self, config: MoleConfig, layer_idx: int): |
| 608 | super().__init__() |
| 609 | self.hidden_size = config.hidden_size |
| 610 | |
| 611 | self.self_attn = MOLE_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) |
| 612 | self.num_experts = config.num_experts |
| 613 | self.mlp = MoleMLP(config) |
| 614 | self.experts = nn.ModuleList([MoleMLP(config) for _ in range(self.num_experts)]) |
| 615 | self.router = nn.Linear(self.hidden_size, self.num_experts, bias=False) |
| 616 | self.input_layernorm = MoleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 617 | self.expert_layernorm = MoleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 618 | self.post_attention_layernorm = MoleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 619 | |
| 620 | def forward( |
| 621 | self, |
nothing calls this directly
no test coverage detected