| 772 | |
| 773 | |
| 774 | class PackedSwiGLUFFN(nn.Module): |
| 775 | def __init__( |
| 776 | self, |
| 777 | dim, |
| 778 | hidden_dim, |
| 779 | multiple_of, |
| 780 | ffn_dim_multiplier=None, |
| 781 | device=None, |
| 782 | dtype=None, |
| 783 | ): |
| 784 | factory_kwargs = {"device": device, "dtype": dtype} |
| 785 | super().__init__() |
| 786 | hidden_dim = int(2 * hidden_dim / 3) |
| 787 | # custom dim factor multiplier |
| 788 | if ffn_dim_multiplier is not None: |
| 789 | hidden_dim = int(ffn_dim_multiplier * hidden_dim) |
| 790 | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| 791 | |
| 792 | self.w13 = nn.Linear(dim, 2 * hidden_dim, bias=False, **factory_kwargs) |
| 793 | self.w2 = nn.Linear(hidden_dim, dim, bias=False, **factory_kwargs) |
| 794 | |
| 795 | def forward(self, x): |
| 796 | x1, x3 = torch.chunk(self.w13(x), 2, dim=-1) |
| 797 | return self.w2(F.silu(x1) * x3) |
| 798 | |
| 799 | |
| 800 | ################################################################################ |
no outgoing calls
no test coverage detected