code: def __init__(self, tp_world_size):
      self.rank = deepspeed.comm.get_rank()
      self.world_size = deepspeed.comm.get_world_size()
      self.tp_world_size = tp_world_size
  
      for i in range(0, self.world_size, tp_world_size):
          ranks = range(i, i + tp_world_size)
          group = deepspeed.comm.new_group(ranks)
          if self.rank in ranks:
              self.tp_group = group
  
      for i in range(0, tp_world_size):
          ranks = range(i, self.world_size, tp_world_size)
          group = deepspeed.comm.new_group(ranks)
          if self.rank in ranks:
              self.dp_group = group
  
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 to create a new MPU object  with example: obj = MPU(tp_world_size)