| 743 | |
| 744 | class SwiGLUFFN(nn.Module): |
| 745 | def __init__( |
| 746 | self, |
| 747 | dim, |
| 748 | hidden_dim, |
| 749 | multiple_of, |
| 750 | ffn_dim_multiplier=None, |
| 751 | device=None, |
| 752 | dtype=None, |
| 753 | ): |
| 754 | factory_kwargs = {"device": device, "dtype": dtype} |
| 755 | super().__init__() |
| 756 | hidden_dim = int(2 * hidden_dim / 3) |
| 757 | # custom dim factor multiplier |
| 758 | if ffn_dim_multiplier is not None: |
| 759 | hidden_dim = int(ffn_dim_multiplier * hidden_dim) |
| 760 | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| 761 | |
| 762 | self.w1 = nn.Linear(dim, hidden_dim, bias=False, **factory_kwargs) |
| 763 | self.w2 = nn.Linear(hidden_dim, dim, bias=False, **factory_kwargs) |
| 764 | self.w3 = nn.Linear(dim, hidden_dim, bias=False, **factory_kwargs) |
| 765 | |
| 766 | def forward(self, x): |
| 767 | return self.w2(F.silu(self.w1(x)) * self.w3(x)) |