code: def _init_group(self):
      world_size = gpc.get_world_size(ParallelMode.GLOBAL)
      dp_size = gpc.get_world_size(ParallelMode.DATA)
      pp_size = gpc.get_world_size(ParallelMode.PIPELINE)
      rank = gpc.get_global_rank()
      num_dp_groups = world_size // dp_size
      num_pp_stages = num_dp_groups // pp_size
      for i in range(dp_size):
          for j in range(num_pp_stages):
              pipeline_ranks = list(
                  range(
                      i * num_dp_groups + j,
                      (i + 1) * num_dp_groups,
                      num_pp_stages,
                  )
              )
              sub_ranks = [
                  pipeline_ranks[idx] for idx in self.pipeline_ranks
              ]
              group = dist.new_group(sub_ranks)
              if rank in sub_ranks:
                  self.group = group
                  self.ranks_in_group = sub_ranks
  
file_path: ColossalAI/colossalai/legacy/nn/layer/wrapper/pipeline_wrapper.py
imports: from typing import List, Tuple, Union
  import torch.distributed
  import torch.nn
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
label_desc: Use image node_id 2 for calling the PipelineSharedModuleWrapper obj's underlying member method code with example usage: obj._init_group() without return types