| 224 | return causal_mask_bool |
| 225 | |
| 226 | class FeedForward(nn.Module): |
| 227 | def __init__( |
| 228 | self, |
| 229 | dim: int, |
| 230 | hidden_dim: int, |
| 231 | multiple_of: int, |
| 232 | ffn_dim_multiplier: Optional[float], |
| 233 | ): |
| 234 | super().__init__() |
| 235 | hidden_dim = int(2 * hidden_dim / 3) |
| 236 | # custom dim factor multiplier |
| 237 | if ffn_dim_multiplier is not None: |
| 238 | hidden_dim = int(ffn_dim_multiplier * hidden_dim) |
| 239 | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| 240 | |
| 241 | self.w1 = ColumnParallelLinear( |
| 242 | dim, hidden_dim, bias=False, gather_output=False, init_method=default_linear_init |
| 243 | ) |
| 244 | self.w2 = RowParallelLinear( |
| 245 | hidden_dim, dim, bias=False, input_is_parallel=True, init_method=default_linear_init |
| 246 | ) |
| 247 | self.w3 = ColumnParallelLinear( |
| 248 | dim, hidden_dim, bias=False, gather_output=False, init_method=default_linear_init |
| 249 | ) |
| 250 | |
| 251 | # @torch.compile |
| 252 | def _silu_gating(self, x, y): |
| 253 | return F.silu(x) * y |
| 254 | |
| 255 | def forward(self, x): |
| 256 | return self.w2(self._silu_gating(self.w1(x), self.w3(x))) |
| 257 | |
| 258 | |
| 259 | class TransformerBlock(nn.Module): |