class_docstr: CheckpointIO for Hybrid Parallel Training.
  
  Args:
      dp_group (ProcessGroup): Process group along data parallel dimension.
      pp_group (ProcessGroup): Process group along pipeline parallel dimension.
      tp_group (ProcessGroup): Process group along tensor parallel dimension.
      zero_stage (int): The zero stage of plugin. Should be in [0, 1, 2].
      verbose (bool, optional): Whether to print logging massage when saving/loading has been succesfully executed. Defaults to True.
class_docstr_tok: ["CheckpointIO","for","Hybrid","Parallel","Training",".","Args",":","dp_group","(","ProcessGroup",")",":","Process","group","along","data","parallel","dimension",".","pp_group","(","ProcessGroup",")",":","Process","group","along","pipeline","parallel","dimension",".","tp_group","(","ProcessGroup",")",":","Process","group","along","tensor","parallel","dimension",".","zero_stage","(","int",")",":","The","zero","stage","of","plugin",".","Should","be","in","[","0",",","1",",","2","]",".","verbose","(","bool",",","optional",")",":","Whether","to","print","logging","massage","when","saving\/loading","has","been","succesfully","executed",".","Defaults","to","True","."]
docstr: Save model state dict to a single file with given checkpointing path.
  
  Args:
      model (nn.Module): Model on local device to be saved.
      checkpoint (str): Checkpointing path which should be a file path. Can be absolute or relative path.
      gather_dtensor (bool, optional): Whether to gather dtensor, currently not used. Defaults to True.
      use_safetensors (bool, optional): Whether to use safe tensors. Defaults to False.
docstr_tok: ["Save","model","state","dict","to","a","single","file","with","given","checkpointing","path",".","Args",":","model","(","nn.Module",")",":","Model","on","local","device","to","be","saved",".","checkpoint","(","str",")",":","Checkpointing","path","which","should","be","a","file","path",".","Can","be","absolute","or","relative","path",".","gather_dtensor","(","bool",",","optional",")",":","Whether","to","gather","dtensor",",","currently","not","used",".","Defaults","to","True",".","use_safetensors","(","bool",",","optional",")",":","Whether","to","use","safe","tensors",".","Defaults","to","False","."]
code: def save_unsharded_model(
      self,
      model: ModelWrapper,
      checkpoint: str,
      gather_dtensor: bool,
      use_safetensors: bool,
  ):
      """
      Save model state dict to a single file with given checkpointing path.
  
      Args:
          model (nn.Module): Model on local device to be saved.
          checkpoint (str): Checkpointing path which should be a file path. Can be absolute or relative path.
          gather_dtensor (bool, optional): Whether to gather dtensor, currently not used. Defaults to True.
          use_safetensors (bool, optional): Whether to use safe tensors. Defaults to False.
      """
      if self.coordinator.is_master():
          logging.warning(
              "Please avoid using unsharded checkpointing methods when dealing with large models!"
          )
  
      assert isinstance(
          model, ModelWrapper
      ), "Please boost the model before saving!"
      model = model.unwrap()
  
      if self.dp_rank != 0:
          return
  
      # The logic of collecting parameter shards along tp degree
      # has been implemented by _save_to_state_dict method of ParallelModule in Shardformer.
      state_dict = model.state_dict()
  
      if self.pp_size == 1:
          # When pipeline is not used, let master rank directly save the collected state_dict.
          if self.tp_rank == 0:
              save_state_dict(state_dict, checkpoint, use_safetensors)
      else:
          # When pipeline is used, first collect state_dict from every pipeline stage, then save the complete state_dict.
          state_dict_list = [None for _ in range(self.pp_size)]
          dist.barrier(self.pp_group)
          dist.all_gather_object(
              state_dict_list, state_dict, self.pp_group
          )
  
          # Only the master rank do the saving.
          if self.coordinator.is_master():
              complete_state_dict = dict()
              for _state_dict in state_dict_list:
                  complete_state_dict.update(_state_dict)
              save_state_dict(
                  complete_state_dict, checkpoint, use_safetensors
              )
  
imports: import copy
  from functools import reduce
  import logging
  import os
  from pathlib import Path
  from shutil import rmtree
  from typing import Dict, Iterator, Optional, OrderedDict, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed import ProcessGroup
  from torch.optim.lr_scheduler import _LRScheduler
  from colossalai.cluster import DistCoordinator
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from .general_checkpoint_io import GeneralCheckpointIO
  from .index_file import CheckpointIndexFile
  from .utils import StateDictSharder, gather_distributed_param, get_model_base_filenames, get_optimizer_base_filenames, is_safetensors_available, load_shard_state_dict, load_state_dict, load_state_dict_into_model, load_states_into_optimizer, save_config_file, save_param_groups, save_state_dict, save_state_dict_shards, search_tp_partition_dim, sharded_optimizer_loading_epilogue
label_desc: Use image node_id 8 for calling the HybridParallelCheckpointIO obj's underlying member method code with example usage: obj.save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors) without return types