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 optimizer state dict to a file with given path.
  
  Args:
      optimizer (OptimizerWrapper): Optimizer to save sharded state_dict.
      checkpoint (str): Path to save optimizer state_dict.
      gather_dtensor (bool): Whether to gather_dtensor, not used.
docstr_tok: ["Save","optimizer","state","dict","to","a","file","with","given","path",".","Args",":","optimizer","(","OptimizerWrapper",")",":","Optimizer","to","save","sharded","state_dict",".","checkpoint","(","str",")",":","Path","to","save","optimizer","state_dict",".","gather_dtensor","(","bool",")",":","Whether","to","gather_dtensor",",","not","used","."]
code: def save_unsharded_optimizer(
      self,
      optimizer: OptimizerWrapper,
      checkpoint: str,
      gather_dtensor: bool,
  ):
      """
      Save optimizer state dict to a file with given path.
  
      Args:
          optimizer (OptimizerWrapper): Optimizer to save sharded state_dict.
          checkpoint (str): Path to save optimizer state_dict.
          gather_dtensor (bool): Whether to gather_dtensor, not used.
      """
      if self.coordinator.is_master():
          logging.warning(
              "Please avoid using unsharded checkpointing methods when dealing with large models!"
          )
  
      assert isinstance(
          optimizer, OptimizerWrapper
      ), "Please boost the optimizer before saving!"
  
      # optimizer states of parameters kept by local device('s pipeline stage)
      local_states = dict()
  
      for param, state in optimizer.optim.state.items():
          if param is None:
              continue
  
          # working param is needed for obtaining correct param_id
          master_to_working_map = optimizer.get_master_to_working_map()
          if master_to_working_map is not None:
              working_param = master_to_working_map[id(param)]
          else:
              working_param = param
  
          # gather complete state from tp shards & dp shards
          param_id = optimizer.param_info["param2id"][id(working_param)]
          original_shape = optimizer.param_info["param2shape"][
              id(working_param)
          ]
          local_states[
              param_id
          ] = HybridParallelCheckpointIO.gather_from_sharded_optimizer_state(
              state,
              working_param,
              original_shape=original_shape,
              dp_group=self.dp_group,
              tp_group=self.tp_group,
              use_zero=self.use_zero,
              inplace=False,
              device=torch.device("cuda"),
          )
  
      if self.pp_size == 1:
          # When pipeline is not used, let master rank directly save the collected state_dict.
          state_dict = {
              "param_groups": optimizer.param_info["param_groups"],
              "state": local_states,
          }
          if self.coordinator.is_master():
              save_state_dict(
                  state_dict, checkpoint, use_safetensors=False
              )
      else:
          # When pipeline is used, first collect state_dict from every pipeline stage, then save the complete state_dict.
          states_list = [None for _ in range(self.pp_size)]
          dist.barrier(self.pp_group)
          dist.all_gather_object(
              states_list, local_states, self.pp_group
          )
  
          # Only the master rank do the saving.
          if self.coordinator.is_master():
              state_dict = {
                  "param_groups": optimizer.param_info["param_groups"],
                  "state": dict(),
              }
              for _states in states_list:
                  state_dict["state"].update(_states)
              save_state_dict(
                  state_dict, checkpoint, use_safetensors=False
              )
  
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 10 for calling the HybridParallelCheckpointIO obj's underlying member method code with example usage: obj.save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor) without return types