code: def save_state_dict_shards(
      sharded_state_dict: Iterator[Tuple[OrderedDict, int]],
      checkpoint: str,
      index_file: "CheckpointIndexFile",
      base_filename: str,
      is_master: bool,
      use_safetensors: bool = False,
      use_pp_format: bool = False,
  ) -> int:
      """
      Save sharded state dict only on master rank, this method can be used by both model and optimizer states.
      Args:
          sharded_state_dict (Iterator[Tuple[OrderedDict, int]]): a generator of shards, each shard contains state dict and shard size.
          checkpoint (str): The path of checkpoint directory as string.
          index_file (CheckpointIndexFile): The index file object to be updated.
          base_filename (str): Decides the prefix of filenames of shards.
          is_master (bool): Whether current rank is main process.
          use_safetensors (bool, optional): Whether to use safetensors to save checkpoint. Defaults to False.
          use_pp_format: (bool, optional): Whether to save the files in pipeline format including stage information. Defaults to False.
  
      Returns:
          int: the total size of shards
      """
  
      total_size = 0
      shard_filenames = []
      for idx, shard_pair in enumerate(sharded_state_dict):
          shard, current_size = shard_pair
          if not is_master:
              del shard
              continue
          shard_file = get_shard_filename(base_filename, idx)
          total_size = total_size + current_size
          for key in shard.keys():
              index_file.append_weight_map(key, shard_file)
          checkpoint_file_path = os.path.join(checkpoint, shard_file)
  
          # Only save on master rank.
          save_state_dict(
              shard,
              checkpoint_file_path,
              use_safetensors=use_safetensors,
          )
          shard_filenames.append(shard_file)
          del shard
  
      # Clean folder, deleted unneeded files.
      clean_folder(
          checkpoint,
          base_filename,
          shard_filenames,
          is_master=is_master,
          use_pp_format=use_pp_format,
      )
  
      return total_size
  
imports: import os
  import re
  from collections import abc
  from collections import defaultdict
  from itertools import chain
  from pathlib import Path
  from typing import Iterator, List, Mapping, Optional, OrderedDict, Tuple
  import torch
  import torch.nn
  from packaging.version import Version
  from torch.optim import Optimizer
  from torch.utils._pytree import tree_map
  from colossalai.tensor.d_tensor import is_customized_distributed_tensor, is_distributed_tensor, to_global, to_global_for_customized_distributed_tensor
label_desc: Use image node_id 7 for calling a global function with example usage: save_state_dict_shards(sharded_state_dict, checkpoint, index_file, base_filename, is_master, use_safetensors, use_pp_format) and returns: total_size