docstr: Loading sharded optimizer from checkpoint folder, with index file given.
  For each process, only loading optimizer states of parameters it controls.
docstr_tok: ["Loading","sharded","optimizer","from","checkpoint","folder",",","with","index","file","given",".","For","each","process",",","only","loading","optimizer","states","of","parameters","it","controls","."]
code: def load_sharded_optimizer(
      self,
      optimizer: GeminiOptimizer,
      checkpoint_index_file: Path,
      prefix: str,
  ):
      """
      Loading sharded optimizer from checkpoint folder, with index file given.
      For each process, only loading optimizer states of parameters it controls.
      """
      assert isinstance(
          optimizer, GeminiOptimizer
      ), "Please boost the optimizer before loading!"
      if not os.path.isfile(checkpoint_index_file):
          logging.error(
              f"Provided path ({checkpoint_index_file}) should be a file"
          )
  
      assert isinstance(optimizer, GeminiOptimizer)
  
      # Read checkpoint index file.
      ckpt_index_file = CheckpointIndexFile.from_file(
          checkpoint_index_file
      )
  
      # Load param_groups.
      param_group_path = ckpt_index_file.get_param_group_filename()
      if param_group_path is None:
          raise RuntimeError(
              f"Invalid index file path {checkpoint_index_file} for an optimizer. \
                                 Lacking param group file under current directory."
          )
      saved_param_groups = torch.load(param_group_path)
      optimizer.load_param_groups(saved_param_groups)
  
      checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames()
  
      # Load optimizer states from shard files under checkpoint path.
      # For each file, only load the states managed by current process.
      for shard_file in checkpoint_files:
          state_dict_shard = load_shard_state_dict(
              Path(shard_file), use_safetensors=False
          )
          optimizer.load_param_states(state_dict_shard)
          del state_dict_shard
          gc.collect()
  
      optimizer.optimizer_loading_epilogue()
  
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 9 for calling the GeminiCheckpointIO obj's underlying member method code with example usage: obj.load_sharded_optimizer(optimizer, checkpoint_index_file, prefix) without return types