class_docstr: Make the FMoETransformerMLP layer that distributes experts across
  communication group `group` to replace the original MLP layer in Megatron.
class_docstr_tok: ["Make","the","FMoETransformerMLP","layer","that","distributes","experts","across","communication","group","`","group","`","to","replace","the","original","MLP","layer","in","Megatron","."]
code: def __init__(self, args, mp_group, moe_group, layer_idx):
      assert (
          args.seq_length
          * args.micro_batch_size
          % args.tensor_model_parallel_size
          == 0
      ), "Batch size x sequence length should be multiple of mp size"
      if not args.distributed_experts:
          world_size = 1
      else:
          world_size = (
              args.tensor_model_parallel_size * args.data_parallel_size
          )
      gate = None
      if not args.balance_strategy or args.balance_strategy == "naive":
          from fmoe.gates import NaiveGate
  
          gate = NaiveGate
      elif args.balance_strategy == "noisy":
          from fmoe.gates import NoisyGate
  
          gate = NoisyGate
      elif args.balance_strategy == "gshard":
          from fmoe.gates import GShardGate
  
          gate = GShardGate
      elif args.balance_strategy == "switch":
          from fmoe.gates import SwitchGate
  
          gate = SwitchGate
      elif args.balance_strategy == "neighbor":
          from fmoe.gates.neighbor_gate import gen_neighbor_gate
  
          gate = gen_neighbor_gate(moe_group.rank())
      elif args.balance_strategy == "baselayer":
          from fmoe.gates.base_layer_gate import BaseLayerGate
  
          gate = BaseLayerGate
      elif args.balance_strategy == "hir":
          from fmoe.gates.hir_gate import gen_hir_gate
  
          gate = gen_hir_gate(moe_group.rank())
      else:
          assert False, "Undefined balance strategy {}" % (
              args.balance_strategy
          )
  
      super().__init__(
          args.num_experts,
          top_k=args.top_k,
          d_model=args.hidden_size,
          d_hidden=args.hidden_hidden_size,
          world_size=world_size,
          mp_group=mp_group,
          moe_group=moe_group,
          expert_dp_comm="none" if args.distributed_experts else "dp",
          gate_hook=generate_megatron_gate_hook(
              layer_idx, args.num_experts * world_size
          ),
          gate=gate,
      )
      self.hidden_size = args.hidden_size
      if args.distributed_experts:
          self.rank = args.rank
      else:
          self.rank = 0
      self.sigma = args.init_method_std
      self.num_layers = args.num_layers
      self.reset_parameters()
  
imports: import math
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from fmoe.transformer import FMoETransformerMLP
  from fmoe.utils import get_torch_default_comm
  from .balance import reset_gate_hook
  from .balance import generate_megatron_gate_hook
  from .distributed import set_moe_group
label_desc: Use image node_id 1 to create a new MegatronMLP object from inherited base classes: FMoETransformerMLP with example: obj = MegatronMLP(args, mp_group, moe_group, layer_idx)