code: def setup_method(
      self, method, use_cpu_initialization=False, swiglu=True
  ):
      print("============")
      print(
          "Test for use_cpu_initilization={} and swiglu={}.".format(
              use_cpu_initialization, swiglu
          )
      )
      print("============")
      Utils.initialize_model_parallel(1, 1)
      num_layers = 1  # 2
      self.hidden_size = 2  # 12
      self.num_experts = 2
      self.gated_linear_unit = True
      self.use_cpu_initialization = use_cpu_initialization
      self.gated_linear_unit = False
      if swiglu:
          self.gated_linear_unit = True
  
      tf_config = TransformerConfig(
          num_layers=num_layers,
          hidden_size=self.hidden_size,
          num_attention_heads=4,
          num_moe_experts=self.num_experts,
          use_cpu_initialization=self.use_cpu_initialization,
          add_bias_linear=False,
          gated_linear_unit=self.gated_linear_unit,
          bias_gelu_fusion=False,
          bf16=True,
          params_dtype=torch.bfloat16,
      )
  
      self.fc1_ffn_hidden_size = tf_config.ffn_hidden_size
      self.fc2_ffn_hidden_size = tf_config.ffn_hidden_size
      # If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
      if self.gated_linear_unit:
          self.fc1_ffn_hidden_size *= 2
  
      ## Vanilla sequential GEMM
      # Set random seed for reproducability
      _set_random_seed(seed_=123, data_parallel_random_init=False)
      transformer_layer_spec = (
          get_gpt_layer_with_transformer_engine_spec(
              self.num_experts, moe_grouped_gemm=False
          )
      )
      self.switch_mlp_smm = SwitchMLP(
          tf_config, transformer_layer_spec.submodules.mlp.submodules
      )
  
      self.args = parse_args(ignore_unknown_args=True)
      self.args.bf16 = True
      # Bias is not supported in grouped gemm currently, thus we disable the
      # bias in the linear layer.
      self.args.add_bias_linear = False
      self.switch_mlp_smm = Float16Module(
          self.switch_mlp_smm, self.args
      ).module
      print("done intializing for sequential gemm")
  
      ## Grouped GEMM
      _set_random_seed(seed_=123, data_parallel_random_init=False)
      self.switch_mlp_gmm = GroupedMLP(tf_config)
      self.switch_mlp_gmm = Float16Module(
          self.switch_mlp_gmm, self.args
      ).module
      print("done intializing for grouped gemm")
  
file_path: megatron-lm/tests/unit_tests/transformer/moe/test_grouped_mlp.py
imports: import pytest
  import torch
  import torch.nn.functional
  from megatron.arguments import parse_args
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
  from megatron.core.transformer.moe.grouped_mlp import GroupedMLP
  from megatron.core.transformer.moe.switch_mlp import SwitchMLP
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.initialize import _set_random_seed
  from megatron.model import Float16Module
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 1 for calling the TestParallelGroupedMLP obj's underlying member method code with example usage: obj.setup_method(method, use_cpu_initialization, swiglu) without return types