code: def fmoefy(
      model,
      num_experts=None,
      distributed_experts=True,
      hidden_hidden_size=None,
      top_k=None,
  ):
      r"""
      Replace MLP layers in a transformer-based model in Megatron by MoE.
      * `model` should be a standard Megatron model that has
      `model.language_model.transformer.layers` as transformer layers, which is an
      array of transformer blocks that contain an `mlp` member.
      * `distributed_expert` is set to True if different experts are located in
      different workers. Otherwise, the experts on the workers are identical, and
      they are trained in data-parallel mode. This can be useful when testing on
      small models that do not require high training throughput or large parameter
      capacity.
      Note that pipeline parallel is not supported yet. When distributed experts
      are enabled, their communicator should be Megatron's
      tensor_model_parall_comm x data_parallel_comm, which is not created.
      """
      from megatron import get_args
      from megatron import mpu
  
      args = get_args()
      if num_experts is not None:
          args.num_experts = num_experts
      assert (
          "num_experts" in args
      ), "num_experts should be specified in arguments or fmoefy function"
  
      if hidden_hidden_size is not None:
          args.hidden_hidden_size = hidden_hidden_size
      elif not hasattr(args, "hidden_hidden_size"):
          args.hidden_hidden_size = args.hidden_size * 4
  
      if top_k is not None:
          args.top_k = top_k
      elif not hasattr(args, "top_k"):
          args.top_k = 2
  
      # Set distributed_experts to None to use default setting in args
      if distributed_experts is not None:
          args.distributed_experts = distributed_experts
  
      if hasattr(mpu, "get_tensor_model_parallel_group"):
          mp_group = mpu.get_tensor_model_parallel_group()
      else:
          # For compatibility to older versions of Megatron-LM
          mp_group = mpu.get_model_parallel_group()
      if args.pipeline_model_parallel_size == 1:
          moe_group = get_torch_default_comm()
      else:
          # Create a comm prependicular to pipeline group
          stage_size = (
              args.world_size // args.pipeline_model_parallel_size
          )
          for i in range(0, args.world_size, stage_size):
              ranks = range(i, i + stage_size)
              group = torch.distributed.new_group(ranks)
              if args.rank in ranks:
                  moe_group = group
          set_moe_group(moe_group)
      for idx, l in enumerate(model.language_model.transformer.layers):
          if args.balance_strategy == "baseorig":
              l.mlp = BaseMLP(args)
          else:
              l.mlp = MegatronMLP(args, mp_group, moe_group, idx)
  
      # initialize gate hook
      num_layers = len(model.language_model.transformer.layers)
      reset_gate_hook(num_layers)
  
      return model
  
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 3 for calling a global function with example usage: fmoefy(model, num_experts, distributed_experts, hidden_hidden_size, top_k) and returns: model