class_docstr: A complete MoE MLP module in a Transformer block.
  * `activation` is the activation function to be used in MLP in each expert.
  * `d_hidden` is the dimension of the MLP layer.
class_docstr_tok: ["A","complete","MoE","MLP","module","in","a","Transformer","block",".","*","`","activation","`","is","the","activation","function","to","be","used","in","MLP","in","each","expert",".","*","`","d_hidden","`","is","the","dimension","of","the","MLP","layer","."]
code: def __init__(
      self,
      num_expert=32,
      d_model=1024,
      d_hidden=4096,
      world_size=1,
      mp_group=None,
      moe_group=None,
      activation=torch.nn.GELU(),
      gate=NaiveGate,
      top_k=2,
      expert_dp_comm="none",
      gate_hook=None,
      bias=True,
      mask=None,
      mask_dict=None,
  ):
      super().__init__(
          num_expert=num_expert,
          d_model=d_model,
          gate=gate,
          top_k=top_k,
          world_size=world_size,
          mp_group=mp_group,
          moe_group=moe_group,
          gate_hook=gate_hook,
          mask=mask,
          policy_fn=_gen_policy(d_hidden / d_model),
          mask_dict=mask_dict,
      )
      self.experts = _Expert(
          num_expert,
          d_model,
          d_hidden,
          activation,
          rank=self.mp_rank,
          bias=bias,
      )
      self.mark_parallel_comm(expert_dp_comm)
  
file_path: thu-pacman-faster-moe/fmoe/transformer.py
imports: import torch
  import torch.nn
  from .gates import NaiveGate
  from .layers import FMoE, FMoELinear
label_desc: Use image node_id 1 to create a new FMoETransformerMLP object from inherited base classes: FMoE with example: obj = FMoETransformerMLP(num_expert, d_model, d_hidden, world_size, mp_group, moe_group, activation, gate, top_k, expert_dp_comm, gate_hook, bias, mask, mask_dict)