class_docstr: CheckpointIO is the base class for all checkpoint IO classes. It defines the interface for checkpoint IO.
  
  
  Examples:
      >>> from colossalai.checkpoint_io import GeneralCheckpointIO
      >>> checkpoint_io = CheckpointIO()
      >>>
      >>> # load model from checkpoint
      >>> model = checkpoint_io.load_model(model, 'model.pt')
      >>>
      >>> # save model to checkpoint, any distributed tensor is gathered by default
      >>> checkpoint_io.save_model(model, 'model.pt')
      >>>
      >>> # if the model contains distributed tensor, and you don't want to gather it
      >>> # each rank will save its own shard of the distributed tensor
      >>> checkpoint_io.save_model(model, 'model.pt', gather_dtensor=False)
      >>>
      >>> # save model to sharded checkpoints
      >>> checkpoint_io.save_model(model, './checkpoints/', shard=True)
      >>>
      >>> # save model to sharded  and assume we don't want to gather distributed tensors
      >>> checkpoint_io.save_model(model, './checkpoints/', shard=True, gather_dtensor=False)
      >>>
      >>> # Note:
      >>> # 1. we don't support loading from distributed tensors, conversion from distributed tensors
      >>> # checkpoints to full tensor checkpoint should be done offline via our CLI
      >>> # 2. you don't have to specify whether the model is sharded or not when loading the model
      >>> # as it will be automatically detected
      >>>
      >>> # load model from sharded checkpoints
      >>> model = checkpoint_io.load_model(model, './checkpoints/')
      >>>
      >>> # load model from unsharded checkpoints
      >>> model = checkpoint_io.load_model(model, './checkpoints/')
      >>>
      >>> # load optimizer from checkpoint
      >>> optimizer = checkpoint_io.load_optimizer(optimizer, 'optimizer.pt')
      >>>
      >>> # save optimizer to checkpoint
      >>> checkpoint_io.save_optimizer(optimizer, 'optimizer.pt')
docstr: Save optimizer to sharded checkpoint.
  
  Args:
      optimizer (Optimizer): optimizer to be saved.
      checkpoint (Path): checkpoint path. It should be a directory path.
      gather_dtensor (bool): whether to gather the distributed tensor to the first device.
      prefix (str): prefix for the optimizer checkpoint.
      size_per_shard (int): size per shard in MB.
docstr_tok: ["Save","optimizer","to","sharded","checkpoint",".","Args",":","optimizer","(","Optimizer",")",":","optimizer","to","be","saved",".","checkpoint","(","Path",")",":","checkpoint","path",".","It","should","be","a","directory","path",".","gather_dtensor","(","bool",")",":","whether","to","gather","the","distributed","tensor","to","the","first","device",".","prefix","(","str",")",":","prefix","for","the","optimizer","checkpoint",".","size_per_shard","(","int",")",":","size","per","shard","in","MB","."]
code: def save_sharded_optimizer(
      self,
      optimizer: Optimizer,
      checkpoint: Path,
      gather_dtensor: bool,
      prefix: str,
      size_per_shard: int,
  ):
      """
      Save optimizer to sharded checkpoint.
  
      Args:
          optimizer (Optimizer): optimizer to be saved.
          checkpoint (Path): checkpoint path. It should be a directory path.
          gather_dtensor (bool): whether to gather the distributed tensor to the first device.
          prefix (str): prefix for the optimizer checkpoint.
          size_per_shard (int): size per shard in MB.
      """
  
imports: from abc import ABC, abstractmethod
  from pathlib import Path
  from typing import Optional, Union
  import torch
  import torch.nn
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from colossalai.interface import ModelWrapper
  from .utils import has_index_file
label_desc: Use image node_id 11 for calling the CheckpointIO obj's underlying member method code with example usage: obj.save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard) without return types