docstr: Save sharded optimizer state dict to checkpoint folder.
  As there is communication when getting state dict, this must be called on all processes.
docstr_tok: ["Save","sharded","optimizer","state","dict","to","checkpoint","folder",".","As","there","is","communication","when","getting","state","dict",",","this","must","be","called","on","all","processes","."]
code: def save_sharded_optimizer(
      self,
      optimizer: GeminiOptimizer,
      checkpoint: Path,
      gather_dtensor: bool,
      prefix: str,
      size_per_shard: int,
  ):
      """
      Save sharded optimizer state dict to checkpoint folder.
      As there is communication when getting state dict, this must be called on all processes.
      """
      assert isinstance(
          optimizer, GeminiOptimizer
      ), "Please boost the optimizer before saving!"
  
      if os.path.isfile(checkpoint):
          logging.error(
              f"Provided path ({checkpoint}) should be a directory, not a file"
          )
          return
  
      Path(checkpoint).mkdir(parents=True, exist_ok=True)
  
      # Preparing file paths and index file.
      (
          states_name,
          save_index_file,
          param_group_file,
      ) = get_optimizer_base_filenames(prefix)
      index_file = CheckpointIndexFile(checkpoint)
      index_file.append_meta_data("param_groups", param_group_file)
  
      # Store the information of param groups to param_group_file.
      if self.coordinator.is_master():
          group_file_path = os.path.join(checkpoint, param_group_file)
          param_groups = optimizer.get_param_groups_for_saving()
          torch.save(param_groups, group_file_path)
  
      # States are broken into shards within max_shard_size.
      state_dict_shard = optimizer.state_shard(
          prefix=prefix, max_shard_size=size_per_shard, only_rank_0=True
      )
  
      # Save shards of optimizer states.
      total_size = save_state_dict_shards(
          sharded_state_dict=state_dict_shard,
          checkpoint=checkpoint,
          index_file=index_file,
          base_filename=states_name,
          is_master=self.coordinator.is_master(),
          use_safetensors=False,
      )
  
      # Wrap up index file. Only save it on master rank.
      if self.coordinator.is_master():
          index_file.append_meta_data("total_size", total_size)
          index_file.write_index_file(save_index_file)
          logging.info(
              f"The optimizer is going to be split to checkpoint shards. "
              f"You can find where each parameters has been saved in the "
              f"index located at {save_index_file}."
          )
  
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 8 for calling the GeminiCheckpointIO obj's underlying member method code with example usage: obj.save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard) without return types