class_docstr: 实现了 **oneflow** 框架训练功能的基本 ``Driver``。这个类被以下子类继承：
  
      1. :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` ：实现了使用单卡和 ``cpu`` 训练的具体功能；
      2. :class:`~fastNLP.core.drivers.oneflow_driver.OneflowDDPDriver` ：实现了使用 ``DistributedDataParallel`` 启动 **oneflow** 分布式训练的功能；
  
  .. warning::
  
      您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver``，而不是
      该类本身。
  
  .. note::
  
      您可以在使用 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 时使用 ``OneflowDriver`` 提供的接口。
  
  :param model: 训练使用的模型
  :param fp16: 该参数暂时无效
  :param oneflow_kwargs:
class_docstr_tok: ["\u5b9e\u73b0\u4e86","*","*","oneflow","*","*","\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.oneflow_driver.OneflowSingleDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528\u5355\u5361\u548c","``","cpu","``","\u8bad\u7ec3\u7684\u5177\u4f53\u529f\u80fd\uff1b","2.",":","class",":","`","~fastNLP.core.drivers.oneflow_driver.OneflowDDPDriver","`","\uff1a\u5b9e\u73b0\u4e86\u4f7f\u7528","``","DistributedDataParallel","``","\u542f\u52a8","*","*","oneflow","*","*","\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","``","OneflowSingleDriver","``","\u548c","``","OneflowDDPDriver","``","\uff0c\u800c\u4e0d\u662f","\u8be5\u7c7b\u672c\u8eab\u3002","..","note",":",":","\u60a8\u53ef\u4ee5\u5728\u4f7f\u7528","``","OneflowSingleDriver","``","\u548c","``","OneflowDDPDriver","``","\u65f6\u4f7f\u7528","``","OneflowDriver","``","\u63d0\u4f9b\u7684\u63a5\u53e3\u3002",":","param","model",":","\u8bad\u7ec3\u4f7f\u7528\u7684\u6a21\u578b",":","param","fp16",":","\u8be5\u53c2\u6570\u6682\u65f6\u65e0\u6548",":","param","oneflow_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 的保存。
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 参数是 trainer 的 dataloader 属性，因为 driver 的所有 dataloader 我们是不会去改变它的，而是通过改变
      #  trainer.dataloader 来改变 dataloader 的状态，从而适配训练或者评测环境；
  
      # 1. sampler 的状态；
      num_consumed_batches = states.pop("num_consumed_batches")
      states["sampler_states"] = self.get_sampler_state(
          dataloader, num_consumed_batches
      )
  
      # 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()
      logger.debug("Save optimizer state dict.")
  
      # # 4. 保存fp16的状态
      # if not isinstance(self.grad_scaler, DummyGradScaler):
      #     grad_scaler_state_dict = self.grad_scaler.state_dict()
      #     states['grad_scaler_state_dict'] = grad_scaler_state_dict
  
      oneflow.save(
          states, Path(folder).joinpath(FASTNLP_CHECKPOINT_FILENAME)
      )
  
imports: import os
  from typing import Union, Dict, Optional, Callable, Tuple
  from functools import partial
  import numpy
  import random
  from dataclasses import dataclass
  from fastNLP.envs.imports import _NEED_IMPORT_ONEFLOW
  from pathlib import Path
  from .utils import optimizer_state_to_device, DummyGradScaler
  from fastNLP.core.drivers.driver import Driver
  from fastNLP.core.utils.utils import _get_fun_msg, nullcontext
  from fastNLP.core.utils import apply_to_collection, oneflow_move_data_to_device, auto_param_call
  from fastNLP.envs import rank_zero_call
  from fastNLP.envs import FASTNLP_GLOBAL_RANK, FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME
  from fastNLP.core.log import logger
  from fastNLP.core.samplers import ReproducibleBatchSampler, ReproducibleSampler, ReproduceBatchSampler, RandomSampler
  from fastNLP.core.dataloaders import OverfitDataLoader
label_desc: Use image node_id 11 for calling the OneflowDriver obj's underlying member method code with example usage: obj.save_checkpoint(folder, states, dataloader, only_state_dict, should_save_model) without return types