code: def gather_tensor_parallel_state_dict(
      state_dict: OrderedDict,
      parallel_mode: ParallelMode,
      dims: dict = dict(),
      partition_states: dict = dict(),
      keep_vars: bool = False,
  ):
      dst_rank = gpc.get_ranks_in_group(parallel_mode)[0]
      depth = gpc.get_world_size(parallel_mode)
  
      for key in list(state_dict.keys()):
          param = state_dict.pop(key)
          param = param if keep_vars else param.detach()
          dim = dims.get(key, 0)
          do_partition = partition_states.get(key, True)
          if do_partition:
              temp = param.transpose(0, dim).contiguous()
              gather_list = None
              if gpc.get_local_rank(parallel_mode) == 0:
                  shape = list(param.shape)
                  shape[0], shape[dim] = shape[dim], shape[0]
                  shape[0] *= depth
                  param = torch.empty(
                      shape, dtype=param.dtype, device=param.device
                  )
                  gather_list = list(torch.chunk(param, depth, dim=0))
              dist.gather(
                  temp,
                  gather_list,
                  dst=dst_rank,
                  group=gpc.get_cpu_group(parallel_mode),
              )
              param = torch.transpose(param, 0, dim)
          # update params in state_dict only on local rank 0
          if gpc.get_local_rank(parallel_mode) == 0:
              state_dict[key] = param
  
      return state_dict
  
file_path: ColossalAI/colossalai/legacy/utils/checkpointing.py
imports: from collections import OrderedDict
  from itertools import chain
  import torch
  import torch.distributed
  from colossalai.legacy.constants import IS_TENSOR_PARALLEL
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from .common import is_using_pp
label_desc: Use image node_id 3 for calling a global function with example usage: gather_tensor_parallel_state_dict(state_dict, parallel_mode, dims, partition_states, keep_vars) and returns: state_dict