| 544 | |
| 545 | |
| 546 | class SwiGLUMixin(BaseMixin): |
| 547 | def __init__(self, num_layers, in_features, hidden_features, bias=False): |
| 548 | super().__init__() |
| 549 | self.w2 = nn.ModuleList( |
| 550 | [ |
| 551 | ColumnParallelLinear( |
| 552 | in_features, |
| 553 | hidden_features, |
| 554 | gather_output=False, |
| 555 | bias=bias, |
| 556 | module=self, |
| 557 | name="dense_h_to_4h_gate", |
| 558 | ) |
| 559 | for i in range(num_layers) |
| 560 | ] |
| 561 | ) |
| 562 | |
| 563 | def mlp_forward(self, hidden_states, **kw_args): |
| 564 | x = hidden_states |
| 565 | origin = self.transformer.layers[kw_args["layer_id"]].mlp |
| 566 | x1 = origin.dense_h_to_4h(x) |
| 567 | x2 = self.w2[kw_args["layer_id"]](x) |
| 568 | hidden = origin.activation_func(x2) * x1 |
| 569 | x = origin.dense_4h_to_h(hidden) |
| 570 | return x |
| 571 | |
| 572 | |
| 573 | # * Main Transformer Layer |