code: def gather_pipeline_parallel_state_dict(state_dict):
      gathered_states = (
          [
              None
              for _ in range(gpc.get_world_size(ParallelMode.PIPELINE))
          ]
          if gpc.get_local_rank(ParallelMode.PIPELINE) == 0
          else None
      )
      dist.gather_object(
          state_dict,
          gathered_states,
          dst=gpc.get_ranks_in_group(ParallelMode.PIPELINE)[0],
          group=gpc.get_cpu_group(ParallelMode.PIPELINE),
      )
  
      state_dict = (
          OrderedDict(
              chain.from_iterable(
                  state.items() for state in gathered_states
              )
          )
          if gpc.get_local_rank(ParallelMode.PIPELINE) == 0
          else OrderedDict()
      )
  
      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 7 for calling a global function with example usage: gather_pipeline_parallel_state_dict(state_dict) and returns: state_dict