class_docstr: 实现了 **jittor** 框架训练功能的基本 ``Driver``。这个类被以下子类继承：
  
      1. :class:`~fastNLP.core.drivers.jittor_driver.JittorSingleDriver` ：实现了使用单卡和 ``cpu`` 训练的具体功能；
      2. :class:`~fastNLP.core.drivers.jittor_driver.JittorMPIDriver` ：实现了使用 ``mpi`` 启动 **jittor** 分布式训练的功能；
  
  .. warning::
  
      您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``JittorSingleDriver`` 和 ``TorchDDPDriver``，而不是
      该类本身。
  
  .. note::
  
      您可以在使用 ``JittorSingleDriver`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口。
  
  :param model: 训练时使用的 **jittor** 模型
  :param fp16: 是否开启混合精度训练
  :param jittor_kwargs:
class_docstr_tok: ["\u5b9e\u73b0\u4e86","*","*","jittor","*","*","\u6846\u67b6\u8bad\u7ec3\u529f\u80fd\u7684\u57fa\u672c","``","Driver","``","\u3002\u8fd9\u4e2a\u7c7b\u88ab\u4ee5\u4e0b\u5b50\u7c7b\u7ee7\u627f\uff1a","1.",":","class",":","`","~fastNLP.core.drivers.jittor_driver.JittorSingleDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528\u5355\u5361\u548c","``","cpu","``","\u8bad\u7ec3\u7684\u5177\u4f53\u529f\u80fd\uff1b","2.",":","class",":","`","~fastNLP.core.drivers.jittor_driver.JittorMPIDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528","``","mpi","``","\u542f\u52a8","*","*","jittor","*","*","\u5206\u5e03\u5f0f\u8bad\u7ec3\u7684\u529f\u80fd\uff1b","..","warning",":",":","\u60a8\u4e0d\u5e94\u5f53\u76f4\u63a5\u521d\u59cb\u5316\u8be5\u7c7b\uff0c\u7136\u540e\u4f20\u5165\u7ed9","``","Trainer","``","\uff0c\u6362\u53e5\u8bdd\u8bf4\uff0c\u60a8\u5e94\u5f53\u4f7f\u7528\u8be5\u7c7b\u7684\u5b50\u7c7b","``","JittorSingleDriver","``","\u548c","``","TorchDDPDriver","``","\uff0c\u800c\u4e0d\u662f","\u8be5\u7c7b\u672c\u8eab\u3002","..","note",":",":","\u60a8\u53ef\u4ee5\u5728\u4f7f\u7528","``","JittorSingleDriver","``","\u548c","``","JittorMPIDriver","``","\u65f6\u4f7f\u7528","``","JittorDriver","``","\u63d0\u4f9b\u7684\u63a5\u53e3\u3002",":","param","model",":","\u8bad\u7ec3\u65f6\u4f7f\u7528\u7684","*","*","jittor","*","*","\u6a21\u578b",":","param","fp16",":","\u662f\u5426\u5f00\u542f\u6df7\u5408\u7cbe\u5ea6\u8bad\u7ec3",":","param","jittor_kwargs",":"]
docstr: 断点重训的保存函数，该函数会负责保存 **优化器** 和 **sampler** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
  
  :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
      ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
      中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
  :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
  :param dataloader: 正在使用的 dataloader。
  :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
  :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
docstr_tok: ["\u65ad\u70b9\u91cd\u8bad\u7684\u4fdd\u5b58\u51fd\u6570\uff0c\u8be5\u51fd\u6570\u4f1a\u8d1f\u8d23\u4fdd\u5b58","*","*","\u4f18\u5316\u5668","*","*","\u548c","*","*","sampler","*","*","\u7684\u72b6\u6001\uff0c\u4ee5\u53ca","*","*","\u6a21\u578b","*","*","\uff08\u82e5","``","should_save_model","``","\u4e3a","``","True","``","\uff09",":","param","folder",":","\u4fdd\u5b58\u65ad\u70b9\u91cd\u8bad\u7684\u72b6\u6001\u7684\u6587\u4ef6\u5939\uff1b",":","meth",":","`","save_checkpoint","`","\u51fd\u6570\u5e94\u8be5\u5728\u8be5\u8def\u5f84\u4e0b\u9762\u4e0b\u9762\u65b0\u589e\u540d\u4e3a","``","FASTNLP_CHECKPOINT_FILENAME","``","\u4e0e","``","FASTNLP_MODEL_FILENAME","``","\uff08\u5982\u679c","``","should_save_model","``","\u4e3a","``","True","``","\uff09\u7684\u6587\u4ef6\u3002\u628a","model","\u76f8\u5173\u7684\u5185\u5bb9\u653e\u5165\u5230","``","FASTNLP_MODEL_FILENAME","``","\u6587\u4ef6","\u4e2d\uff0c\u5c06\u4f20\u5165\u7684","``","states","``","\u4ee5\u53ca\u81ea\u8eab\u4ea7\u751f\u7684\u5176\u5b83\u72b6\u6001\u4e00\u5e76\u4fdd\u5b58\u5728","``","FASTNLP_CHECKPOINT_FILENAME","``","\u91cc\u9762\u3002",":","param","states",":","\u7531",":","class",":","`","~fastNLP.core.controllers.Trainer","`","\u4f20\u5165\u7684\u4e00\u4e2a\u5b57\u5178\uff0c\u5176\u4e2d\u5df2\u7ecf\u5305\u542b\u4e86\u4e3a\u4e86\u5b9e\u73b0\u65ad\u70b9\u91cd\u8bad\u6240\u9700\u8981\u4fdd\u5b58\u7684\u5176\u5b83\u5bf9\u8c61\u7684\u72b6\u6001\u3002",":","param","dataloader",":","\u6b63\u5728\u4f7f\u7528\u7684","dataloader\u3002",":","param","only_state_dict",":","\u662f\u5426\u53ea\u4fdd\u5b58\u6a21\u578b\u7684\u53c2\u6570\uff0c\u5f53","``","should_save_model","``","\u4e3a","``","False","``","\uff0c\u8be5\u53c2\u6570\u65e0\u6548\u3002",":","param","should_save_model",":","\u662f\u5426\u5e94\u8be5\u4fdd\u5b58\u6a21\u578b\uff0c\u5982\u679c\u4e3a","``","False","``","\uff0cDriver","\u5c06\u4e0d\u8d1f\u8d23","model","\u7684\u4fdd\u5b58\u3002"]
code: def save_checkpoint(
      self,
      folder: Path,
      states: Dict,
      dataloader,
      only_state_dict: bool = True,
      should_save_model: bool = True,
      **kwargs,
  ):
      r"""
      断点重训的保存函数，该函数会负责保存 **优化器** 和 **sampler** 的状态，以及 **模型** （若 ``should_save_model`` 为 ``True``）
  
      :param folder: 保存断点重训的状态的文件夹；:meth:`save_checkpoint` 函数应该在该路径下面下面新增名为 ``FASTNLP_CHECKPOINT_FILENAME`` 与
          ``FASTNLP_MODEL_FILENAME`` （如果 ``should_save_model`` 为 ``True`` ）的文件。把 model 相关的内容放入到 ``FASTNLP_MODEL_FILENAME`` 文件
          中，将传入的 ``states`` 以及自身产生的其它状态一并保存在 ``FASTNLP_CHECKPOINT_FILENAME`` 里面。
      :param states: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。
      :param dataloader: 正在使用的 dataloader。
      :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
      :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
      """
      dataloader_args = self.get_dataloader_args(dataloader)
      if dataloader_args.sampler:
          sampler = dataloader_args.sampler
      else:
          raise RuntimeError(
              "This condition is not supposed to appear. Please report a bug to us."
          )
  
      num_consumed_batches = states.pop("num_consumed_batches")
      if hasattr(sampler, "state_dict") and callable(
          sampler.state_dict
      ):
          sampler_states = sampler.state_dict()
          if dataloader_args.batch_size is not None:
              sampler_states["num_consumed_samples"] = (
                  sampler.num_replicas
                  * dataloader_args.batch_size
                  * num_consumed_batches
              )
          else:
              logger.rank_zero_warning(
                  "fastNLP cannot get batch_size, we have to save based on `num_consumed_samples`, "
                  "it may cause missing some samples when reload."
              )
  
          states["sampler_states"] = sampler_states
      else:
          raise RuntimeError(
              "The sampler has no `state_dict()` method, fastNLP cannot save the training "
              "state."
          )
  
      # 2. 保存模型的状态；
      if should_save_model:
          if not os.path.exists(folder):
              os.mkdir(folder)
          model_path = folder.joinpath(FASTNLP_MODEL_FILENAME)
          self.save_model(model_path, only_state_dict=only_state_dict)
  
      # 3. 保存 optimizers 的状态；
      states["optimizers_state_dict"] = self.get_optimizer_state()
  
      # 4. 保存fp16的状态
  
      logger.debug("Save optimizer state dict")
      jt.save(
          states, Path(folder).joinpath(FASTNLP_CHECKPOINT_FILENAME)
      )
  
imports: import os
  from pathlib import Path
  from typing import Union, Optional, Dict
  from dataclasses import dataclass
  from fastNLP.envs.imports import _NEED_IMPORT_JITTOR
  from fastNLP.core.drivers.driver import Driver
  from fastNLP.core.dataloaders import JittorDataLoader
  from fastNLP.core.dataloaders import OverfitDataLoader
  from fastNLP.core.samplers import ReproducibleSampler, RandomSampler
  from fastNLP.core.log import logger
  from fastNLP.core.utils import apply_to_collection, nullcontext
  from fastNLP.envs import FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME
label_desc: Use image node_id 9 for calling the JittorDriver obj's underlying member method code with example usage: obj.save_checkpoint(folder, states, dataloader, only_state_dict, should_save_model) without return types