docstr: Save unsharded optimizer state dict to checkpoint.
  After calling optimizer.state_dict(), the complete optimizer states will be collected on master rank.
  As there is communication when getting state dict, optimizer.state_dict() must be called on all processes.
  The saving process will only be executed by master rank.
docstr_tok: ["Save","unsharded","optimizer","state","dict","to","checkpoint",".","After","calling","optimizer.state_dict","(",")",",","the","complete","optimizer","states","will","be","collected","on","master","rank",".","As","there","is","communication","when","getting","state","dict",",","optimizer.state_dict","(",")","must","be","called","on","all","processes",".","The","saving","process","will","only","be","executed","by","master","rank","."]
code: def save_unsharded_optimizer(
      self,
      optimizer: GeminiOptimizer,
      checkpoint: str,
      gather_dtensor: bool,
  ):
      """
      Save unsharded optimizer state dict to checkpoint.
      After calling optimizer.state_dict(), the complete optimizer states will be collected on master rank.
      As there is communication when getting state dict, optimizer.state_dict() must be called on all processes.
      The saving process will only be executed by master rank.
      """
      assert isinstance(
          optimizer, GeminiOptimizer
      ), "Please boost the optimizer before saving!"
      state_dict = optimizer.state_dict()
      if self.coordinator.is_master():
          save_state_dict(state_dict, checkpoint, use_safetensors=False)
  
imports: import gc
  import logging
  import os
  import random
  from pathlib import Path
  from typing import Callable, Iterator, List, Optional, Tuple
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO
  from colossalai.checkpoint_io.utils import get_model_base_filenames, get_optimizer_base_filenames, load_shard_state_dict, save_config_file, save_state_dict, save_state_dict_shards
  from colossalai.cluster import DistCoordinator, ProcessGroupMesh
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
  from colossalai.zero import GeminiDDP, GeminiOptimizer
  from colossalai.zero.gemini.memory_tracer import MemStats
  from .dp_plugin_base import DPPluginBase
label_desc: Use image node_id 4 for calling the GeminiCheckpointIO obj's underlying member method code with example usage: obj.save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor) without return types