docstr: Save sharded model.
  As there is communication when getting state dict, model.state_dict() must be called on all processes.
docstr_tok: ["Save","sharded","model",".","As","there","is","communication","when","getting","state","dict",",","model.state_dict","(",")","must","be","called","on","all","processes","."]
code: def save_sharded_model(
      self,
      model: GeminiDDP,
      checkpoint_path: str,
      gather_dtensor: bool = False,
      prefix: Optional[str] = None,
      max_shard_size: int = 1024,
      use_safetensors: bool = False,
  ):
      """
      Save sharded model.
      As there is communication when getting state dict, model.state_dict() must be called on all processes.
      """
      assert isinstance(
          model, GeminiDDP
      ), "Please boost the model before saving!"
      if os.path.isfile(checkpoint_path):
          logging.error(
              f"Provided path ({checkpoint_path}) should be a directory, not a file"
          )
          return
  
      Path(checkpoint_path).mkdir(parents=True, exist_ok=True)
  
      state_dict_shard = model.state_dict_shard(
          max_shard_size=max_shard_size, only_rank_0=True
      )
      weights_name, save_index_file = get_model_base_filenames(
          prefix, use_safetensors
      )
      index_file = CheckpointIndexFile(checkpoint_path)
  
      # Save shards of optimizer states.
      is_master = self.coordinator.is_master()
      total_size = save_state_dict_shards(
          sharded_state_dict=state_dict_shard,
          checkpoint=checkpoint_path,
          index_file=index_file,
          base_filename=weights_name,
          is_master=is_master,
          use_safetensors=use_safetensors,
      )
  
      # only save the index file on the master rank
      if self.coordinator.is_master():
          index_file.append_meta_data("total_size", total_size)
          index_file.write_index_file(save_index_file)
          save_config_file(model.unwrap(), checkpoint_path)
          logging.info(
              f"The model is split into 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 6 for calling the GeminiCheckpointIO obj's underlying member method code with example usage: obj.save_sharded_model(model, checkpoint_path, gather_dtensor, prefix, max_shard_size, use_safetensors) without return types