code: def test(self, ep_size, tp_size, enable_expert_tp, use_residual):
      # TODO: replace this with a true parallel mlp in the future
      # and run convergence tests
      if not required_torch_version(min_version=1.8):
          pytest.skip(
              "DeepSpeed MoE tests need torch 1.8 or higher to run correctly"
          )
  
      config_dict = {
          "train_batch_size": 8,
          "steps_per_print": 1,
          "fp16": {"enabled": True},
      }
      hidden_dim = 16
  
      tensor_parallel_expert = torch.nn.Sequential(
          torch.nn.Linear(hidden_dim, 4 * hidden_dim // tp_size),
          torch.nn.ReLU(),
          torch.nn.Linear(4 * hidden_dim // tp_size, hidden_dim),
      )
  
      # set num experts to world size
      world_size = deepspeed.comm.get_world_size()
      model = MoE(
          hidden_size=hidden_dim,
          expert=tensor_parallel_expert,
          num_experts=world_size,
          ep_size=ep_size,
          use_residual=use_residual,
          enable_expert_tensor_parallelism=enable_expert_tp,
      )
      optimizer = torch.optim.AdamW(params=model.parameters())
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          optimizer=optimizer,
          dist_init_required=False,
          mpu=MPU(tp_size),
      )
  
      assert model.num_local_experts == world_size // ep_size
      if enable_expert_tp:
          assert (
              deepspeed.utils.groups._get_expert_model_parallel_world_size()
              == tp_size
          )
      else:
          assert (
              deepspeed.utils.groups._get_expert_model_parallel_world_size()
              == 1
          )
  
file_path: DeepSpeed/tests/unit/moe/test_moe_tp.py
imports: import torch
  import deepspeed
  import pytest
  from unit.common import DistributedTest
  from deepspeed.runtime.utils import required_torch_version
  from deepspeed.moe.layer import MoE
label_desc: Use image node_id 1 for calling the TestMOETensorParallel obj's underlying member method code with example usage: obj.test(ep_size, tp_size, enable_expert_tp, use_residual) without return types