| 439 | """ |
| 440 | |
| 441 | def __init__(self, config: ChatGLMConfig, device=None): |
| 442 | super(MLP, self).__init__() |
| 443 | |
| 444 | self.add_bias = config.add_bias_linear |
| 445 | |
| 446 | # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf |
| 447 | self.dense_h_to_4h = nn.Linear( |
| 448 | config.hidden_size, |
| 449 | config.ffn_hidden_size * 2, |
| 450 | bias=self.add_bias, |
| 451 | device=device, |
| 452 | **_config_to_kwargs(config), |
| 453 | ) |
| 454 | |
| 455 | def swiglu(x): |
| 456 | x = torch.chunk(x, 2, dim=-1) |
| 457 | return F.silu(x[0]) * x[1] |
| 458 | |
| 459 | self.activation_func = swiglu |
| 460 | |
| 461 | # Project back to h. |
| 462 | self.dense_4h_to_h = nn.Linear( |
| 463 | config.ffn_hidden_size, config.hidden_size, bias=self.add_bias, device=device, **_config_to_kwargs(config) |
| 464 | ) |
| 465 | |
| 466 | def forward(self, hidden_states): |
| 467 | # [s, b, 4hp] |