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 model to checkpoint.
  
  Examples:
      >>> from colossalai.checkpoint_io import GeneralCheckpointIO
      >>> checkpoint_io = CheckpointIO()
      >>>
      >>> # save model to a single file
      >>> save_model(model, 'model.pt')
      >>>
      >>> # save model to a sharded checkpoint
      >>> save_model(model, './checkpoints/', shard=True)
  
  Args:
      model (nn.Module): model to be saved.
      checkpoint (str): checkpoint path. The checkpoint path can be :
          1. a file path, e.g. 'model.pt'
          2. a directory path to save the sharded checkpoint, e.g. './checkpoints/' when shard = True.
      shard (bool): whether to shard the checkpoint. Default: False. If set to True, the checkpoint will be sharded into
          multiple files. The model shards will be specified by a `model.index.json` file. When shard = True, please ensure
          that the checkpoint path is a directory path instead of a file path.
      gather_dtensor (bool): whether to gather the distributed tensor to the first device. Default: True.
      prefix (str): If specified, weights are saved in the format pytorch_model.<prefix>.bin. Default: None.
      size_per_shard (int): size per shard in MB. Default: 1024. This value is only used when shard = True.
      use_safetensors (bool): whether to use safe tensors. Default: False. If set to True, the checkpoint will be saved
docstr_tok: ["Save","model","to","checkpoint",".","Examples",":",">",">",">","from","colossalai.checkpoint_io","import","GeneralCheckpointIO",">",">",">","checkpoint_io","=","CheckpointIO","(",")",">",">",">",">",">",">","#","save","model","to","a","single","file",">",">",">","save_model","(","model",",","'model.pt","'",")",">",">",">",">",">",">","#","save","model","to","a","sharded","checkpoint",">",">",">","save_model","(","model",",","'.\/checkpoints\/","'",",","shard=True",")","Args",":","model","(","nn.Module",")",":","model","to","be","saved",".","checkpoint","(","str",")",":","checkpoint","path",".","The","checkpoint","path","can","be",":","1.","a","file","path",",","e.g",".","'model.pt'","2.","a","directory","path","to","save","the","sharded","checkpoint",",","e.g",".","'.\/checkpoints\/","'","when","shard","=","True",".","shard","(","bool",")",":","whether","to","shard","the","checkpoint",".","Default",":","False",".","If","set","to","True",",","the","checkpoint","will","be","sharded","into","multiple","files",".","The","model","shards","will","be","specified","by","a","`","model.index.json","`","file",".","When","shard","=","True",",","please","ensure","that","the","checkpoint","path","is","a","directory","path","instead","of","a","file","path",".","gather_dtensor","(","bool",")",":","whether","to","gather","the","distributed","tensor","to","the","first","device",".","Default",":","True",".","prefix","(","str",")",":","If","specified",",","weights","are","saved","in","the","format","pytorch_model.","<","prefix",">",".bin",".","Default",":","None",".","size_per_shard","(","int",")",":","size","per","shard","in","MB",".","Default",":","1024",".","This","value","is","only","used","when","shard","=","True",".","use_safetensors","(","bool",")",":","whether","to","use","safe","tensors",".","Default",":","False",".","If","set","to","True",",","the","checkpoint","will","be","saved"]
code: def save_model(
      self,
      model: Union[nn.Module, ModelWrapper],
      checkpoint: str,
      shard: bool = False,
      gather_dtensor: bool = True,
      prefix: str = None,
      size_per_shard: int = 1024,
      use_safetensors: bool = False,
  ):
      """
      Save model to checkpoint.
  
      Examples:
          >>> from colossalai.checkpoint_io import GeneralCheckpointIO
          >>> checkpoint_io = CheckpointIO()
          >>>
          >>> # save model to a single file
          >>> save_model(model, 'model.pt')
          >>>
          >>> # save model to a sharded checkpoint
          >>> save_model(model, './checkpoints/', shard=True)
  
      Args:
          model (nn.Module): model to be saved.
          checkpoint (str): checkpoint path. The checkpoint path can be :
              1. a file path, e.g. 'model.pt'
              2. a directory path to save the sharded checkpoint, e.g. './checkpoints/' when shard = True.
          shard (bool): whether to shard the checkpoint. Default: False. If set to True, the checkpoint will be sharded into
              multiple files. The model shards will be specified by a `model.index.json` file. When shard = True, please ensure
              that the checkpoint path is a directory path instead of a file path.
          gather_dtensor (bool): whether to gather the distributed tensor to the first device. Default: True.
          prefix (str): If specified, weights are saved in the format pytorch_model.<prefix>.bin. Default: None.
          size_per_shard (int): size per shard in MB. Default: 1024. This value is only used when shard = True.
          use_safetensors (bool): whether to use safe tensors. Default: False. If set to True, the checkpoint will be saved
      """
  
      if shard:
          self.save_sharded_model(
              model,
              checkpoint,
              gather_dtensor,
              prefix,
              size_per_shard,
              use_safetensors,
          )
      else:
          self.save_unsharded_model(
              model, checkpoint, gather_dtensor, use_safetensors
          )
  
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 2 for calling the CheckpointIO obj's underlying member method code with example usage: obj.save_model(model, checkpoint, shard, gather_dtensor, prefix, size_per_shard, use_safetensors) without return types