code: def _get_mlp_module_spec(
      use_te: bool = True,
      num_experts: int = None,
      moe_grouped_gemm: bool = False,
  ) -> ModuleSpec:
      if num_experts is None:
          # Dense MLP w/ or w/o TE modules.
          return ModuleSpec(
              module=MLP,
              submodules=MLPSubmodules(
                  linear_fc1=TELayerNormColumnParallelLinear
                  if use_te
                  else ColumnParallelLinear,
                  linear_fc2=TERowParallelLinear
                  if use_te
                  else RowParallelLinear,
              ),
          )
      elif moe_grouped_gemm:
          # GroupedMLP based MoE with modules in megatron core.
          return GroupedMLP
      else:
          # SwitchMLP based MoE with modules in megatron core.
          return ModuleSpec(
              module=SwitchMLP,
              submodules=MLPSubmodules(
                  linear_fc1=ColumnParallelLinear,
                  linear_fc2=RowParallelLinear,
              ),
          )
  
file_path: megatron-lm/megatron/core/models/gpt/gpt_layer_specs.py
imports: from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
  from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
  from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
  from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
  from megatron.core.transformer.custom_layers.transformer_engine import TEDotProductAttention, TELayerNormColumnParallelLinear, TENorm, TERowParallelLinear
  from megatron.core.transformer.dot_product_attention import DotProductAttention
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityOp
  from megatron.core.transformer.mlp import MLP, MLPSubmodules
  from megatron.core.transformer.moe.grouped_mlp import GroupedMLP
  from megatron.core.transformer.moe.switch_mlp import SwitchMLP
  from megatron.core.transformer.spec_utils import ModuleSpec
  from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
label_desc: Use image node_id 3 for calling a global function with example usage: _get_mlp_module_spec(use_te, num_experts, moe_grouped_gemm) and returns: ModuleSpec, GroupedMLP, ModuleSpec