class_docstr: 2d tensor parallel initialization among cols.
  
  Args:
      num_group (int): The number of all tensor groups.
      summa_dim (int): The dimension of SUMMA.
      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: ["2d","tensor","parallel","initialization","among","cols",".","Args",":","num_group","(","int",")",":","The","number","of","all","tensor","groups",".","summa_dim","(","int",")",":","The","dimension","of","SUMMA",".","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 2D tensor row parallel groups, and assign local_ranks and groups to each gpu.
  
  Returns:
      Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode):
          2D tensor col parallelism's information in a tuple.
docstr_tok: ["Initialize","2D","tensor","row","parallel","groups",",","and","assign","local_ranks","and","groups","to","each","gpu",".","Returns",":","Tuple","(","local_rank",",","group_world_size",",","process_group",",","ranks_in_group",",","mode",")",":","2D","tensor","col","parallelism","'s","information","in","a","tuple","."]
code: def init_dist_group(self):
      """Initialize 2D tensor row parallel groups, and assign local_ranks and groups to each gpu.
  
      Returns:
          Tuple (local_rank, group_world_size, process_group, ranks_in_group, mode):
              2D tensor col parallelism's information in a tuple.
      """
      local_rank = None
      ranks_in_group = None
      process_group = None
      cpu_group = None
      group_world_size = None
      mode = ParallelMode.PARALLEL_2D_COL
  
      for i in range(self.num_group):
          for j in range(self.summa_dim):
              ranks = [
                  i * self.tensor_parallel_size + j + k * self.summa_dim
                  for k in range(self.summa_dim)
              ]
              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,
      )
  
file_path: ColossalAI/colossalai/legacy/context/process_group_initializer/initializer_2d.py
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_2D_Col 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