| 375 | """ |
| 376 | |
| 377 | def __init__(self, config: ChatGLMConfig, device=None): |
| 378 | super(MLP, self).__init__() |
| 379 | |
| 380 | self.add_bias = config.add_bias_linear |
| 381 | |
| 382 | # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf |
| 383 | self.dense_h_to_4h = nn.Linear( |
| 384 | config.hidden_size, |
| 385 | config.ffn_hidden_size * 2, |
| 386 | bias=self.add_bias, |
| 387 | device=device, |
| 388 | **_config_to_kwargs(config) |
| 389 | ) |
| 390 | |
| 391 | def swiglu(x): |
| 392 | x = torch.chunk(x, 2, dim=-1) |
| 393 | return F.silu(x[0]) * x[1] |
| 394 | |
| 395 | self.activation_func = swiglu |
| 396 | |
| 397 | # Project back to h. |
| 398 | self.dense_4h_to_h = nn.Linear( |
| 399 | config.ffn_hidden_size, |
| 400 | config.hidden_size, |
| 401 | bias=self.add_bias, |
| 402 | device=device, |
| 403 | **_config_to_kwargs(config) |
| 404 | ) |
| 405 | |
| 406 | def forward(self, hidden_states): |
| 407 | # [s, b, 4hp] |