docstr: 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 (float, optional): Custom multiplier for hidden dimension. Defaults to None.
  
  Attributes:
      w1 (ColumnParallelLinear): Linear transformation for the first layer.
      w2 (RowParallelLinear): Linear transformation for the second layer.
      w3 (ColumnParallelLinear): Linear transformation for the third layer.
code: def __init__(
      self,
      dim: int,
      hidden_dim: int,
      multiple_of: int,
      ffn_dim_multiplier: Optional[float],
  ):
      """
      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 (float, optional): Custom multiplier for hidden dimension. Defaults to None.
  
      Attributes:
          w1 (ColumnParallelLinear): Linear transformation for the first layer.
          w2 (RowParallelLinear): Linear transformation for the second layer.
          w3 (ColumnParallelLinear): Linear transformation for the third layer.
  
      """
      super().__init__()
      hidden_dim = int(2 * hidden_dim / 3)
      # custom dim factor multiplier
      if ffn_dim_multiplier is not None:
          hidden_dim = int(ffn_dim_multiplier * hidden_dim)
      hidden_dim = multiple_of * (
          (hidden_dim + multiple_of - 1) // multiple_of
      )
  
      from fairscale.nn.model_parallel.layers import (
          ColumnParallelLinear,
          RowParallelLinear,
      )
  
      self.w1 = ColumnParallelLinear(
          dim,
          hidden_dim,
          bias=False,
          gather_output=False,
          init_method=lambda x: x,
      )
      self.w2 = RowParallelLinear(
          hidden_dim,
          dim,
          bias=False,
          input_is_parallel=True,
          init_method=lambda x: x,
      )
      self.w3 = ColumnParallelLinear(
          dim,
          hidden_dim,
          bias=False,
          gather_output=False,
          init_method=lambda x: x,
      )
  
file_path: mixtral-kit/mixtralkit/layers/ffn.py
imports: import torch
  import torch.nn.functional
  from torch import nn
  from typing import Optional, Tuple
label_desc: Use image node_id 1 to create a new FairScaleFFN object from inherited base classes: nn with example: obj = FairScaleFFN(dim, hidden_dim, multiple_of, ffn_dim_multiplier)