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_load_model`` 为 True）以及其它
  在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint` 
  接受到的 ``states`` ）。
  
  该函数应该在所有 rank 上执行。
  
  :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
      （如果 should_load_model 为True）。
  :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
      以及 ``'batch_idx_in_epoch'`` 这两个值。
  :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
      False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
  :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
      找到对应的模型状态，则报错。
  :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
  
      * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
        不一致时报错。
      * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
        ``batch_size`` 可能有变化，而应该符合以下等式::
  
          返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
        
        由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
        这个值来使等式成立。一个简单的计算原则如下：
  
          * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
          * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
code: def load_checkpoint(
      self,
      folder: Path,
      dataloader,
      only_state_dict: bool = True,
      should_load_model: bool = True,
      **kwargs,
  ) -> Dict:
      r"""
      断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 的状态和 **模型** （如果 ``should_load_model`` 为 True）以及其它
      在 :meth:`save_checkpoint` 函数中执行的保存操作，然后将一个 state 字典返回给 :class:`~fastNLP.core.controllers.Trainer` （ 内容为 :meth:`save_checkpoint`
      接受到的 ``states`` ）。
  
      该函数应该在所有 rank 上执行。
  
      :param folder: 读取该 folder 下的 ``FASTNLP_CHECKPOINT_FILENAME`` 文件与 ``FASTNLP_MODEL_FILENAME``
          （如果 should_load_model 为True）。
      :param dataloader: 当前给定 dataloader，需要根据保存的 dataloader 状态合理设置。若该值为 ``None`` ，则不需要返回 ``'dataloader'``
          以及 ``'batch_idx_in_epoch'`` 这两个值。
      :param only_state_dict: 是否仅读取模型的 state_dict ，当 ``should_save_model`` 为 ``False`` ，该参数无效。如果为 ``True`` ，说明保存的内容为权重；如果为
          False 说明保存的是模型，但也是通过当前 Driver 的模型去加载保存的模型的权重，而不是使用保存的模型替换当前模型。
      :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
          找到对应的模型状态，则报错。
      :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
  
          * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
            不一致时报错。
          * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
            ``batch_size`` 可能有变化，而应该符合以下等式::
  
              返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
  
            由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
            这个值来使等式成立。一个简单的计算原则如下：
  
              * drop_last 为 ``True`` 时，等同于 floor(sample_in_this_rank/batch_size) - floor(num_left_samples/batch_size)；
              * drop_last 为 ``False`` 时，等同于 ceil(sample_in_this_rank/batch_size) - ceil(num_left_samples/batch_size)。
      """
      states = oneflow.load(
          folder.joinpath(FASTNLP_CHECKPOINT_FILENAME)
      )
  
      # 1. 加载 optimizers 的状态；
      optimizers_state_dict = states.pop("optimizers_state_dict")
      self.load_optimizer_state(optimizers_state_dict)
  
      # 2. 加载模型状态；
      if should_load_model:
          self.load_model(
              filepath=folder.joinpath(FASTNLP_MODEL_FILENAME),
              only_state_dict=only_state_dict,
          )
  
      # # 3. 加载 fp16 的状态
      # if "grad_scaler_state_dict" in states:
      #     grad_scaler_state_dict = states.pop("grad_scaler_state_dict")
      #     if not isinstance(self.grad_scaler, DummyGradScaler):
      #         self.grad_scaler.load_state_dict(grad_scaler_state_dict)
      #         logger.debug("Load grad_scaler state dict...")
      # elif not isinstance(self.grad_scaler, DummyGradScaler):
      #     logger.rank_zero_warning(f"Checkpoint {folder} is not trained with fp16=True, while resume to a fp16=True training, "
      #                    f"the training process may be unstable.")
  
      # 4. 恢复 sampler 的状态；
      sampler_states = states.pop("sampler_states")
      states_ret = self.load_sampler_state(dataloader, sampler_states)
      states.update(states_ret)
  
      return states
  
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 16 for calling the OneflowDriver obj's underlying member method code with example usage: obj.load_checkpoint(folder, dataloader, only_state_dict, should_load_model) and returns: states