Initialize the FeedForward module. Args: dim (int): Input dimension. hidden_dim (int): Hidden dimension of the feedforward layer. multiple_of (int): Value to ensure hidden dimension is a multiple of this value. ffn_dim_multiplier (flo
(
self,
dim: int,
hidden_dim: int,
multiple_of: int,
ffn_dim_multiplier: Optional[float],
)
| 306 | |
| 307 | class FeedForward(nn.Module): |
| 308 | def __init__( |
| 309 | self, |
| 310 | dim: int, |
| 311 | hidden_dim: int, |
| 312 | multiple_of: int, |
| 313 | ffn_dim_multiplier: Optional[float], |
| 314 | ): |
| 315 | """ |
| 316 | Initialize the FeedForward module. |
| 317 | |
| 318 | Args: |
| 319 | dim (int): Input dimension. |
| 320 | hidden_dim (int): Hidden dimension of the feedforward layer. |
| 321 | multiple_of (int): Value to ensure hidden dimension is a multiple of this value. |
| 322 | ffn_dim_multiplier (float, optional): Custom multiplier for hidden dimension. Defaults to None. |
| 323 | |
| 324 | Attributes: |
| 325 | w1 (ColumnParallelLinear): Linear transformation for the first layer. |
| 326 | w2 (RowParallelLinear): Linear transformation for the second layer. |
| 327 | w3 (ColumnParallelLinear): Linear transformation for the third layer. |
| 328 | |
| 329 | """ |
| 330 | super().__init__() |
| 331 | hidden_dim = int(2 * hidden_dim / 3) |
| 332 | # custom dim factor multiplier |
| 333 | if ffn_dim_multiplier is not None: |
| 334 | hidden_dim = int(ffn_dim_multiplier * hidden_dim) |
| 335 | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| 336 | |
| 337 | self.w1 = ColumnParallelLinear( |
| 338 | dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x |
| 339 | ) |
| 340 | self.w2 = RowParallelLinear( |
| 341 | hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x |
| 342 | ) |
| 343 | self.w3 = ColumnParallelLinear( |
| 344 | dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x |
| 345 | ) |
| 346 | |
| 347 | def forward(self, x): |
| 348 | return self.w2(F.silu(self.w1(x)) * self.w3(x)) |