class_docstr: 3D tensor parallel initialization among input.
  
  Args:
      num_group (int): The number of all tensor groups.
      depth (int): Depth of 3D parallelism.
      rank (int): The rank of current process.
      world_size (int): Size of whole communication world.
      config (Config): Running configuration.
      data_parallel_size (int): Size of data parallel.
      pipeline_parallel_size (int): Size of pipeline parallel.
      tensor_parallel_size (int): Size of tensor parallel.
class_docstr_tok: ["3D","tensor","parallel","initialization","among","input",".","Args",":","num_group","(","int",")",":","The","number","of","all","tensor","groups",".","depth","(","int",")",":","Depth","of","3D","parallelism",".","rank","(","int",")",":","The","rank","of","current","process",".","world_size","(","int",")",":","Size","of","whole","communication","world",".","config","(","Config",")",":","Running","configuration",".","data_parallel_size","(","int",")",":","Size","of","data","parallel",".","pipeline_parallel_size","(","int",")",":","Size","of","pipeline","parallel",".","tensor_parallel_size","(","int",")",":","Size","of","tensor","parallel","."]
docstr: Initialize 3D tensor parallel groups among input, and assign local_ranks and groups to each gpu.
  
  Returns:
      Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode):
          3D tensor parallelism's information among input in a tuple.
docstr_tok: ["Initialize","3D","tensor","parallel","groups","among","input",",","and","assign","local_ranks","and","groups","to","each","gpu",".","Returns",":","Tuple","(","local_rank",",","group_world_size",",","process_group",",","ranks_in_group",",","mode",")",":","3D","tensor","parallelism","'s","information","among","input","in","a","tuple","."]
code: def init_dist_group(self):
      """Initialize 3D tensor parallel groups among input, and assign local_ranks and groups to each gpu.
  
      Returns:
          Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode):
              3D tensor parallelism's information among input in a tuple.
      """
      local_rank = None
      ranks_in_group = None
      process_group = None
      cpu_group = None
      group_world_size = None
      mode = ParallelMode.PARALLEL_3D_INPUT
      env.input_group_3d = mode
  
      for h in range(self.num_group):
          for i in range(self.depth):
              for k in range(self.depth):
                  ranks = [
                      h * self.depth**3
                      + i
                      + self.depth * (j + self.depth * k)
                      for j in range(self.depth)
                  ]
                  group = dist.new_group(ranks)
                  group_cpu = (
                      dist.new_group(ranks, backend="gloo")
                      if dist.get_backend() != "gloo"
                      else group
                  )
  
                  if self.rank in ranks:
                      local_rank = ranks.index(self.rank)
                      group_world_size = len(ranks)
                      process_group = group
                      cpu_group = group_cpu
                      ranks_in_group = ranks
  
      return (
          local_rank,
          group_world_size,
          process_group,
          cpu_group,
          ranks_in_group,
          mode,
      )
  
imports: import math
  import torch.distributed
  from colossalai.legacy.global_variables import tensor_parallel_env
  from colossalai.legacy.registry import DIST_GROUP_INITIALIZER
  from ..parallel_mode import ParallelMode
  from .process_group_initializer import ProcessGroupInitializer
label_desc: Use image node_id 2 for calling the Initializer_3D_Input obj's underlying member method code with example usage: obj.init_dist_group() and returns: local_rank,  group_world_size,  process_group,  cpu_group,  ranks_in_group,  mode