SwiGLU Feed-Forward Network. Reference: GLU Variants Improve Transformer (https://arxiv.org/abs/2002.05202)
| 49 | |
| 50 | |
| 51 | class SwiGLUFFN(nn.Module, ListForwardMixin): |
| 52 | """SwiGLU Feed-Forward Network. |
| 53 | |
| 54 | Reference: GLU Variants Improve Transformer (https://arxiv.org/abs/2002.05202) |
| 55 | """ |
| 56 | |
| 57 | def __init__( |
| 58 | self, |
| 59 | in_features: int, |
| 60 | hidden_features: Optional[int] = None, |
| 61 | out_features: Optional[int] = None, |
| 62 | act_layer: Optional[Callable[..., nn.Module]] = None, |
| 63 | drop: float = 0.0, |
| 64 | bias: bool = True, |
| 65 | align_to: int = 8, |
| 66 | device=None, |
| 67 | ) -> None: |
| 68 | super().__init__() |
| 69 | out_features = out_features or in_features |
| 70 | hidden_features = hidden_features or in_features |
| 71 | d = int(hidden_features * 2 / 3) |
| 72 | swiglu_hidden_features = d + (-d % align_to) |
| 73 | self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) |
| 74 | self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) |
| 75 | self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device) |
| 76 | |
| 77 | def forward(self, x: Tensor) -> Tensor: |
| 78 | x1 = self.w1(x) |
| 79 | x2 = self.w2(x) |
| 80 | hidden = F.silu(x1) * x2 |
| 81 | return self.w3(hidden) |
nothing calls this directly
no outgoing calls
no test coverage detected