code: def partition_pipeline_parallel_state_dict(model, state_dict):
      pipeline_state = OrderedDict()
  
      if gpc.get_local_rank(ParallelMode.TENSOR) == 0:
          # receive all states from prev stage
          if not gpc.is_first_rank(ParallelMode.PIPELINE):
              state_dict = _recv_state_dict(
                  gpc.get_prev_global_rank(ParallelMode.PIPELINE),
                  ParallelMode.PIPELINE,
              )
          # move states to output
          for name, _ in model.named_parameters(recurse=True):
              if name in state_dict:
                  pipeline_state[name] = state_dict.pop(name)
          for name, _ in model.named_buffers(recurse=True):
              if name in state_dict:
                  pipeline_state[name] = state_dict.pop(name)
          for name, _ in model.named_modules():
              extra_state_key = name + "." + _EXTRA_STATE_KEY_SUFFIX
              if extra_state_key in state_dict:
                  pipeline_state[extra_state_key] = state_dict.pop(
                      extra_state_key
                  )
          # send rest states to next stage
          if not gpc.is_last_rank(ParallelMode.PIPELINE):
              _send_state_dict(
                  state_dict,
                  gpc.get_next_global_rank(ParallelMode.PIPELINE),
                  ParallelMode.PIPELINE,
              )
  
      return pipeline_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 6 for calling a global function with example usage: partition_pipeline_parallel_state_dict(model, state_dict) and returns: pipeline_state