| 492 | |
| 493 | class TelechatMLP(nn.Module): |
| 494 | def __init__(self, config: TelechatConfig): |
| 495 | super().__init__() |
| 496 | hidden_size = config.hidden_size |
| 497 | self.gate_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False) |
| 498 | self.up_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False) |
| 499 | self.down_proj = nn.Linear(config.ffn_hidden_size, hidden_size, bias=True) |
| 500 | self.hidden_dropout = config.hidden_dropout |
| 501 | |
| 502 | def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: |
| 503 | intermediate_output = self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) |