code: def partition_tensor_parallel_state_dict(
      state_dict: OrderedDict,
      parallel_mode: ParallelMode,
      dims: dict = dict(),
      partition_states: dict = dict(),
  ):
      src_rank = gpc.get_ranks_in_group(parallel_mode)[0]
      depth = gpc.get_world_size(parallel_mode)
      group = gpc.get_cpu_group(parallel_mode)
      is_rank0 = gpc.get_local_rank(parallel_mode) == 0
      partition_info = [None]
      if is_rank0:
          partition_info_dict = OrderedDict()
          for key, param in state_dict.items():
              dim = dims[key]
              is_partitioned = partition_states[key]
              shape = list(param.shape)
              if is_partitioned:
                  shape[dim] = shape[dim] // depth
              partition_info_dict[key] = (
                  is_partitioned,
                  param.dtype,
                  shape,
                  dim,
              )
          partition_info[0] = partition_info_dict
      dist.broadcast_object_list(partition_info, src_rank, group=group)
      partitioned_state = OrderedDict()
      for key, (is_partitioned, dtype, shape, dim) in partition_info[
          0
      ].items():
          if is_partitioned:
              output = torch.empty(shape, dtype=dtype)
              if is_rank0:
                  scatter_list = [
                      t.contiguous()
                      for t in state_dict[key].chunk(depth, dim)
                  ]
              else:
                  scatter_list = None
              dist.scatter(output, scatter_list, src_rank, group=group)
          else:
              if is_rank0:
                  output = state_dict[key]
              else:
                  output = torch.empty(shape, dtype=dtype)
              dist.broadcast(output, src_rank, group=group)
          partitioned_state[key] = output
      return partitioned_state
  
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 2 for calling a global function with example usage: partition_tensor_parallel_state_dict(state_dict, parallel_mode, dims, partition_states) and returns: partitioned_state