| 166 | return F.relu(x) ** 2 |
| 167 | |
| 168 | class FeedForward(nn.Module): |
| 169 | def __init__( |
| 170 | self, |
| 171 | dim: int, |
| 172 | hidden_dim: int, |
| 173 | norm_eps: float, |
| 174 | use_kernel: bool, |
| 175 | ): |
| 176 | super().__init__() |
| 177 | |
| 178 | Linear = BitLinearKernel if use_kernel else BitLinear |
| 179 | |
| 180 | self.w13 = Linear( |
| 181 | dim, |
| 182 | 2 * hidden_dim, |
| 183 | bias=False, |
| 184 | ) |
| 185 | self.w2 = Linear( |
| 186 | hidden_dim, |
| 187 | dim, |
| 188 | bias=False, |
| 189 | ) |
| 190 | self.ffn_sub_norm = RMSNorm(hidden_dim, norm_eps) |
| 191 | |
| 192 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 193 | x13 = self.w13(x) |
| 194 | x1, x3 = x13.chunk(2, -1) |
| 195 | inner = self.ffn_sub_norm(squared_relu(x1) * x3) |
| 196 | output = self.w2(inner) |
| 197 | return output |
| 198 | |
| 199 | |
| 200 | class TransformerBlock(nn.Module): |