code: def get_mpu_ranks(
      tp_size=1, pp_size=1, dp_size=1, virtual_pp_size=None
  ):
      """
      Initialize model data parallel groups.
  
      Arguments:
          tp_size: number of GPUs used to parallelize model tensor.
          pp_size: number of GPUs used to parallelize model pipeline.
          dp_size: number of GPUs used to parallelize model data.
  
      Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
      use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
      the model pipeline. The present function will
      create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
      and 8 data-parallel groups as:
          8 data_parallel groups:
              [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
          8 tensor model-parallel groups:
              [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
          4 pipeline model-parallel groups:
              [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
      Note that for efficiency, the caller should make sure adjacent ranks
      are on the same DGX box. For example if we are using 2 DGX-1 boxes
      with a total of 16 GPUs, rank 0 to 7 belong to the first box and
      ranks 8 to 15 belong to the second box.
      """
  
      world_size = tp_size * pp_size * dp_size
  
      print(
          f"\n\n*** tp={tp_size}, pp={pp_size}, dp={dp_size}, world={world_size}"
      )
  
      tensor_model_parallel_size = min(tp_size, world_size)
      pipeline_model_parallel_size = min(pp_size, world_size)
      data_parallel_size = world_size // (
          tensor_model_parallel_size * pipeline_model_parallel_size
      )
  
      num_tensor_model_parallel_groups = (
          world_size // tensor_model_parallel_size
      )
      num_pipeline_model_parallel_groups = (
          world_size // pipeline_model_parallel_size
      )
      num_data_parallel_groups = world_size // data_parallel_size
  
      # Build the data-parallel groups.
      all_dp_group_ranks = []
      for i in range(pipeline_model_parallel_size):
          start_rank = i * num_pipeline_model_parallel_groups
          end_rank = (i + 1) * num_pipeline_model_parallel_groups
          for j in range(tensor_model_parallel_size):
              ranks = range(
                  start_rank + j, end_rank, tensor_model_parallel_size
              )
              all_dp_group_ranks.append(list(ranks))
  
      print("DP", all_dp_group_ranks)
  
      # Build the model-parallel groups.
      all_pp_group_ranks = []
      for i in range(data_parallel_size):
          ranks = [
              data_parallel_group_ranks[i]
              for data_parallel_group_ranks in all_dp_group_ranks
          ]
          all_pp_group_ranks.append(list(ranks))
  
      print(f"PP", all_pp_group_ranks)
  
      # Build the tensor model-parallel groups.
      all_tp_group_ranks = []
      for i in range(num_tensor_model_parallel_groups):
          ranks = range(
              i * tensor_model_parallel_size,
              (i + 1) * tensor_model_parallel_size,
          )
          all_tp_group_ranks.append(list(ranks))
  
      print(f"TP", all_tp_group_ranks)
  
      return all_tp_group_ranks, all_pp_group_ranks, all_dp_group_ranks
  
file_path: deepspeed/deepspeed/checkpoint/reshape_meg_2d.py
imports: from .reshape_utils import partition_data
label_desc: Use image node_id 4 for calling a global function with example usage: get_mpu_ranks(tp_size, pp_size, dp_size, virtual_pp_size) and returns: all_tp_group_ranks,  all_pp_group_ranks,  all_dp_group_ranks