调整 fastNLP/core/drivers 的文档

fastNLP/core/callbacks/callback.py

@@ -264,7 +264,7 @@ class Callback:
         r"""
         ``callback`` 的名称，我们会使用该名称从 ``checkpoint`` 中读取的相应的 ``state`` 并传递给 :meth:`on_load_checkpoint` 函数。
 
-        :return: 返回用于区分该 ``callback`` 实例的名称；
+        :return: 用于区分该 ``callback`` 实例的名称；
         """
         return self.__class__.__name__
 

fastNLP/core/callbacks/topk_saver.py

@@ -62,7 +62,7 @@ class Saver:
 
         :param trainer: Trainer 对象
         :param folder_name: 保存的 folder 名称，将被创建。
-        :return: 返回实际发生保存的 folder 绝对路径。如果为 None 则没有创建。
+        :return: 实际发生保存的 folder 绝对路径。如果为 None 则没有创建。
         """
         folder = self.timestamp_path.joinpath(folder_name)
         folder.mkdir(parents=True, exist_ok=True)

fastNLP/core/callbacks/utils.py

@@ -14,7 +14,7 @@ def _get_monitor_value(monitor: Union[callable, str], real_monitor: Optional[str
     :param monitor:
     :param real_monitor:
     :param res:
-    :return: 返回两个值（str, value)，其中str就是最终要到的key，value就是这个key对应的value。如果value为None说明当前results中没有
+    :return: 两个值（str, value)，其中str就是最终要到的key，value就是这个key对应的value。如果value为None说明当前results中没有
         找到对应的 monitor
     """
     if len(res) == 0 or monitor is None:

fastNLP/core/collators/collator.py

@@ -211,7 +211,7 @@ class Collator:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 Collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回 Collator 自身；
+        :return: Collator 自身；
         """
         self._renew()
 
@@ -298,7 +298,7 @@ class Collator:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回 Collator 自身；
+        :return: Collator 自身；
         """
         self._renew()
         input_field_names = [(field, field) if isinstance(field, tuple) else ((field,), field)

fastNLP/core/controllers/evaluator.py

@@ -257,7 +257,7 @@ class Evaluator:
                ``metric_indicator_name#metric_name#dataloader_name``，其中 metric_indicator_name 可能不存在；
 
         :param num_eval_batch_per_dl: 每个 dataloader 测试前多少个 batch 的数据，-1 为测试所有数据。
-        :return: 返回评测得到的结果，是一个没有嵌套的字典；
+        :return: 评测得到的结果，是一个没有嵌套的字典；
         """
         assert isinstance(num_eval_batch_per_dl, int), "num_eval_batch_per_dl must be of int type."
         assert num_eval_batch_per_dl > 0 or num_eval_batch_per_dl == -1, "num_eval_batch_per_dl must be -1 or larger than 0."

fastNLP/core/controllers/trainer.py

@@ -73,6 +73,10 @@ class Trainer(TrainerEventTrigger):
             这意味着当您传入一个 ``Driver`` 实例时，您传入给 ``Trainer`` 的 ``model`` 参数将会被忽略；也就是说模型在训练时使用的真正的模型是
             您传入的 ``Driver`` 实例中的模型；
 
+        .. note::
+
+            如果您选择使用 :mod:`deepspeed` 或 :mod:`fairscale` 进行训练，请不要将 ``driver`` 的值设为 ``'auto'`` 。
+
     :param train_dataloader: 训练数据集，注意其必须是单独的一个数据集，不能是 :class:`List` 或者 :class:`Dict`；
 
         .. warning::
@@ -885,7 +889,7 @@ class Trainer(TrainerEventTrigger):
         :param marker: 用来标记该 callback 函数属于哪几个具体的 trainer 实例；两个特殊情况：1.当 ``marker`` 为 None（默认情况）时，
          表示该 callback 函数只属于代码下方最近的一个 trainer 实例；2.当 ``marker`` 为 'all' 时，该 callback 函数会被所有的 trainer
          实例使用；
-        :return: 返回原函数；
+        :return: 原函数；
         """
 
         def wrapper(fn: Callable) -> Callable:
@@ -1001,7 +1005,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def driver(self):
         """
-        :return: 返回 ``trainer`` 中的 ``driver`` 实例；
+        :return:  ``trainer`` 中的 ``driver`` 实例；
         """
         return self._driver
 
@@ -1012,7 +1016,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def train_batch_loop(self):
         """
-        :return: 返回 ``trainer`` 中的 ``train_batch_loop`` 实例；
+        :return:  ``trainer`` 中的 ``train_batch_loop`` 实例；
         """
         return self._train_batch_loop
 
@@ -1275,7 +1279,7 @@ class Trainer(TrainerEventTrigger):
             ``trainer.backward / zero_grad / step`` 函数的作用类似；
 
         :param batch: 一个 batch 的数据；
-        :return: 返回模型的前向传播函数所返回的结果；
+        :return: 模型的前向传播函数所返回的结果；
         """
         with self.driver.auto_cast():
             outputs = self.driver.model_call(batch, self._train_step, self._train_step_signature_fn)
@@ -1328,7 +1332,7 @@ class Trainer(TrainerEventTrigger):
         用来从用户模型的输出对象中抽取 ``loss`` 对象；
         目前支持 `outputs` 对象为 ``dict`` 或者 ``dataclass``；
 
-        :return: 返回被抽取出来的 ``loss`` 对象，例如如果是 ``pytorch``，那么返回的就是一个 tensor；
+        :return: 被抽取出来的 ``loss`` 对象，例如如果是 ``pytorch``，那么返回的就是一个 tensor；
         """
         if isinstance(outputs, Dict):
             try:
@@ -1375,7 +1379,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def n_epochs(self) -> int:
         r"""
-        :return: 返回当前训练的总体的 epoch 的数量；
+        :return: 当前训练的总体的 epoch 的数量；
         """
         return self.trainer_state.n_epochs
 
@@ -1386,7 +1390,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def cur_epoch_idx(self) -> int:
         r"""
-        :return: 返回当前正在第几个 epoch；
+        :return: 当前正在第几个 epoch；
         """
         return self.trainer_state.cur_epoch_idx
 
@@ -1397,7 +1401,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def global_forward_batches(self) -> int:
         """
-        :return: 返回从训练开始到当前总共训练了多少 batch 的数据；
+        :return: 从训练开始到当前总共训练了多少 batch 的数据；
         """
         return self.trainer_state.global_forward_batches
 
@@ -1408,7 +1412,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def batch_idx_in_epoch(self) -> int:
         r"""
-        :return: 返回在从当前的这个 epoch 开始，到现在共训练了多少 batch 的数据；
+        :return: 在从当前的这个 epoch 开始，到现在共训练了多少 batch 的数据；
         """
         return self.trainer_state.batch_idx_in_epoch
 
@@ -1419,7 +1423,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def num_batches_per_epoch(self) -> int:
         r"""
-        :return: 返回每一个 epoch 实际会训练多少个 batch 的数据；
+        :return: 每一个 epoch 实际会训练多少个 batch 的数据；
         """
         return self.trainer_state.num_batches_per_epoch
 
@@ -1430,7 +1434,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def n_batches(self) -> int:
         r"""
-        :return: 返回整体的训练中实际会训练多少个 batch 的数据；
+        :return: 整体的训练中实际会训练多少个 batch 的数据；
         """
         return self.trainer_state.n_batches
 
@@ -1443,7 +1447,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def model_device(self):
         r"""
-        :return: 返回当前模型所在的设备；注意该值在当且仅当在少数情况下为 ``None``，例如当使用 ``pytorch`` 时，仅当用户自己初始化 ``init_progress_group`` 时
+        :return: 当前模型所在的设备；注意该值在当且仅当在少数情况下为 ``None``，例如当使用 ``pytorch`` 时，仅当用户自己初始化 ``init_progress_group`` 时
             ``model_device`` 才为 None；
         """
         return self.driver.model_device
@@ -1451,7 +1455,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def data_device(self):
         r"""
-        :return: 返回数据会被迁移到的目的设备；
+        :return: 数据会被迁移到的目的设备；
         """
         return self.driver.data_device
 
@@ -1460,7 +1464,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def train_dataloader(self):
         """
-        :return: 返回用户传入的 ``train_dataloader``，注意该 ``dataloader`` 与用户传入给 ``Trainer`` 的 ``dataloader`` 对象是同一个对象，而我们在
+        :return: 用户传入的 ``train_dataloader``，注意该 ``dataloader`` 与用户传入给 ``Trainer`` 的 ``dataloader`` 对象是同一个对象，而我们在
         实际训练过程中使用的 ``dataloader`` 的状态可能有所更改；
         """
         return self._train_dataloader
@@ -1472,7 +1476,7 @@ class Trainer(TrainerEventTrigger):
     @property
     def evaluate_dataloaders(self):
         """
-        :return: 返回用户传入的 ``evaluate_dataloaders``；
+        :return: 用户传入的 ``evaluate_dataloaders``；
         """
         return self._evaluate_dataloaders
 

fastNLP/core/controllers/utils/state.py

@@ -65,7 +65,7 @@ class TrainerState:
 
     def state_dict(self) -> Dict:
         r"""
-        :return: 返回用于断点重训来保存的状态字典；
+        :return: 用于断点重训来保存的状态字典；
         """
         return {"cur_epoch_idx": self.cur_epoch_idx, "global_forward_batches": self.global_forward_batches,
                 "batch_idx_in_epoch": self.batch_idx_in_epoch}

fastNLP/core/dataloaders/jittor_dataloader/fdl.py

@@ -146,7 +146,7 @@ class JittorDataLoader:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -177,7 +177,7 @@ class JittorDataLoader:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/oneflow_dataloader/fdl.py

@@ -167,7 +167,7 @@ class OneflowDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -198,7 +198,7 @@ class OneflowDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/paddle_dataloader/fdl.py

@@ -199,7 +199,7 @@ class PaddleDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -230,7 +230,7 @@ class PaddleDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataloaders/torch_dataloader/fdl.py

@@ -167,7 +167,7 @@ class TorchDataLoader(DataLoader):
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):
@@ -198,7 +198,7 @@ class TorchDataLoader(DataLoader):
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回使用的 collator
+        :return: 使用的 collator
         """
         collator = self._get_collator()
         if isinstance(collator, Collator):

fastNLP/core/dataset/dataset.py

@@ -637,7 +637,7 @@ class DataSet:
 
         :param progress_desc: 如果不为 ``None``，则会显示当前正在处理的进度条的名称；
         :param progress_bar: 显示进度条的方式，支持 ``["rich", "tqdm", None]``。
-        :return: 返回一个字典
+        :return: 一个字典
         """
         assert len(self) != 0, "Null DataSet cannot use apply_field()."
         if not self.has_field(field_name=field_name):
@@ -762,7 +762,7 @@ class DataSet:
 
         :param progress_desc: 当 progress_bar 不为 ``None`` 时，可以显示当前正在处理的进度条名称
         :param progress_bar: 显示进度条的方式，支持 ``["rich", "tqdm", None]``。
-        :return: 返回一个字典
+        :return: 一个字典
         """
         assert callable(func), "The func is not callable."
         assert len(self) != 0, "Null DataSet cannot use apply()."
@@ -1013,7 +1013,7 @@ class DataSet:
             若 ``pad_val`` 为 ``None`` ，该值无意义 。
         :param pad_fn: 指定当前 field 的 pad 函数，传入该函数则 ``pad_val``, ``dtype``, ``backend`` 等参数失效。``pad_fn`` 的输入为当前 field 的
             batch 形式。 Collator 将自动 unbatch 数据，然后将各个 field 组成各自的 batch 。
-        :return: 返回自身的 collator；
+        :return: 自身的 collator；
         """
         if isinstance(self.collator, Collator):
             self.collator.set_pad(field_name=field_name, pad_val=pad_val, dtype=dtype, pad_fn=pad_fn, backend=backend)
@@ -1031,7 +1031,7 @@ class DataSet:
         :param field_names: field_name: 需要调整的 field 的名称。如果 :meth:`Dataset.__getitem__` 方法返回的是字典类型，则可以直接使用对应的
             field 的 key 来表示，如果是嵌套字典，可以使用元组表示多层次的 key，例如 ``{'a': {'b': 1}}`` 中可以使用 ``('a', 'b')``;
             如果 :meth:`Dataset.__getitem__` 返回的是 Sequence 类型，则可以使用 ``'_0'``, ``'_1'`` 表示序列中第 **0** 或 **1** 个元素。
-        :return: 返回自身的 collator；
+        :return: 自身的 collator；
         """
         if isinstance(self.collator, Collator):
             self.collator.set_ignore(*field_names)

fastNLP/core/drivers/choose_driver.py

@@ -3,15 +3,17 @@ from typing import Union, Optional, List
 from .driver import Driver
 from ..utils import is_torch_module, is_paddle_module, is_jittor_module, is_oneflow_module
 
+__all__ = []
 
 def choose_driver(model, driver: Union[str, Driver], device: Optional[Union[int, List[int], str]], **kwargs) -> Driver:
     r"""
-    根据输入的参数 'gpus' 的格式来决定具体的工作模式;
+    根据输入的参数 ``driver`` 和 ``device`` 的格式来决定具体的工作模式。
 
-    :param model: 运行过程中使用的具体的最原始的模型；
-    :param driver: 应当为字符串或者 `Driver` 实例，表示运行中具体使用的训练/评测模式；
-    :param device: 具体的形式请参见 `fastNLP.core.drivers.torch_driver.utils.initialize_torch_dirver` 的注释；
-    :param kwargs: 其余的传给 `Driver` 的参数；
+    :param model: 运行过程中使用的具体的最原始的模型。
+    :param driver: 训练模型所使用的具体的驱动模式，应当为以下选择中的一个：``["auto", "torch", "paddle", "jittor", "fairscale", "deepspeed", "oneflow"]``，分别对应
+        各种框架。值为 ``'auto'`` 时，将会根据模型的类型进行选择。
+    :param device: 训练使用的设备。详细的格式可以查阅 :class:`~fastNLP.core.controllers.Trainer` 中的说明。
+    :param kwargs: 其余的传给 `Driver` 的参数。
     """
 
     # 如果用户直接传进来一个 driver 实例，我们就直接返回回去，目前用户需要自己保证传进来的 driver 的正确性；

fastNLP/core/drivers/driver.py

@@ -18,9 +18,10 @@ from fastNLP.core.utils import nullcontext
 class Driver(ABC):
     r"""
     用来初始化 `Driver` 的基类，所有定制的 `driver` 都需要继承此类；
-    fastNLP 提供的 driver 实例都会同时被 Trainer 和 Evaluator 调用；
-    :param model: 训练或者评测的模型，需要注意该模型可能为用户已经使用类似 `torch.nn.DataParallel` 或者
-        `torch.nn.parallel.DistributedDataParallel` 包裹过的模型；
+    **fastNLP** 提供的 driver 实例都会同时被 :class:`~fastNLP.core.controllers.Trainer` 和 :class:`~fastNLP.core.controllers.Evaluator` 调用。
+
+    :param model: 训练或者评测的模型，需要注意该模型可能为用户已经使用类似 :class:`torch.nn.DataParallel` 或者
+        :class:`torch.nn.parallel.DistributedDataParallel` 包裹过的模型。
     """
 
     def __init__(self, model):
@@ -33,29 +34,32 @@ class Driver(ABC):
     @abstractmethod
     def setup(self):
         r"""
-        该函数用来初始化训练环境，例如将模型迁移到对应的设备上等；
+        该函数用来初始化训练环境，例如将模型迁移到对应的设备上等。
         多卡的 ``driver`` 的该函数要更为复杂一些，例如其可能需要开启多进程之间的通信环境，以及设置一些环境变量和其余所需要的变量值；
         """
 
     def set_dist_repro_dataloader(self, dataloader, dist=None, reproducible: bool = False):
         r"""
-        根据输入的 ``dataloader`` 得到一个 支持分布式 （``distributed``） 与 可复现的 (``reproducible``) 的 dataloader。
-
-        :param dataloader: 根据 ``dataloader`` 设置其对应的分布式版本以及可复现版本；
-        :param dist: 应当为一个字符串，其值应当为以下之一：``[None, "dist", "unrepeatdist"]``；为 ``None`` 时，表示不需要考虑当前 dataloader
-            切换为分布式状态；为 ``dist`` 时，表示该 dataloader 应该保证每个 gpu 上返回的 batch 的数量是一样多的，允许出现少量 sample ，在
-            不同 gpu 上出现重复；为 ``unrepeatdist`` 时，表示该 dataloader 应该保证所有 gpu 上迭代出来的数据合并起来应该刚好等于原始的
-            数据，允许不同 gpu 上 batch 的数量不一致。
-            其中 trainer 中 kwargs 的参数 ``use_dist_sampler`` 为 ``True`` 时，该值为 ``dist``；
-            否则为 ``None``，evaluator 中的 kwargs 的参数 ``use_dist_sampler`` 为 ``True`` 时，该值为 ``unrepeatdist``，否则为 ``None``；
-            注意当 dist 为 ReproducibleSampler, ReproducibleBatchSampler 时，是断点重训加载时 driver.load_checkpoint 函数在调用；
-            当 dist 为 str 或者 None 时，是 trainer 在初始化时调用该函数；
+        根据输入的 ``dataloader`` 得到一个 支持分布式 （**distributed**） 与 可复现的 (**reproducible**) 的 dataloader。
+
+        :param dataloader: 根据 ``dataloader`` 设置其对应的分布式版本以及可复现版本。
+        :param dist: 应当为一个字符串，其值应当为以下之一：``[None, "dist", "unrepeatdist"]``，并且根据在 :class:`~fastNLP.core.controllers.Trainer`
+            和 :class:`~fastNLP.core.controllers.Evaluator` 中 *kwargs* 的参数 ``use_dist_sampler`` 和调用时机不同，对应不同的值：
+        
+            * 当 ``use_dist_sampler`` 为 ``False`` ，且在 :class:`~fastNLP.core.controllers.Trainer` 或 :class:`~fastNLP.core.controllers.Evaluator`
+              **初始化** 中被调用时，参数值为 ``None`` ，表示不需要考虑当前 ``dataloader`` 切换为分布式状态；
+            * 当 ``use_dist_sampler`` 为 ``True`` ，且在 :class:`~fastNLP.core.controllers.Trainer` **初始化** 中被调用时，参数值为 ``"dist"`` ，表示该
+              ``dataloader`` 应该保证每个 gpu 上返回的 batch 的数量是一样多的，允许出现少量 sample 在不同 gpu 上出现重复；
+            * 当 ``use_dist_sampler`` 为 ``True`` ，且在 :class:`~fastNLP.core.controllers.Evaluator` **初始化** 中被调用时，参数值为 ``"unrepeatdist"`` ，
+              表示该 ``dataloader`` 应该保证所有 gpu 上迭代出来的数据合并起来应该刚好等于原始的数据，允许不同 gpu 上 batch 的数量不一致；
+            * 当 **断点重训加载** 中调用 :meth:`load_checkpoint` 时，该函数也会被调用，且 ``dist`` 值为 :class:`~fastNLP.core.samplers.ReproducibleSampler`
+              或 :class:`~fastNLP.core.samplers.ReproducibleBatchSampler` ，此时表示需要用 ``dist`` 代表的 sampler 或 batch_sampler 重新实例化一个新的 dataloader；
 
         :param reproducible: 如果为 ``False``，不要做任何考虑；如果为 ``True``，需要保证返回的 dataloader 可以保存当前的迭代状态，使得
-            该状态可以加载到一个全新的 dataloader 中然后恢复其状态；
-        :return: 应当返回一个被替换 sampler 后的新的 dataloader 对象 (注意此处一定需要返回一个新的 dataloader 对象) ；此外，
-            如果传入的 dataloader 中是 ReproducibleSampler 或者 ReproducibleBatchSampler 需要重新初始化一个放入返回的
-            dataloader 中。如果 dist 为空，且 reproducible 为 False，可直接返回原对象。
+            该状态可以加载到一个全新的 dataloader 中然后恢复其状态。
+        :return: 应当返回一个被替换 sampler 后的 **新的** dataloader 对象 (注意此处一定需要返回一个新的 dataloader 对象) ；此外，
+            如果传入的 ``dataloader`` 中是 :class:`~fastNLP.core.samplers.ReproducibleSampler` 或者 :class:`~fastNLP.core.samplers.ReproducibleBatchSampler`
+            需要 **重新初始化** 一个放入返回的 dataloader 中。如果 ``dist`` 为空，且 ``reproducible`` 为 ``False``，可直接返回原对象。
         """
         if dist is None and reproducible is False:
             return dataloader
@@ -64,64 +68,69 @@ class Driver(ABC):
 
     def set_deterministic_dataloader(self, dataloader):
         r"""
-        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的；例如对于 ``pytorch`` 的 ``dataloader``，其
-        需要将 ``worker_init_fn`` 替换；
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的；例如对于 **pytorch** 的 ``dataloader``，其
+        需要将 ``worker_init_fn`` 替换。
         """
 
     def set_sampler_epoch(self, dataloader, cur_epoch_idx):
         r"""
-        对于分布式的 ``sampler``，例如 ``pytorch`` 的 ``DistributedSampler``，其需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的；
+        对于分布式的 ``sampler``，例如 **pytorch** 的 :class:`DistributedSampler`，其需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的；
         ``dataloader`` 中可能真正发挥作用的是 ``batch_sampler`` 也可能是 ``sampler``。
 
-        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``；
-        :param cur_epoch_idx: 当前是第几个 ``epoch``；
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
         """
 
     @abstractmethod
     def model_call(self, batch, fn: Callable, signature_fn: Optional[Callable]) -> Dict:
         r"""
         通过调用 ``fn`` 来实现训练时的前向传播过程；
-        注意 ``Trainer`` 和 ``Evaluator`` 会调用该函数来实现网络的前向传播过程，其中传入该函数的参数 ``fn`` 是函数 ``get_model_call_fn`` 所返回的
-        函数；
+        注意 :class:`~fastNLP.core.controllers.Trainer` 和 :class:`~fastNLP.core.controllers.Evaluator` 会调用该函数来
+        实现网络的前向传播过程，其中传入该函数的参数 ``fn`` 是函数 :meth:`get_model_call_fn` 所返回的函数。
 
-        :param batch: 当前的一个 batch 的数据；可以为字典或者其它类型；
+        :param batch: 当前的一个 batch 的数据；可以为字典或者其它类型。
         :param fn: 调用该函数进行一次计算。
-        :param signature_fn: 由 ``Trainer`` 传入的用于网络前向传播一次的签名函数，因为当 batch 是一个 ``Dict`` 的时候，我们会自动调用 ``auto_param_call`` 函
-            数，而一些被包裹的模型需要暴露其真正的函数签名，例如 ``DistributedDataParallel`` 的调用函数是 ``forward``，但是需要其函数签名为 ``model.module.forward``；
-        :return: 返回由 ``fn`` 返回的结果（应当为一个 ``dict`` 或者 ``dataclass``，但是不需要我们去检查）；
+        :param signature_fn: 由 :class:`~fastNLP.core.controllers.Trainer` 传入的用于网络前向传播一次的签名函数，因为当
+            batch 是一个 :class:`Dict` 的时候，我们会自动调用 :func:`fastNLP.core.utils.auto_param_call` 函数，而一些被
+            包裹的模型需要暴露其真正的函数签名，例如 :class:`DistributedDataParallel` 的调用函数是 ``forward``，但是需要其
+            函数签名为 ``model.module.forward``。
+        :return: 由 ``fn`` 返回的结果（应当为一个 :class:`dict` 或者 :class:`dataclass` ，但是不需要我们去检查）。
         """
         raise NotImplementedError("Each specific driver should implemented its own `model_call` function.")
 
     @abstractmethod
     def get_model_call_fn(self, fn: str) -> Tuple:
         r"""
-        该函数会接受 ``Trainer`` 的 ``train_fn`` 或者 ``Evaluator`` 的 ``evaluate_fn``，返回一个实际用于调用 ``driver.model_call`` 时传入的函数参数；
-        该函数会在 ``Trainer`` 和 ``Evaluator`` 在 ``driver.setup`` 函数之后调用；
+        该函数会接受 :class:`~fastNLP.core.controllers.Trainer` 的 ``train_fn`` 或者 :class:`~fastNLP.core.controllers.Evaluator`
+        的 ``evaluate_fn``，返回一个实际用于调用 :meth:`model_call` 时传入的函数参数；该函数会由 :class:`~fastNLP.core.controllers.Trainer`
+        和 :class:`~fastNLP.core.controllers.Evaluator` 在 :func:`driver.setup` 函数之后调用。
 
-        之所以设置该函数的目的在于希望将具体的 model_call function 从 driver 中抽离出来，然后将其附着在 Trainer 或者 Evaluator 身上；
+        之所以设置该函数的目的在于希望将具体的 model_call function 从 driver 中抽离出来，然后将其附着在 ``Trainer`` 或者 ``Evaluator`` 身上；
         这样是因为在新版的设计中，使用 model 的哪种方法来进行 ``train step`` 或者 ``evaluate step`` 是通过额外的参数 ``train_fn`` 和
-        ``evaluate_fn`` 来确定的，而二者又分别是通过 Trainer 和 Evaluator 来控制的；因此不能将确定具体的 ``train step fn`` 和
-        ``evaluate step fn`` 的逻辑放在每一个 driver 的初始化的时候（因此在 Trainer 初始化第一个 driver 时，Evaluator 还没有初始化，但是
-        ``evaluate step fn`` 的确定却需要 Evaluator 的初始化），因此我们将这一逻辑抽象到这一函数当中；
+        ``evaluate_fn`` 来确定的，而二者又分别是通过 ``Trainer`` 和 ``Evaluator`` 来控制的；因此不能将确定具体的 ``train step fn`` 和
+        ``evaluate step fn`` 的逻辑放在每一个 driver 的初始化的时候（因此在 ``Trainer`` 初始化第一个 driver 时，``Evaluator`` 还没有初始化，但是
+        ``evaluate step fn`` 的确定却需要 Evaluator 的初始化），因此我们将这一逻辑抽象到这一函数当中.
 
         这一函数应当通过参数 ``fn`` 来判断应当返回的实际的调用的函数，具体逻辑如下所示：
-        1. 如果 fn == "train_step" or "evaluate_step"，那么对传入的模型进行检测，如果模型没有定义方法 ``fn``，则默认调用模型的 ``forward``
-        函数，然后给出 warning；
-        2. 如果 fn 是其他字符串，那么如果模型没有定义方法 ``fn`` 则直接报错；
 
-        注意不同的 driver 需要做额外的检测处理，例如在 DDPDriver 中，当传入的模型本身就是 DistributedDataParallel 中，我们只能调用模型的
-        forward 函数，因此需要额外的 warning；这一点特别需要注意的问题在于 driver 自己在 setup 时也会对模型进行改变（DDPDriver），因此
-        可能需要额外标记最初传入 driver 的模型是哪种形式的；
+            1. 如果 ``fn`` == "train_step" or "evaluate_step"，那么对传入的模型进行检测，如果模型没有定义方法 ``fn``，则默认调用模型的 :meth:`forward`
+               函数，然后给出 warning；
+            2. 如果 ``fn`` 是其他字符串，那么如果模型没有定义方法 ``fn`` 则直接报错；
+
+        注意不同的 driver 需要做额外的检测处理，例如在 :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` 中，当传入的模型本身就是
+        :class:`DistributedDataParallel` 时，我们只能调用模型的 :meth:`forward` 函数，因此需要额外的 warning；这一点特别需要注意的问题在于
+        driver 自己在 setup 时也会对模型进行改变（ :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` ），因此可能需要额外标记最初
+        传入 driver 的模型是哪种形式的.
 
-        :param fn: 应当为一个字符串，该函数通过该字符串判断要返回模型的哪种方法；
-        :return: 返回一个元组，包含两个函数，用于在调用 driver.model_call 时传入；
+        :param fn: 一个字符串，该函数通过该字符串判断要返回模型的哪种方法
+        :return: 一个元组，包含两个函数，用于在调用 :meth:`model_call` 时传入
         """
         raise NotImplementedError("Each specific driver should implemented its own `get_model_call_fn` function.")
 
     @property
     def model(self):
         r"""
-        :return: 返回 driver 中在实际训练或者评测时所使用的模型；
+        :return:  driver 中在实际训练或者评测时所使用的模型。
         """
         return self._model
 
@@ -132,8 +141,8 @@ class Driver(ABC):
     @property
     def optimizers(self) -> List:
         r"""
-        如下所示，driver 返回的 optimizers 一定是一个 List，如果用户直接向 Trainer 传入一个单独的 optimizer，我们会使用一个 List 将其
-        包裹；
+        如下所示，driver 返回的 :attr:`optimizers` 一定是一个 :class:`List`，如果用户直接向 :class:`~fastNLP.core.controllers.Trainer` 传入一个单独的 optimizer，
+        我们会使用一个 List 将其包裹；
 
         :return: List[optimizer0, optimizer1, optimizer2, ...]
         """
@@ -150,7 +159,7 @@ class Driver(ABC):
     @property
     def model_device(self):
         r"""
-        :return: 返回 driver 中模型实际所在的设备；
+        :return: driver 中模型实际所在的设备。
         """
         return self._model_device
 
@@ -161,55 +170,54 @@ class Driver(ABC):
     @property
     def data_device(self):
         """
-        :return: 返回 driver 中数据默认会被迁移到的设备；
+        :return: driver 中数据默认会被迁移到的设备。
         """
         return self.model_device
 
     @staticmethod
     def _check_optimizer_legality(optimizers):
         r"""
-        对于用户传入 trainer 的每一个 optimizer，检测其是否合理，因为不同的深度学习框架所使用的的 optimizer 是不相同的；
+        对于用户传入 trainer 的每一个 optimizer，检测其是否合理，因为不同的深度学习框架所使用的的 optimizer 是不相同的。
 
-        :param optimizers: 需要检测的 `optimizers`；
+        :param optimizers: 需要检测的 `optimizers`。
         """
         raise NotImplementedError(
             "Each specific driver should implemented its own `_check_optimizer_legality` function.")
 
     def check_dataloader_legality(self, dataloader):
         """
-        检测 DataLoader 是否合法，如果不合法，会 raise TypeError 。
+        检测 ``dataloader`` 是否合法，如果不合法，会 ``raise TypeError`` 。
 
         :param dataloder:
-        :return:
         """
 
     def set_optimizers(self, optimizers=None):
         r"""
-        trainer 会调用该函数将用户传入的 optimizers 挂载到 driver 实例上；
+        trainer 会调用该函数将用户传入的 ``optimizers`` 挂载到 driver 实例上。
         """
         self.optimizers = optimizers
 
     @abstractmethod
     def backward(self, loss):
         r"""
-        实现深度学习中的反向传播过程；
+        实现深度学习中的反向传播过程。
 
-        :param loss: 用来实现反向传播的损失函数值；
+        :param loss: 用来实现反向传播的损失函数值
         """
         raise NotImplementedError("Each specific driver should implemented its own `backward` function.")
 
     @abstractmethod
     def step(self):
         r"""
-        实现深度学习中的参数的优化更新过程，应当直接通过优化器 optimizers 来更新参数；
+        实现深度学习中的参数的优化更新过程，应当直接通过优化器 :attr:`optimizers` 来更新参数。
         """
         raise NotImplementedError("Each specific driver should implemented its own `step` function.")
 
     @abstractmethod
     def zero_grad(self):
         r"""
-        实现深度学习中的梯度的置零操作，应当直接通过优化器 optimizers 来将梯度置零；
-        注意梯度累积不需要在这里实现，trainer 已经在内部实现了梯度累积；
+        实现深度学习中的梯度的置零操作，应当直接通过优化器 :attr:`optimizers` 来将梯度置零；
+        注意梯度累积不需要在这里实现，trainer 已经在内部实现了梯度累积。
 
         """
         raise NotImplementedError("Each specific driver should implemented its own `zero_grad` function.")
@@ -217,26 +225,26 @@ class Driver(ABC):
     def get_model_no_sync_context(self):
         r"""
         返回一个用于关闭多进程之间 model 中的自动互相同步操作的 context 上下文对象；只有多卡的 driver 需要单独实现该函数，
-        单卡的 driver 不需要；
+        单卡的 driver 不需要。
 
-        :return: 返回一个类似于 DistributedDataParallel(model).no_sync 的 context 上下文对象；
+        :return: 一个类似于 ``DistributedDataParallel(model).no_sync`` 的 context 上下文对象
         """
         return nullcontext
 
     def get_evaluate_context(self):
         r"""
-        返回一个不计算梯度的环境用来对模型进行评测；
+        返回一个不计算梯度的环境用来对模型进行评测。
 
-        :return: 一个类似 `torch.no_grad` 的 context 上下文对象；
+        :return: 一个类似 ``torch.no_grad`` 的 context 上下文对象
         """
         return nullcontext
 
     @property
     def auto_cast(self):
         r"""
-        fp16 的上下文环境；
+        fp16 的上下文环境。
 
-        :return: 返回一个用于 fp16 计算的上下文环境；
+        :return: 一个用于 fp16 计算的上下文环境
         """
         return self._auto_cast
 
@@ -247,20 +255,19 @@ class Driver(ABC):
     @abstractmethod
     def save_model(self, filepath: Union[str, Path, BytesIO], only_state_dict: bool = True, **kwargs):
         r"""
-        保存模型的函数；注意函数 `save` 是用来进行断点重训的函数；
+        保存模型的函数；注意函数 :meth:`save_checkpoint` 是用来进行断点重训的函数。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）或一个 BytesIO 对象；
-        :param only_state_dict: 是否只保存模型的 `state_dict`；
-        :param model_save_fn: 用户传入的用来代替该函数本身保存逻辑的函数；如果该参数不为 None，那么我们会调用 model_save_fn(path)；
+        :param filepath: 保存文件的文件位置（需要包括文件名）或一个 BytesIO 对象
+        :param only_state_dict: 是否只保存模型的 `state_dict`
         """
         raise NotImplementedError("Each specific driver should implemented its own `save_model` function.")
 
     @abstractmethod
     def load_model(self, filepath: Union[str, Path, BytesIO], only_state_dict: bool = False, **kwargs):
         r"""
-        加载模型的函数；将 filepath 中的模型加载并赋值给当前 model 。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 model 。
 
-        :param filepath: 需要被加载的对象的文件位置（需要包括文件名）或一个 ``BytesIO`` 对象；
+        :param filepath: 需要被加载的对象的文件位置（需要包括文件名）或一个 ``BytesIO`` 对象。
         :param load_state_dict: 保存的文件是否只是模型的权重，还是完整的模型。即便是保存的完整的模型，此处也只能使用尝试加载filepath
             模型中的权重到自身模型，而不会直接替代当前 Driver 中的模型。
         """
@@ -269,18 +276,18 @@ class Driver(ABC):
     @abstractmethod
     def save_checkpoint(self, folder, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True,
              **kwargs):
+
         r"""
-        断点重训的保存函数，该函数会负责保存模型和 optimizers, fp16 的 state_dict；以及模型的保存（若 should_save_model 为 True）
+        断点重训的保存函数，该函数会负责保存优化器、fp16 状态和 sampler 的状态，以及模型的保存（若 ``should_save_model`` 为 ``True``）
 
-        :param folder: 保存断点重训的状态的文件夹；save_checkpoint 函数应该在下面新增两（一）个文件 的 FASTNLP_CHECKPOINT_FILENAME 文件与
-            FASTNLP_MODEL_FILENAME （如果 should_save_model 为 True ）。把 model 相关的内容放入到 FASTNLP_MODEL_FILENAME 文件
-            中，将传入的 states 以及自身产生其它状态一并保存在 FASTNLP_CHECKPOINT_FILENAME 里面。
-        :param states: 由 trainer 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态，Driver 应该只需要保存
-            该对象即可， Driver 应该不需要理解该对象，同时在 driver.load_checkpoint() 的时候，需要将 states 返回回去，load_checkpoint() 返回的值与这里的
-            传入的值保持一致。
+        :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` 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态。Driver 应该
+            只需要保存该对象而不需要理解该对象，同时在 :meth:`load_checkpoint` 的时候需要将 ``states`` 返回回去，返回的值与这里传入的值保持一致。
         :param dataloader: 正在使用的 dataloader，需要保存里面的状态使得之后可以从当前迭代的位置恢复。
-        :param only_state_dict: 是否只保存模型的参数，当 should_save_model 为 False ，该参数无效。
-        :param should_save_model: 是否应该保存模型，如果为False，Driver 将不负责 model 的保存。
+        :param only_state_dict: 是否只保存模型的参数，当 ``should_save_model`` 为 ``False`` ，该参数无效。
+        :param should_save_model: 是否应该保存模型，如果为 ``False`` ，Driver 将不负责 model 的保存。
         """
         raise NotImplementedError("Each specific driver should implemented its own `save_checkpoint` function.")
 
@@ -288,113 +295,115 @@ class Driver(ABC):
     def load_checkpoint(self, folder: Union[str, Path], dataloader, only_state_dict: bool = True, should_load_model: bool = True,
              **kwargs) -> Dict:
         r"""
-        断点重训的加载函数，注意该函数会负责读取数据，并且恢复 optimizers , fp16 的 state_dict 和 模型（根据 should_load_model ）和；
-        其它在 Driver.save_checkpoint() 函数中执行的保存操作，然后将一个 state 字典返回给 trainer （ 内容为Driver.save_checkpoint() 接受到的 states ）。
-
+        断点重训的加载函数，该函数会负责读取数据，并且恢复优化器 、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
+        :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: 读取的，当 should_save_model 为 False ，该参数无效。如果为 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 ，但在保存的状态中没有
+        :param should_load_model: 是否应该加载模型，如果为 ``False`` ，Driver 将不负责加载模型。若该参数为 ``True`` ，但在保存的状态中没有
             找到对应的模型状态，则报错。
-        :return: 需要返回 save_checkpoint 函数输入的 states 内容
+        :return: :meth:`save_checkpoint` 函数输入的 ``states`` 内容。除此之外，还返回的内容有：
+
+            * *dataloader* -- 根据传入的 ``dataloader`` 与读取出的状态设置为合理状态的 dataloader。在当前 ``dataloader`` 样本数与读取出的 sampler 样本数
+              不一致时报错。
+            * *batch_idx_in_epoch* -- :class:`int` 类型的数据，表明当前 epoch 进行到了第几个 batch 。请注意，该值不能仅通过保存的数据中读取的，因为前后两次运行的
+              ``batch_size`` 可能有变化，而应该符合以下等式::
 
-            * *dataloader* -- 返回的是根据传入的 dataloader 与 保存的状态一起设置为合理的状态，可以返回的对象与传入的dataloader是同一个。
-             在保存与当前传入 data sample 数目不一致时报错。
+                返回的 dataloader 还会产生的 batch 数量 + batch_idx_in_epoch = 原来不断点训练时的 batch 的总数
+              
+              由于 ``返回的 dataloader 还会产生的batch数`` 在 ``batch_size`` 与 ``drop_last`` 参数给定的情况下，无法改变，因此只能通过调整 ``batch_idx_in_epoch``
+              这个值来使等式成立。一个简单的计算原则如下：
 
-            * *batch_idx_in_epoch* -- 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)。
+                * 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)。
         """
-        raise NotImplementedError("Each specific driver should implemented its own `load_checkpoint` function.")
 
     @staticmethod
     def tensor_to_numeric(tensor, reduce: Optional[str] = None):
         r"""
         将一个 ``tensor`` 对象（仅处理当前 driver 使用的 tensor 即可）转换为 python 的 ``numeric`` 对象；如果 ``tensor`` 只包含一个
-        元素则返回 ``float`` 或 ``int``；
+        元素则返回 ``float`` 或 ``int``。
 
-        :param tensor: 需要被转换的 `tensor` 对象；
+        :param tensor: 需要被转换的 ``tensor`` 对象
         :param reduce: 可选 ``['sum', 'max', 'mea', 'min']``，如果不为 ``None`` 将使用该 ``reduce`` 方法来处理当前 ``tensor`` 再返回
-            ``float`` 或 ``int`` 对象；
-        :return: 转换后返回的结果；
+            :class:`float` 或 :class:`int` 对象
+        :return: 转换后返回的结果
         """
         raise NotImplementedError("Each specific driver should implemented its own `tensor_to_numeric` function.")
 
     @abstractmethod
     def set_model_mode(self, mode: str):
         r"""
-        设置模型为 `train` / `eval` 的模式；目的是为切换模型训练和推理（会关闭dropout等）模式；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
 
-        :param mode: 应为二者之一：["train", "eval"]；
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
 
     def unwrap_model(self):
         r"""
         保证用户拿到的模型一定是最原始的模型；
-        注意因为我们把保存模型的主要逻辑和代码移到了 `Driver` 中，因此在 `save_model` 函数中，一定要先调用此函数来保证我们保存的模型一定是
+        注意因为我们把保存模型的主要逻辑和代码移到了 `Driver` 中，因此在 :meth:`save_model` 函数中，一定要先调用此函数来保证我们保存的模型一定是
         最为原始的模型；
-        需要注意用户本身传入的模型就是经过类似 `torch.nn.DataParallel` 或者 `torch.nn.parallel.DistributedDataParallel` 包裹的模型，
-        因此在该函数内需要先判断模型的类别；
+        需要注意用户本身传入的模型就是经过类似 :class:`torch.nn.DataParallel` 或者 :class:`torch.nn.parallel.DistributedDataParallel` 包裹的模型，
+        因此在该函数内需要先判断模型的类别。
 
-        :return: 返回最原始的模型，例如没有被 `DistributedDataParallel` 包裹的模型；
+        :return: 最原始的模型，例如没有被 :class:`DistributedDataParallel` 包裹的模型。
         """
 
     @staticmethod
     def move_model_to_device(model, device):
         r"""
-        用来将模型转移到指定的 device 上；
-        之所以写成 `staticmethod`，是因为一方面在 `Driver` 中我们要使用 `unwrap_model` 来拿到最原始的模型，另一方面，在 `save_model`
-        中，我们需要先将模型移到 cpu 后，又再移到 gpu 上，因此不适宜在该函数内部调用 `unwrap_model`，而是将 model 作为该函数的参数；
+        用来将模型转移到指定的 ``device`` 上；
+        之所以写成 :class:`staticmethod`，是因为一方面在 `Driver` 中我们要使用 :meth:`unwrap_model` 来拿到最原始的模型，另一方面，在 :meth`save_model`
+        中，我们需要先将模型移到 cpu 后，又再移到 gpu 上，因此不适宜在该函数内部调用 :meth:`unwrap_model`，而是将 ``model`` 作为该函数的参数。
         """
 
     @abstractmethod
     def move_data_to_device(self, batch):
         r"""
-        将数据迁移到指定的机器上；batch 可能是 list 也可能 dict ，或其嵌套结构；
+        将数据迁移到指定的机器上；``batch`` 是包含了张量的数据集合，可以是 **List**、**Dict** 等嵌套类型。
 
-        :return: 将移动到指定机器上的 batch 对象返回；
+        :return: 移动到指定机器上的 ``batch`` 对象
         """
 
     def get_local_rank(self) -> int:
         r"""
-        返回当前的local_rank，本函数的返回值只在运行分布式训练的时候有实际含义；
+        返回当前的 ``local_rank``，本函数的返回值只在运行分布式训练的时候有实际含义。
 
-        :return: 一个整数值，表示当前进程在当前这台机器上的序号；
+        :return: 一个整数值，表示当前进程在当前这台机器上的序号
         """
         return 0
 
     def barrier(self):
         r"""
         用于在多进程工作时同步各进程的工作进度，运行快的进程运行到这里会等待运行慢的进程，只有所有进程都运行到此函数时，所有的进程才会继续运行；
-        仅在多分布式训练场景中有使用；
+        仅在多分布式训练场景中有使用。
 
-        注意，该函数的行为会受到 FASTNLP_NO_SYNC 的影响。仅当 FASTNLP_NO_SYNC 在 os.environ 中不存在，或小于 1 时才真的执行 barrier；
+        注意，该函数的行为会受到环境变量 ``FASTNLP_NO_SYNC`` 的影响。仅当 ``FASTNLP_NO_SYNC`` 在 ``os.environ`` 中不存在，或小于 **1** 时
+        才真的执行 :meth:`barrier`。
         """
 
     def is_distributed(self) -> bool:
         r"""
-        当前的 driver 实例是否是分布式的；
+        当前的 driver 实例是否是分布式的。
 
-        :return: 返回一个 bool 值，如果当前的 driver 实例是用于分布式的，那么返回 True；
+        :return: 一个 bool 值，如果当前的 driver 实例是用于分布式的，那么返回 ``True``
         """
         return False
 
     def on_exception(self):
         r"""
-        该函数用于在训练或者预测过程中出现错误时正确地关掉其它的进程，这一点是通过在多进程 driver 调用 open_subprocess 的时候将每一个进程
-        的 pid 记录下来，然后在出现错误后，由出现错误的进程手动地将其它进程 kill 掉；
+        该函数用于在训练或者预测过程中出现错误时正确地关掉其它的进程，这一点是通过在多进程 driver 调用 :meth:`open_subprocess` 的时候将每一个进程
+        的 pid 记录下来，然后在出现错误后，由出现错误的进程手动地将其它进程 kill 掉。
 
-        因此，每一个多进程 driver 如果想要该函数能够正确地执行，其需要在自己的 open_subprocess（开启多进程的函数）中正确地记录每一个进程的
-        pid 的信息；
+        因此，每一个多进程 driver 如果想要该函数能够正确地执行，其需要在自己的 :meth:`open_subprocess` （开启多进程的函数）中正确地记录每一个进程的
+        pid 的信息；单卡 driver 不需要这个函数。
         """
         # 单卡 driver 不需要这个函数；
         if self._pids is not None:
@@ -419,10 +428,10 @@ class Driver(ABC):
         从 ``src`` 端将 ``obj`` 对象（可能是 ``tensor``，可能是 ``object`` ）broadcast 到其它所有进程。如果是非 ``tensor`` 的对象会尝试使用 ``pickle`` 进行打包进行
         传输，然后再 ``dst`` 处再加载回来。仅在分布式的 ``driver`` 中有实际意义。
 
-        :param obj: obj，可能是 ``Tensor`` 或 嵌套类型的数据；
-        :param src: source 的 ``global rank``；
-        :param group: 所属的通信组；
-        :return: 输入的 ``obj``；
+        :param obj: obj，可能是 ``Tensor`` 或 嵌套类型的数据
+        :param src: source 的 ``global rank``
+        :param group: 所属的通信组
+        :return: 输入的 ``obj``
         """
         if not self.is_distributed():
             return obj
@@ -431,12 +440,12 @@ class Driver(ABC):
 
     def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
-        :param obj: 可以是 ``float/int/bool/np.ndarray/{}/[]/Tensor`` 等；
-        :param group: 用于不同进程之间互相通信的通信组；
-        :return: 返回值应该是 ``[obj0, obj1, ...]``，其中 ``obj1`` 是 ``rank0`` 上的对象，``obj1`` 是 ``rank1`` 上的对象；
+        :param obj: 可以是 ``float/int/bool/np.ndarray/{}/[]/Tensor`` 等类型的数据
+        :param group: 用于不同进程之间互相通信的通信组
+        :return: 返回值应该是 ``[obj0, obj1, ...]``，其中 ``obj0`` 是 ``rank0`` 上的对象，``obj1`` 是 ``rank1`` 上的对象。以此类推
         """
         if not self.is_distributed():
             return [obj]

fastNLP/core/drivers/jittor_driver/initialize_jittor_driver.py

@@ -18,9 +18,9 @@ def initialize_jittor_driver(driver: str, device: Union[str, int, List[int]], mo
 
         创建多卡的 driver
 
-    :param driver: 该参数的值应为以下之一：``["jittor"]``；
-    :param device: ``jittor`` 运行的设备；
-    :param model: 训练或者评测的具体的模型；
+    :param driver: 该参数的值应为以下之一：``["jittor"]``
+    :param device: ``jittor`` 运行的设备
+    :param model: 训练或者评测的具体的模型
     :param kwargs: 
 
     :return: :class:`~fastNLP.core.JittorSingleDriver` 或 :class:`~fastNLP.core.JittorMPIDriver` 实例；

fastNLP/core/drivers/jittor_driver/jittor_driver.py

@@ -40,19 +40,22 @@ __all__ = [
 
 class JittorDriver(Driver):
     r"""
-    ``Jittor`` 框架的 ``Driver``，是 ``JittorSingleDevice`` 和 ``JittorMPIDriver`` 的父类。
+    实现了 **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::
 
-        您可以在使用 ``JittorSingleDevice`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口；
+        您可以在使用 ``JittorSingleDriver`` 和 ``JittorMPIDriver`` 时使用 ``JittorDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **jittor** 模型；
-    :param fp16: 是否开启混合精度训练;
+    :param model: 训练时使用的 **jittor** 模型
+    :param fp16: 是否开启混合精度训练
     :param jittor_kwargs:
     """
     def __init__(self, model, fp16: bool = False, jittor_kwargs: Dict = None, **kwargs):
@@ -73,6 +76,11 @@ class JittorDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.JittorDataLoader`、 :class:`jittor.dataset.Dataset` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, (Dataset, JittorDataLoader, OverfitDataLoader)):
             raise TypeError(f"{Dataset} or {JittorDataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -87,14 +95,23 @@ class JittorDriver(Driver):
                                  f"not {type(each_optimizer)}.")
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             optimizer.step()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """        
         for optimizer in self.optimizers:
             optimizer.backward(loss)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.zero_grad()
 
@@ -102,8 +119,8 @@ class JittorDriver(Driver):
         r"""
         将模型保存到 ``filepath`` 中。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
-        :param only_state_dict: 在 **Jittor** 中，该参数无效，**Jittor** 仅支持保存模型的 ``state_dict``。
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 在 **Jittor** 中，该参数无效，因为 **Jittor** 仅支持保存模型的 ``state_dict``。
         """
         if not only_state_dict:
             logger.rank_zero_warning(
@@ -119,7 +136,7 @@ class JittorDriver(Driver):
         r"""
         加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
+        :param filepath: 保存文件的文件位置
         :param load_state_dict: 在 **Jittor** 中，该参数无效，**Jittor** 仅支持加载模型的 ``state_dict``。
         """
         if not only_state_dict:
@@ -133,6 +150,17 @@ class JittorDriver(Driver):
         model.load(filepath)
 
     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
@@ -185,7 +213,36 @@ class JittorDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     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 = jt.load(str(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME)))
 
         # 1. 加载 optimizers 的状态；
@@ -232,6 +289,11 @@ class JittorDriver(Driver):
         return states
 
     def get_evaluate_context(self):
+        r"""
+        返回一个不计算梯度的上下文环境用来对模型进行评测；
+
+        :return: 上下文对象 ``jittor.no_grad``
+        """
         return jt.no_grad
 
     @staticmethod
@@ -241,6 +303,12 @@ class JittorDriver(Driver):
         """
         ...
 
+    def move_data_to_device(self, batch: 'jt.Var'):
+        """
+        将数据迁移到指定的机器上；**jittor** 会自动为变量分配设备无需手动迁移，因此这个函数只是简单地返回 ``batch``。
+        """
+        return batch
+
     def move_data_to_device(self, batch):
         """
         将数据 ``batch`` 转移到指定的设备上。由于 **Jittor** 会自动为数据分配设备，因此该函数实际上无效。
@@ -250,11 +318,11 @@ class JittorDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce=None):
         r"""
-        将一个 :class:`jittor.Var` 对象转换为 转换成 python 中的数值类型；
+        将一个 :class:`jittor.Var` 对象转换为 转换成 python 中的数值类型。
 
-        :param tensor: :class:`jittor.Var` 类型的对象；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: :class:`jittor.Var` 类型的对象
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等
         """
         if tensor is None:
             return None
@@ -274,29 +342,45 @@ class JittorDriver(Driver):
         )
 
     def set_model_mode(self, mode: str):
+        r"""
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
+        """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
 
     @property
     def data_device(self):
-        return self.model_device
-
-    def move_data_to_device(self, batch: 'jt.Var'):
         """
-        **jittor** 暂时没有提供数据迁移的函数，因此这个函数只是简单地返回 **batch**
+        :return: 数据默认会被迁移到的设备
         """
-        return batch
+        return self.model_device
 
     def set_deterministic_dataloader(self, dataloader: Union["JittorDataLoader", "Dataset"]):
+        r"""
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 **jittor** 暂时不提供
+        该功能。
+        """
         ...
 
     def set_sampler_epoch(self, dataloader: Union["JittorDataLoader", "Dataset"], cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
 
     @staticmethod
     def get_dataloader_args(dataloader: Union["JittorDataLoader", "Dataset"]):
+        """
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
+        """
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/jittor_driver/mpi.py

@@ -145,13 +145,13 @@ class JittorMPIDriver(JittorDriver):
 
     def is_distributed(self):
         """
-        判断是否为分布式的 **Driver** ，在 ``JittorSingleDriver`` 中，返回 ``True``。
+        判断是否为分布式的 **Driver** ，在 ``JittorMPIDriver`` 中，返回 ``True``。
         """
         return True
 
     @property
     def data_device(self) -> str:
         """
-        :return: 数据所在的设备；
+        :return: 数据所在的设备。
         """
         return self.model_device

fastNLP/core/drivers/jittor_driver/single_device.py

@@ -25,16 +25,21 @@ class JittorSingleDriver(JittorDriver):
     r"""
     ``Jittor`` 框架下用于 ``cpu`` 和单卡 ``gpu`` 运算的 ``Driver``。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: 训练和模型所在的设备，在 **Jittor** 中，应当为以下值之一：``[None, 'cpu', 'gpu', 'cuda']``；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数。
+    :param device: 训练和模型所在的设备，在 **Jittor** 中，应当为以下值之一：``[None, 'cpu', 'gpu', 'cuda']``：
         
-        * 为 ``None`` 或 ``cpu`` 时
-         表示在 ``cpu`` 上进行训练；
-        * 为 ``gpu`` 或 ``cuda`` 时
-         表示在显卡设备上进行训练；
+        * 为 ``None`` 或 ``cpu`` 时，表示在 ``cpu`` 上进行训练；
+        * 为 ``gpu`` 或 ``cuda`` 时，表示在显卡设备上进行训练；
 
-    :param fp16: 是否开启 fp16；
+    :param fp16: 是否开启 fp16 混合精度训练。
     :param jittor_kwargs:
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(self, model, device=None, fp16: bool = False, jittor_kwargs: Dict = None, **kwargs):
@@ -50,7 +55,7 @@ class JittorSingleDriver(JittorDriver):
 
     def setup(self):
         r"""
-        初始化训练环境；根据传入的 ``device`` 值设置模型的训练场景为 ``cpu`` 或 ``gpu``；
+        初始化训练环境；根据传入的 ``device`` 值设置模型的训练场景为 ``cpu`` 或 ``gpu``。
         """
         if self.model_device in ["cpu", None]:
             jt.flags.use_cuda = 0   # 使用 cpu
@@ -130,16 +135,3 @@ class JittorSingleDriver(JittorDriver):
             return replace_batch_sampler(dataloader, batch_sampler)
         else:
             return dataloader
-
-    def unwrap_model(self):
-        """
-        返回训练使用的模型。
-        """
-        return self.model
-
-    @property
-    def data_device(self) -> str:
-        """
-        :return: 数据和模型所在的设备；
-        """
-        return self.model_device

fastNLP/core/drivers/jittor_driver/utils.py

@@ -31,7 +31,7 @@ def jittor_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tru
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min

fastNLP/core/drivers/oneflow_driver/__init__.py

@@ -1,15 +1,14 @@
 __all__ = [
-    "OneflowDDPDriver",
-    "OneflowSingleDriver",
     "OneflowDriver",
+    "OneflowSingleDriver",
+    "OneflowDDPDriver",
     "oneflow_seed_everything",
-    "optimizer_state_to_device"
 ]
 
 from .ddp import OneflowDDPDriver
 from .single_device import OneflowSingleDriver
 from .oneflow_driver import OneflowDriver
-from .utils import oneflow_seed_everything, optimizer_state_to_device
+from .utils import oneflow_seed_everything
 
 
 

fastNLP/core/drivers/oneflow_driver/ddp.py

@@ -35,7 +35,7 @@ class OneflowDDPDriver(OneflowDriver):
 
     .. note::
 
-        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练；
+        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练。
 
         ``OneflowDDPDriver`` 目前支持两种启动方式：
         
@@ -43,13 +43,20 @@ class OneflowDDPDriver(OneflowDriver):
             2. 用户将模型通过 ``DistributedDataParallel`` 处理后，通过运行 ``python -m oneflow.distributed.launch --nproc_per_node 2 train.py`` 启动；
 
         注意多机的启动强制要求用户在每一台机器上使用 ``python -m oneflow.distributed.launch`` 启动；因此我们不会在 ``OneflowDDPDriver`` 中保存
-        任何当前有多少台机器的信息；
+        任何当前有多少台机器的信息。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 该参数无效，**fastNLP** 会自动获取当前进程的设备；
-    :param fp16: 是否开启 fp16 训练；目前该参数无效；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param parallel_device: 该参数无效，**fastNLP** 会自动获取当前进程的设备
+    :param fp16: 是否开启 fp16 训练；目前该参数无效
     :param oneflow_kwargs: 
-        * *ddp_kwargs* -- 用于 ``DistributedDataParallel`` 的其它参数，详情可查阅 **oneflow** 的官方文档；
+        * *ddp_kwargs* -- 用于 ``DistributedDataParallel`` 的其它参数，详情可查阅 **oneflow** 的官方文档
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(
@@ -88,7 +95,7 @@ class OneflowDDPDriver(OneflowDriver):
 
     def setup(self):
         r"""
-        将模型用 ``DistributedDataParallel`` 进行处理；
+        将模型用 ``DistributedDataParallel`` 进行处理。
         """
         if self._has_setup:
             return
@@ -121,14 +128,23 @@ class OneflowDDPDriver(OneflowDriver):
 
     @property
     def master_address(self) -> str:
+        """
+        分布式训练中的地址 ``MASTER_ADDR``
+        """
         return os.environ.get("MASTER_ADDR")
 
     @property
     def master_port(self) -> str:
+        """
+        分布式训练使用的端口 ``MASTER_PORT``
+        """
         return os.environ.get("MASTER_PORT")
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WORLD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -137,6 +153,9 @@ class OneflowDDPDriver(OneflowDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -144,11 +163,18 @@ class OneflowDDPDriver(OneflowDriver):
         self._global_rank = rank
 
     @property
-    def local_rank(self) -> int:  # 这个不会受到 all_rank_call_context 的影响
+    def local_rank(self) -> int:
+        """
+        当前进程的局部编号 ``local_rank``
+        """
         return int(os.environ.get("LOCAL_RANK", 0))
 
     @property
     def data_device(self):
+        """
+        数据所在的设备。由于 **oneflow** 可以通过 :func:`oneflow.cuda.current_device` 获取当前进程的设备，因此
+        该属性和 ``model_device`` 表现相同。
+        """
         return self._data_device
 
     def set_dist_repro_dataloader(self, dataloader,
@@ -240,13 +266,13 @@ class OneflowDDPDriver(OneflowDriver):
 
     def is_global_zero(self):
         r"""
-        :return: 返回当前的进程是否在全局上是进程 0 ；
+        :return: 当前的进程是否在全局上是进程 0 
         """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；该功能暂时无效，返回一个空的上下文环境；
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；该功能暂时无效，返回一个空的上下文环境
         """
         # TODO 暂时没有在 oneflow 中找到类似的功能；
         from fastNLP.core.utils import nullcontext
@@ -255,68 +281,64 @@ class OneflowDDPDriver(OneflowDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原始模型；
+        :return: 使用的原始模型
         """
         return self.model
 
     def get_local_rank(self) -> int:
         r"""
-        :return: 返回当前进程局部的进程编号；
+        :return: 当前进程局部的进程编号
         """
         return self.local_rank
 
     def barrier(self):
         r"""
-        通过使用该函数来使得各个进程之间同步操作；
+        同步各个进程之间的操作
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             comm.barrier()
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``OneflowDDPDriver`` 来说，该函数一定返回 ``True``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``OneflowDDPDriver`` 来说，该函数一定返回 ``True``
         """
         return True
 
-    def broadcast_object(self, obj, src: int = 0, **kwargs):
+    def broadcast_object(self, obj, src: int = 0, group=None, **kwargs):
         r"""
-        从 src 端将 obj 对象（可能是 tensor ，可能是 object ）发送到 dst 处。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
-        传输，然后再 dst 处再加载回来。仅在分布式的 driver 中有实际意义。
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
 
         :param obj: obj，可能是 Tensor 或 嵌套类型的数据
-        :param int src: source 的 global rank 。
-        :param int dst: target 的 global rank，可以是多个目标 rank
-        :param group: 所属的 group
-        :return: 如果当前不是分布式 driver 直接返回输入的 obj 。如果当前 rank 是接收端（其 global rank 包含在了 dst 中），则返回
-            接收到的参数；如果是 source 端则返回发射的内容；既不是发送端、又不是接收端，则返回 None 。
+        :param src: 发送方的 ``global_rank``
+        :param group: 该参数无效
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC == 2 直接返回。
             return
         return fastnlp_oneflow_broadcast_object(obj, src, device=self.data_device)
 
-    def all_gather(self, obj) -> List:
+    def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
         example::
 
-            obj = {
-                'a': [1, 1],
-                'b': [[1, 2], [1, 2]],
-                'c': {
-                    'd': [1, 2]
-                }
-            }
-            ->
-            [
-                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-            ]
-
-        :param obj: 需要传输的对象，在每个rank上都应该保持相同的结构。
-        :param group:
-        :return:
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 该参数无效。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC 表示不执行
             return [obj]

fastNLP/core/drivers/oneflow_driver/dist_utils.py

@@ -13,6 +13,8 @@ if _NEED_IMPORT_ONEFLOW:
 
 PROTOCOL_VERSION = 1
 
+__all__ = []
+
 def _validate_output_list_for_rank(my_rank, dst, gather_list):
     if dst == my_rank:
         if not gather_list:
@@ -176,18 +178,15 @@ def fastnlp_oneflow_all_gather(obj: Any, device=None) ->List:
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/oneflow_driver/initialize_oneflow_driver.py

@@ -18,11 +18,11 @@ def initialize_oneflow_driver(driver: str, device: Optional[Union[str, "oneflow.
     r"""
     用来根据参数 ``driver` 和 ``device`` 来确定并且初始化一个具体的 ``Driver`` 实例然后返回回去；
 
-    :param driver: 该参数的值应为以下之一：``["oneflow"]``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
+    :param driver: 该参数的值应为以下之一：``["oneflow"]``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
     :param model: 训练或者评测的具体的模型；
 
-    :return: 返回一个 :class:`~fastNLP.core.OneflowSingleDriver` 或 :class:`~fastNLP.core.OneflowDDPDriver` 实例；
+    :return: 一个 :class:`~fastNLP.core.OneflowSingleDriver` 或 :class:`~fastNLP.core.OneflowDDPDriver` 实例；
     """
     # world_size 和 rank
     if FASTNLP_BACKEND_LAUNCH in os.environ:

fastNLP/core/drivers/oneflow_driver/oneflow_driver.py

@@ -36,17 +36,23 @@ from fastNLP.core.dataloaders import OverfitDataLoader
 
 class OneflowDriver(Driver):
     r"""
-    专属于 ``oneflow`` 的 ``driver``，是 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 的父类；
+    实现了 **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`` 提供的接口；
+        您可以在使用 ``OneflowSingleDriver`` 和 ``OneflowDDPDriver`` 时使用 ``OneflowDriver`` 提供的接口。
 
+    :param model: 训练使用的模型
+    :param fp16: 该参数暂时无效
+    :param oneflow_kwargs:
     """
     def __init__(self, model, fp16: Optional[bool] = False, oneflow_kwargs: Dict = None, **kwargs):
         super(OneflowDriver, self).__init__(model)
@@ -66,19 +72,33 @@ class OneflowDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.zero_grad(self.set_grad_to_none)
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         loss.backward()
         # self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.OneflowDataLoader`、 :class:`oneflow.utils.data.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -95,11 +115,11 @@ class OneflowDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce: str = None):
         r"""
-        将 ``oneflow.Tensor`` 转换成 python 中的数值类型；
+        将 ``oneflow.Tensor`` 转换成 python 中的数值类型。
 
-        :param tensor: ``oneflow.Tensor``；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: ``oneflow.Tensor``
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
 
         if tensor is None:
@@ -120,8 +140,9 @@ class OneflowDriver(Driver):
 
     def set_model_mode(self, mode: str):
         r"""
-        设置模型的状态是 ``train`` 还是 ``eval``；
-        :param mode: ``'train'`` 或 ``'eval'``；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
@@ -129,10 +150,10 @@ class OneflowDriver(Driver):
     @rank_zero_call
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = True, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存当前 driver 的模型到 ``filepath``。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，该参数只能为 ``True``；
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 是否只保存权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，该参数只能为 ``True``
         :return:
         """
         model = self.unwrap_model()
@@ -155,12 +176,11 @@ class OneflowDriver(Driver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重。
-        :param kwargs:
-        :return:
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重；如果使用 ``DistributedDataParallel`` 启动分布式训练的话，
+            该参数只能为 ``True``
         """
         model = self.unwrap_model()
         res = oneflow.load(filepath)
@@ -176,6 +196,17 @@ class OneflowDriver(Driver):
 
     @rank_zero_call
     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 的状态，从而适配训练或者评测环境；
 
@@ -280,6 +311,36 @@ class OneflowDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     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 的状态；
@@ -309,7 +370,9 @@ class OneflowDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        :return: 返回 ``oneflow.no_grad`` 这个 context；
+        返回一个不计算梯度的上下文环境用来对模型进行评测。
+
+        :return: 上下文对象 ``oneflow.no_grad``
         """
         return oneflow.no_grad
 
@@ -335,17 +398,17 @@ class OneflowDriver(Driver):
     @staticmethod
     def move_model_to_device(model: "oneflow.nn.Module", device: "oneflow.device"):
         r"""
-        将模型迁移到对应的设备上；
+        将模型迁移到对应的设备上。
         """
         if device is not None:
             model.to(device)
 
     def move_data_to_device(self, batch):
         """
-        将一个 batch 的数据迁移到对应的设备上；
+        将一个 ``batch`` 的数据迁移到对应的设备上。
 
-        :param batch: 一个 batch 的数据，可以是 ``list、dict`` 等；
-        :return:
+        :param batch: 包含 :class:`oneflow.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型
+        :return: 移动到指定机器后的 ``batch``
         """
         return oneflow_move_data_to_device(batch, self.data_device)
 
@@ -366,11 +429,20 @@ class OneflowDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader: "DataLoader"):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function,
                                                 rank=int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)))
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
@@ -378,9 +450,9 @@ class OneflowDriver(Driver):
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
         """
-        获取 dataloader 的 shuffle 和 drop_last 属性；
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
         """
-
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/oneflow_driver/single_device.py

@@ -21,12 +21,19 @@ from fastNLP.core.log import logger
 
 class OneflowSingleDriver(OneflowDriver):
     r"""
-    用于执行 ``oneflow`` 动态图 cpu 和 单卡 gpu 运算的 ``driver``；
+    用于执行 ``oneflow`` 动态图 cpu 和 单卡 gpu 运算的 ``driver``。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: oneflow.device，当前进程所使用的设备；
-    :param fp16: 是否开启 fp16；目前动态图的单卡下该参数无效；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param device: oneflow.device，当前进程所使用的设备
+    :param fp16: 是否开启 fp16；目前动态图的单卡下该参数无效。
     :param oneflow_kwargs:
+    :kwargs:
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为。
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
+
     """
 
     def __init__(self, model, device: "oneflow.device", fp16: bool = False, oneflow_kwargs: Dict = None, **kwargs):
@@ -54,7 +61,7 @@ class OneflowSingleDriver(OneflowDriver):
 
     def setup(self):
         r"""
-        将模型迁移到相应的设备上；
+        将模型迁移到相应的设备上。
         """
         if self.model_device is not None:
             self.model.to(self.model_device)
@@ -96,19 +103,19 @@ class OneflowSingleDriver(OneflowDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回模型
+        :return: 训练使用的模型
         """
         return self.model
 
     @property
     def data_device(self):
         r"""
-        :return: 数据和模型所在的设备；
+        :return: 数据和模型所在的设备。
         """
         return self.model_device
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，在 ``OneflowSingleDriver`` 中返回 ``False``；
+        :return: 当前使用的 driver 是否是分布式的 driver，在 ``OneflowSingleDriver`` 中返回 ``False``。
         """
         return False

fastNLP/core/drivers/oneflow_driver/utils.py

@@ -31,7 +31,6 @@ else:
 
 __all__ = [
     'oneflow_seed_everything',
-    'optimizer_state_to_device'
 ]
 
 def oneflow_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True) -> int:
@@ -40,7 +39,7 @@ def oneflow_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tr
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -263,11 +262,11 @@ def replace_batch_sampler(dataloader, new_batch_sampler):
 
 def optimizer_state_to_device(state, device):
     r"""
-    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备；
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
 
-    :param state: ``optimzier.state_dict()``；
-    :param device: 要迁移到的目的设备；
-    :return: 返回迁移后的新的 state_dict；
+    :param state: :func:`optimzier.state_dict` 获取的 state_dictt
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
     """
     new_state = {}
     for name, param in state.items():

fastNLP/core/drivers/paddle_driver/dist_utils.py

@@ -175,18 +175,15 @@ def fastnlp_paddle_all_gather(obj: Any, device=None, group=None) ->List:
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/paddle_driver/fleet.py

@@ -61,7 +61,7 @@ r"""
 .. note::
 
     多机的启动强制要求用户在每一台机器上使用 ``python -m paddle.distributed.launch`` 启动；因此我们不会在 ``PaddleFleetDriver``
-    中保存任何当前有多少台机器的信息；
+    中保存任何当前有多少台机器的信息。
 
 """
 import os
@@ -82,7 +82,6 @@ from fastNLP.core.utils import (
     auto_param_call,
     check_user_specific_params,
     is_in_paddle_dist,
-    is_in_paddle_dist,
     get_paddle_device_id,
 )
 from fastNLP.core.utils.paddle_utils import _convert_data_device
@@ -120,26 +119,26 @@ __all__ = [
 
 class PaddleFleetDriver(PaddleDriver):
     """
-    :param model: 训练使用的模型；
+    :param model: 训练使用的模型。
 
         * 如果不想自己初始化分布式环境，类型应为 :class:`paddle.nn.Layer`；
         * 如果已经在外面初始化了分布式环境，类型应为 :class:`paddle.DataParallel`；
 
     :param parallel_device: 多卡训练时使用的设备，必须是一个列表。
-        当使用 ``python -m paddle.distributed.launch`` 启动时，该参数无效；
+        当使用 ``python -m paddle.distributed.launch`` 启动时，该参数无效。
     :param is_pull_by_paddle_run: 标记当前进程是否为通过 ``python -m paddle.distributed.launch`` 启动的。
         这个参数仅在 :class:`~fastNLP.core.Trainer` 中初始化 driver 时使用
-    :param fp16: 是否开启混合精度训练；
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
         * *fleet_kwargs* -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括：
             
-            * *is_collective* -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况；
-            * *role_maker* -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``；
-            * 其它用于初始化 ``DataParallel`` 的参数；
-        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数;
+            * *is_collective* -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况。
+            * *role_maker* -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``。
+            * 其它用于初始化 ``DataParallel`` 的参数。
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数
     
     :kwargs:
-        * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
 
         .. note::
 
@@ -176,9 +175,9 @@ class PaddleFleetDriver(PaddleDriver):
         self.parallel_device = parallel_device
         # 在初始化时，如果发现 is_pull_by_paddle_run ，则将 parallel_device 设置成当前进程的gpu
         if is_pull_by_paddle_run:
-            self._model_device = parallel_device
+            self.model_device = parallel_device
         else:
-            self._model_device = parallel_device[self.local_rank]
+            self.model_device = parallel_device[self.local_rank]
 
         # 如果用户自己在外面初始化了并行模型；
         self.outside_fleet = False
@@ -311,6 +310,9 @@ class PaddleFleetDriver(PaddleDriver):
         self.global_rank = paddledist.get_rank()
 
     def barrier(self):
+        """
+        同步进程之间的操作
+        """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             paddledist.barrier()
 
@@ -329,6 +331,9 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WOLRD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -337,6 +342,9 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -345,20 +353,16 @@ class PaddleFleetDriver(PaddleDriver):
 
     @property
     def local_rank(self) -> int:
-        return int(os.getenv("PADDLE_RANK_IN_NODE", "0"))
-
-    @property
-    def model_device(self):
         """
-        :return: 模型所在的设备；
+        当前进程的局部编号 ``local_rank``
         """
-        return self._model_device
+        return int(os.getenv("PADDLE_RANK_IN_NODE", "0"))
 
     @property
     def data_device(self):
         """
-        :return: 数据所在的设备；由于 **PaddlePaddle** 可以通过环境变量获取当前进程的设备，因此该属性
-            和 ``model_device`` 表现相同；
+        数据所在的设备；由于 **PaddlePaddle** 可以通过环境变量获取当前进程的设备，因此该属性
+        和 ``model_device`` 表现相同。
         """
         return self.model_device
 
@@ -484,9 +488,15 @@ class PaddleFleetDriver(PaddleDriver):
             raise ValueError("Parameter `dist_sampler` can only be one of three values: ('dist', 'unrepeatdist', None).")
 
     def is_global_zero(self) -> bool:
+        r"""
+        :return: 当前的进程是否在全局上是进程 0
+        """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
+        r"""
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步。
+        """
         return self.model.no_sync
 
     def unwrap_model(self) -> "paddle.nn.Layer":
@@ -500,11 +510,14 @@ class PaddleFleetDriver(PaddleDriver):
             return _layers
 
     def get_local_rank(self) -> int:
+        r"""
+        :return: 当前进程局部的进程编号。
+        """
         return self.local_rank
 
     def is_distributed(self) -> bool:
         """
-        判断是否为分布式的 **Driver** ，在 ``PaddleFleetDriver`` 中，返回 ``True``。
+        :return: 当前使用的 driver 是否是分布式的 driver，在 ``PaddleFleetDriver`` 中，返回 ``True``。
         """
         return True
 
@@ -518,9 +531,39 @@ class PaddleFleetDriver(PaddleDriver):
                                 f"not {type(each_optimizer)}.")
 
     def broadcast_object(self, obj, src:int=0, group=None, **kwargs):
+        r"""
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
+
+        :param obj: obj，可能是 Tensor 或 嵌套类型的数据
+        :param src: 发送方的 ``global_rank``
+        :param group: 进程所在的通信组
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
+        """
         # 因为设置了CUDA_VISIBLE_DEVICES，可能会引起错误
         device = _convert_data_device(self.data_device)
         return fastnlp_paddle_broadcast_object(obj, src, device=device, group=group)
 
     def all_gather(self, obj, group=None) -> List:
+        r"""
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
+        pickle 进行序列化，接收到之后再反序列化。
+
+        example::
+
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 进程所在的通信组。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
+        """
         return fastnlp_paddle_all_gather(obj, group=group)

fastNLP/core/drivers/paddle_driver/initialize_paddle_driver.py

@@ -27,8 +27,8 @@ def initialize_paddle_driver(driver: str, device: Optional[Union[str, int, List[
     2. 如果 ``device`` 包含了多个设备，则返回一个 :class:`~fastNLP.core.PaddleFleetDriver` 实例，否则返回
     单卡的 :class:`~fastNLP.core.PaddleSingleDriver` 实例
 
-    :param driver: 使用的 ``driver`` 类型，在这个函数中仅支持 ``paddle``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
+    :param driver: 使用的 ``driver`` 类型，在这个函数中仅支持 ``paddle``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
     :param model: 训练或者评测的具体的模型；
 
     :return: 一个 :class:`~fastNLP.core.PaddleSingleDriver` 或 :class:`~fastNLP.core.PaddleFleetDriver` 实例；

fastNLP/core/drivers/paddle_driver/paddle_driver.py

@@ -47,28 +47,25 @@ if _NEED_IMPORT_PADDLE:
 
 class PaddleDriver(Driver):
     r"""
-    实现了 **PaddlePaddle** 框架训练功能的基本 Driver，实现了单卡和多卡情景下均需要实现的功能，以和 **fastNLP** 的
-    :class:`~fastNLP.core.Trainer` 兼容；通过这个 Driver，可以在 **fastNLP** 中实现从 **Pytorch** 框架到
-    **PaddlePaddle** 深度学习框架的切换。
+    实现了 **PaddlePaddle** 框架训练功能的基本 Driver。
 
     这个类被以下子类继承：
 
-    1. :class:`~fastNLP.core.drivers.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能；
-    2. :class:`~fastNLP.core.drivers.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能；
+    1. :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能；
+    2. :class:`~fastNLP.core.drivers.paddle_driver.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能；
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``PaddleSingleDriver`` 和 ``PaddleDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``PaddleSingleDriver`` 和 ``PaddleFleetDriver`` 时使用 ``PaddleDriver`` 提供的接口；
+        您可以在使用 ``PaddleSingleDriver`` 和 ``PaddleFleetDriver`` 时使用 ``PaddleDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **PaddlePaddle** 模型；
-    :param fp16: 是否开启混合精度训练；
+    :param model: 训练时使用的 **PaddlePaddle** 模型
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
-
     """
     def __init__(self, model: "paddle.nn.Layer", fp16: Optional[bool] = False, paddle_kwargs: Dict = None, **kwargs):
         if not isinstance(model, paddle.nn.Layer):
@@ -87,18 +84,32 @@ class PaddleDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             optimizer.clear_grad()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """  
         self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.PaddleDataLoader`、 :class:`paddle.io.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if dataloader.batch_size is None and dataloader.batch_sampler is None:
@@ -113,7 +124,7 @@ class PaddleDriver(Driver):
         r"""
         对于用户传入 trainer 的每一个 optimizer检测其合法性，必须为`paddle.optimizer.Optimizer`类型。
 
-        :param optimizers: 需要检测的 `optimizers`；
+        :param optimizers: 需要检测的 `optimizers`。
         """
         for each_optimizer in optimizers:
             if not isinstance(each_optimizer, Optimizer):
@@ -123,11 +134,11 @@ class PaddleDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce=None):
         r"""
-        将一个 :class:`paddle.Tensor` 对象转换为 转换成 python 中的数值类型；
+        将一个 :class:`paddle.Tensor` 对象转换为 转换成 python 中的数值类型。
 
-        :param tensor: :class:`paddle.Tensor` 类型的对象；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: :class:`paddle.Tensor` 类型的对象。
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
         if tensor is None:
             return None
@@ -148,6 +159,11 @@ class PaddleDriver(Driver):
         )
 
     def set_model_mode(self, mode: str):
+        r"""
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
+        """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
 
@@ -156,13 +172,13 @@ class PaddleDriver(Driver):
         r"""
         将模型保存到 ``filepath`` 中。
 
-        :param filepath: 保存文件的文件位置（需要包括文件名）；
+        :param filepath: 保存文件的文件位置（需要包括文件名）。
         :param only_state_dict: 是否只保存模型的 ``state_dict``；如果为 ``False``，则会调用 ``paddle.jit.save`` 
-            函数保存整个模型的参数，此时需要传入 ``input_spec`` 参数；
+            函数保存整个模型的参数，此时需要传入 ``input_spec`` 参数。
         :kwargs:
             * *input_spec* -- 描述存储模型 ``forward`` 方法的输入；
-            当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
-            进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_；
+              当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
+              进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_。
         """
         model = self.unwrap_model()
         if isinstance(filepath, Path):
@@ -178,6 +194,12 @@ class PaddleDriver(Driver):
             paddle.jit.save(model, filepath, input_spec)
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
+        """
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
+
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重。
+        """
         model = self.unwrap_model()
         if isinstance(filepath, Path):
             filepath = str(filepath)
@@ -192,25 +214,15 @@ class PaddleDriver(Driver):
     @rank_zero_call
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
         r"""
-        断点重训的保存函数，该函数会负责保存模型和 optimizers, fp16 的 state_dict；以及模型的保存（若 should_save_model 为 True）
-
-        :param folder: 保存断点重训的状态的文件夹；save 函数应该在下面新增两（一）个文件 的 FASTNLP_CHECKPOINT_FILENAME 文件与
-            FASTNLP_MODEL_FILENAME （如果 should_save_model 为 True ）。把 model 相关的内容放入到 FASTNLP_MODEL_FILENAME 文件中，
-            将传入的 states 以及自身产生其它状态一并保存在 FASTNLP_CHECKPOINT_FILENAME 里面。
-        :param states: 由 trainer 传入的一个字典，其中已经包含了为了实现断点重训所需要保存的其它对象的状态，Driver 应该只需要保存该对象即可，
-            Driver 应该不需要理解该对象，同时在 driver.load_checkpoint() 的时候，需要将 states 返回回去，load() 返回的值与这里的传入的值保持一致。
-        :param dataloader: 正在使用的 dataloader，需要保存里面的状态使得之后可以从当前迭代的位置恢复。
-        :param only_state_dict: 是否只保存模型的参数，当 should_save_model 为 False ，该参数无效。
-        :param should_save_model: 是否应该保存模型，如果为False，Driver 将不负责 model 的保存。
-        :kwargs:
-            * input_spec -- 描述存储模型 ``forward`` 方法的输入；
-            当 ``only_state_dict`` 为 ``False`` 时必须传入，否则加载时会报错。您可以通过 ``InputSpec`` 或者示例 ``Tensor``
-            进行描述。详细的使用方法可以参考 **PaddlePaddle** `关于 paddle.jit.save 函数的文档 <https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/jit/save_cn.html#save>`_；
-        
-        .. todo:
-
-            等 Driver 的文档写完
-
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``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 的状态，从而适配训练或者评测环境；
@@ -272,7 +284,36 @@ class PaddleDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
-        
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``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 = paddle.load(str(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME)))
 
         # 1. 加载 optimizers 的状态；
@@ -333,7 +374,7 @@ class PaddleDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        返回一个不计算梯度的环境用来对模型进行评测；
+        返回一个不计算梯度的环境用来对模型进行评测。
 
         :return: 上下文对象 ``paddle.no_grad``;
         """
@@ -342,14 +383,14 @@ class PaddleDriver(Driver):
     @staticmethod
     def move_model_to_device(model: "paddle.nn.Layer", device: Union[str, int, "paddle.CUDAPlace", "paddle.CPUPlace"]):
         r"""
-        用来将模型 ``model`` 转移到指定的设备上；
+        用来将模型 ``model`` 转移到指定的设备上。
 
         .. note::
 
             在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。
 
-        :param model: 需要进行转移的模型；
-        :param device: 目标设备；
+        :param model: 需要进行转移的模型。
+        :param device: 目标设备。
         """
         if device is not None:
             model.to(device)
@@ -362,8 +403,8 @@ class PaddleDriver(Driver):
 
             在 **Paddle** 中使用可能会引起因与设置的设备不一致而产生的问题，请注意。
 
-        :param batch: 包含 :class:`paddle.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型；
-        :return: 移动到指定机器后的 `batch``；
+        :param batch: 包含 :class:`paddle.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型。
+        :return: 移动到指定机器后的 ``batch``。
         """
         device = _convert_data_device(self.data_device)
         return paddle_move_data_to_device(batch, device)
@@ -387,10 +428,19 @@ class PaddleDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function, rank=self.global_rank)
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         if callable(getattr(dataloader.batch_sampler, "set_epoch", None)):
             dataloader.batch_sampler.set_epoch(cur_epoch_idx)
         elif callable(getattr(dataloader.batch_sampler.sampler, "set_epoch", None)):
@@ -398,7 +448,10 @@ class PaddleDriver(Driver):
 
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
-
+        """
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
+        """
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/paddle_driver/single_device.py

@@ -40,13 +40,13 @@ class PaddleSingleDriver(PaddleDriver):
     """
     实现了 **PaddlePaddle** 框架下在单卡或 ``cpu`` 环境下训练功能的 **Driver**。
 
-    :param model: 训练时使用的 **PaddlePaddle** 模型；
-    :param device: 训练使用的设备；
-    :param fp16: 是否开启混合精度训练；
+    :param model: 训练时使用的 **PaddlePaddle** 模型
+    :param device: 训练使用的设备
+    :param fp16: 是否开启混合精度训练
     :param paddle_kwargs:
-        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数;
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数。
     :kwargs:
-        * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
+        * *model_wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为。
 
         .. note::
 
@@ -155,19 +155,19 @@ class PaddleSingleDriver(PaddleDriver):
 
     def unwrap_model(self):
         """
-        返回训练使用的模型。
+        :return: 训练使用的模型。
         """
         return self.model
 
     @property
     def data_device(self) -> str:
         """
-        :return: 数据和模型所在的设备；
+        :return: 数据和模型所在的设备。
         """
         return self.model_device
 
     def is_distributed(self) -> bool:
         """
-        判断是否为分布式的 **Driver** ，在 ``PaddleSingleDriver`` 中，返回 ``False``。
+        :return 是否为分布式的 **Driver** ，在 ``PaddleSingleDriver`` 中，返回 ``False``。
         """
         return False

fastNLP/core/drivers/paddle_driver/utils.py

@@ -32,6 +32,7 @@ else:
 
 __all__ = [
     "paddle_seed_everything",
+    "optimizer_state_to_device",
 ]
 
 def paddle_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True) -> int:
@@ -40,7 +41,7 @@ def paddle_seed_everything(seed: int = None, add_global_rank_to_seed: bool = Tru
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -247,6 +248,13 @@ def replace_sampler(dataloader, new_sampler):
     return replace_batch_sampler(dataloader, new_batch_sampler)
 
 def optimizer_state_to_device(state, device):
+    r"""
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
+
+    :param state: :func:`optimzier.state_dict` 获取的 state_dictt
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
+    """
     new_state = {}
     for name, param in state.items():
         if isinstance(param, dict):

fastNLP/core/drivers/torch_driver/__init__.py

@@ -1,8 +1,8 @@
 __all__ = [
-    'TorchDDPDriver',
+    'TorchDriver',
     'TorchSingleDriver',
+    'TorchDDPDriver',
     'DeepSpeedDriver',
-    'TorchDriver',
     'torch_seed_everything',
     'optimizer_state_to_device'
 ]

fastNLP/core/drivers/torch_driver/ddp.py

@@ -164,11 +164,11 @@ from .utils import _check_dataloader_args_for_distributed
 
 class TorchDDPDriver(TorchDriver):
     r"""
-    ``TorchDDPDriver`` 通过开启多个进程，让每个进程单独使用一个 gpu 设备来实现分布式训练；
+    ``TorchDDPDriver`` 通过开启多个进程，让每个进程单独使用一个 gpu 设备来实现分布式训练。
 
     .. note::
 
-        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练；
+        您在绝大多数情况下不需要自己使用到该类，通过向 ``Trainer`` 传入正确的参数，您可以方便快速地部署您的分布式训练。
 
         ``TorchDDPDriver`` 目前支持的三种启动方式：
 
@@ -229,18 +229,24 @@ class TorchDDPDriver(TorchDriver):
             通过运行 ``python -m torch.distributed.launch --nproc_per_node 2 train.py`` 启动；
 
         注意多机的启动强制要求用户在每一台机器上使用 ``python -m torch.distributed.launch`` 启动；因此我们不会在 ``TorchDDPDriver`` 中保存
-        任何当前有多少台机器的信息；
+        任何当前有多少台机器的信息。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 用于分布式训练的 ``gpu`` 设备；
-    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的；
-    :param fp16: 是否开启 fp16 训练；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的
+    :param fp16: 是否开启 fp16 训练
     :param torch_kwargs: 
         * *ddp_kwargs* -- 用于在使用 ``TorchDDPDriver`` 时指定 ``DistributedDataParallel`` 初始化时的参数；例如传入
-            {'find_unused_parameters': True} 来解决有参数不参与前向运算导致的报错等；
-        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None；
-        * *non_blocking* -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking；
-        * *gradscaler_kwargs* -- 用于 fp16=True 时，提供给 ``torch.amp.cuda.GradScaler`` 的参数;
+          ``{'find_unused_parameters': True}`` 来解决有参数不参与前向运算导致的报错等
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
 
     def __init__(
@@ -329,7 +335,7 @@ class TorchDDPDriver(TorchDriver):
         r"""
         准备分布式环境，该函数主要做以下两件事情：
 
-            1. 开启多进程，每个 gpu 设备对应单独的一个进程；
+            1. 开启多进程，每个 ``gpu`` 设备对应单独的一个进程；
             2. 每个进程将模型迁移到自己对应的 ``gpu`` 设备上；然后使用 ``DistributedDataParallel`` 包裹模型；
         """
         if self._has_setup:
@@ -450,10 +456,16 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def master_address(self) -> str:
+        """
+        分布式训练中的地址 ``MASTER_ADDR``
+        """
         return os.environ.get("MASTER_ADDR", "127.0.0.1")
 
     @property
     def master_port(self) -> str:
+        """
+        分布式训练使用的端口 ``MASTER_PORT``
+        """
         if self.outside_ddp:
             return os.environ.get("MASTER_PORT")
         if self._master_port is None:
@@ -462,6 +474,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def world_size(self) -> int:
+        """
+        分布式训练的进程总数 ``WORLD_SIZE``
+        """
         return self._world_size
 
     @world_size.setter
@@ -470,6 +485,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def global_rank(self) -> int:
+        """
+        当前进程的全局编号 ``global_rank``
+        """
         return self._global_rank
 
     @global_rank.setter
@@ -478,6 +496,9 @@ class TorchDDPDriver(TorchDriver):
 
     @property
     def local_rank(self) -> int:  # 这个不会受到 all_rank_call_context 的影响
+        """
+        当前进程的局部编号 ``local_rank``
+        """
         return int(os.environ.get("LOCAL_RANK", 0))
 
     @property
@@ -609,20 +630,20 @@ class TorchDDPDriver(TorchDriver):
 
     def is_global_zero(self):
         r"""
-        :return: 返回当前的进程是否在全局上是进程 0 ；
+        :return: 当前的进程是否在全局上是进程 0 。
         """
         return self.global_rank == 0
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步。
         """
         # 注意此时的 model 是 "DistributedDataParallel" 对象；
         return self.model.no_sync
 
     def unwrap_model(self):
         r"""
-        :return: 返回没有经过 ``DistributedDataParallel`` 包裹的原始模型；
+        :return: 没有经过 ``DistributedDataParallel`` 包裹的原始模型。
         """
         _module = self.model.module
         if isinstance(_module, _DDPWrappingModel):
@@ -632,34 +653,33 @@ class TorchDDPDriver(TorchDriver):
 
     def get_local_rank(self) -> int:
         r"""
-        :return: 返回当前进程局部的进程编号；
+        :return: 当前进程局部的进程编号。
         """
         return self.local_rank
 
     def barrier(self):
         r"""
-        通过使用该函数来使得各个进程之间同步操作；
+        通过使用该函数来使得各个进程之间同步操作。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) < 1:  # 当 FASTNLP_NO_SYNC 小于 1 时实际执行
             torch.distributed.barrier(async_op=False)
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``TorchDDPDriver`` 来说，该函数一定返回 ``True``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``TorchDDPDriver`` 来说，该函数一定返回 ``True``。
         """
         return True
 
     def broadcast_object(self, obj, src: int = 0, group=None, **kwargs):
         r"""
-        从 src 端将 obj 对象（可能是 tensor ，可能是 object ）发送到 dst 处。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
-        传输，然后再 dst 处再加载回来。仅在分布式的 driver 中有实际意义。
+        从 ``src`` 端将 ``obj`` 对象（可能是 tensor ，可能是 object ）广播到其它进程。如果是非 tensor 的对象会尝试使用 pickle 进行打包进行
+        传输，然后在接收处处再加载回来。仅在分布式的 driver 中有实际意义。
 
         :param obj: obj，可能是 Tensor 或 嵌套类型的数据
-        :param int src: source 的 global rank 。
-        :param int dst: target 的 global rank，可以是多个目标 rank
-        :param group: 所属的 group
-        :return: 如果当前不是分布式 driver 直接返回输入的 obj 。如果当前 rank 是接收端（其 global rank 包含在了 dst 中），则返回
-            接收到的参数；如果是 source 端则返回发射的内容；既不是发送端、又不是接收端，则返回 None 。
+        :param src: 发送方的 ``global_rank``
+        :param group: 进程所在的通信组
+        :return: 如果当前 rank 是接收端，则返回接收到的参数；如果是 source 端则返回发送的内容。如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则
+            返回 ``None``
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC == 2 直接返回。
             return
@@ -667,27 +687,24 @@ class TorchDDPDriver(TorchDriver):
 
     def all_gather(self, obj, group) -> List:
         r"""
-        将 obj 互相传送到其它所有的 rank 上，其中 obj 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，尝试通过
+        将 ``obj`` 互相传送到其它所有的 rank 上，其中 ``obj`` 可能是 Tensor，也可能是嵌套结构的 object 。如果不是基础类型的数据，将会尝试通过
         pickle 进行序列化，接收到之后再反序列化。
 
         example::
 
-            obj = {
-                'a': [1, 1],
-                'b': [[1, 2], [1, 2]],
-                'c': {
-                    'd': [1, 2]
-                }
-            }
-            ->
-            [
-                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-            ]
-
-        :param obj: 需要传输的对象，在每个rank上都应该保持相同的结构。
-        :param group:
-        :return:
+            >>> # rank 0
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+            >>> # rank 1
+            >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            >>> # after all_gather():
+            >>> result = [
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                    {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+                ]
+
+        :param obj: 需要传输的对象，在每个 rank 上都应该保持相同的结构。
+        :param group: 进程所在的通信组。
+        :return: 所有 rank 发送的 ``obj`` 聚合在一起的内容；如果环境变量 ``FASTNLP_NO_SYNC`` 为 **2** 则不会执行，直接返回 ``[obj]`` 。
         """
         if int(os.environ.get(FASTNLP_NO_SYNC, 0)) == 2:  # 如果 FASTNLP_NO_SYNC 表示不执行
             return [obj]
@@ -701,7 +718,7 @@ class TorchDDPDriver(TorchDriver):
 def find_free_network_port() -> str:
     """
     在 localhost 上找到一个空闲端口；
-    当我们不想连接到真正的主节点但必须设置“MASTER_PORT”环境变量时在单节点训练中很有用；
+    当我们不想连接到真正的主节点但必须设置“MASTER_PORT”环境变量时在单节点训练中很有用。
     """
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     s.bind(("", 0))

fastNLP/core/drivers/torch_driver/deepspeed.py

@@ -83,26 +83,36 @@ class DeepSpeedDriver(TorchDDPDriver):
                 )
                 trainer.run()
 
-            通过运行 ``deepspeed train.py`` 启动；
+            通过运行 ``deepspeed train.py`` 启动。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param parallel_device: 用于分布式训练的 ``gpu`` 设备；
-    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的；
-    :param fp16: 是否开启 fp16 训练；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数。
+    :param parallel_device: 用于分布式训练的 ``gpu`` 设备。
+    :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的。
+    :param fp16: 是否开启 fp16 训练。
     :param deepspeed_kwargs: 
         * *strategy* -- 使用 ZeRO 优化的策略，默认为 ``deepspeed``；目前仅支持以下值：
 
             * ``deepspeed`` -- 使用 ZeRO 的第二阶段，等同于 ``deepspeed_stage_2``；
             * ``deepspeed_stage_1`` -- 使用 ZeRO 的第一阶段，仅将 ``optimizer`` 的状态分散到不同设备上；
-            * ``deepspeed_stage_2`` -- 使用 ZeRO 的第二阶段，将 ``optimizer`` 和**梯度**分散到不同设备上；
+            * ``deepspeed_stage_2`` -- 使用 ZeRO 的第二阶段，将 ``optimizer`` 和 **梯度** 分散到不同设备上；
             * ``deepspeed_stage_2_offload`` -- 使用 ZeRO 的第二阶段，并且借助 cpu 的内存来进一步节约显存；
-            * ``deepspeed_stage_3`` -- 使用 ZeRO 的第三阶段，将 ``optimizer`` 、**梯度**和**模型**分散到不同设备上；
+            * ``deepspeed_stage_3`` -- 使用 ZeRO 的第三阶段，将 ``optimizer`` 、**梯度** 和 **模型** 分散到不同设备上；
             * ``deepspeed_stage_3_offload`` -- 使用 ZeRO 的第三阶段，并且借助 cpu 的内存来进一步节约显存；
             * ``deepspeed_stage_3_offload_nvme`` -- 使用 ZeRO 的第三阶段，并且借助 NVMe 硬盘来进一步节约显存；
-        * *logging_level* -- ``deepspeed`` 库的日志等级，默认为 **logging.ERROR**；
+        * *logging_level* -- ``deepspeed`` 库的日志等级，默认为 **logging.ERROR**。
         * *config* -- ``deepspeed`` 的各项设置；**FastNLP** 允许用户传入自己的设置以增强灵活性，但这会使参数
           中的 ``optimizer`` 、``strategy`` 、 ``fp16`` 等失效，即当这个参数存在时，**FastNLP** 会用该参数覆盖
-          其它的设置；
+          其它的设置。
+    :kwargs:
+        * *accumulation_steps* -- 即在 :class:`~fastNLP.core.controllers.Trainer` 传入的 ``accumulation_steps`` 。 deepspeed 会将 ``config`` 的
+          ``gradient_accumulation_steps`` 设置为该值。
+        * *train_dataloader* -- 即在 :class:`~fastNLP.core.controllers.Trainer` 传入的 ``train_dataloader`` 。 ``deepspeed`` 需要通过它来获取
+          数据的 ``batch_size`` 用于设置 ``train_micro_batch_size_per_gpu`` 。如果没有传入的话，则会设置为 **1** 。
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
     # TODO fp16 load_config
     def __init__(
@@ -357,28 +367,38 @@ class DeepSpeedDriver(TorchDDPDriver):
                 self.config["amp"] = {"enabled": True, "opt_level": "O1"}
 
     def zero_grad(self):
+        """
+        进行梯度置零操作；由于 :meth:`DeepSpeedEngine.step` 包含了 :meth:`zero_step` 的功能，因此该接口实际无意义。
+        """
         # DeepSpeedEngine.step 包含了 zero_grad 功能
         pass
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         self.model.backward(loss)
 
     def step(self):
+        """
+        更新模型的参数
+        """
         self.model.step()
 
     def get_model_no_sync_context(self):
         r"""
-        :return: 返回一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；在 ``deepspeed`` 中，返回一个空的上下文
+        :return: 一个 ``context`` 上下文环境，用于关闭各个进程之间的同步；在 ``deepspeed`` 中，返回一个空的上下文
         """
         # 注意此时的 model 是 "DistributedDataParallel" 对象；
         return nullcontext
 
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = False, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存的模型到 ``filepath`` 中。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；在 ``DeepSpeedDriver`` 中该参数无效；
+        :param filepath: 文件路径
+        :param only_state_dict: 是否只保存权重；在 ``DeepSpeedDriver`` 中该参数无效。
+        :param kwargs: 需要传入 **deepspeed** 模型 :meth:`save_checkpoint` 的其它参数。
         :return:
         """
         # deepspeed engine 要求在每个 rank 都调用 save_checkpoint，故去掉了 rank_zero_call 装饰器
@@ -398,11 +418,11 @@ class DeepSpeedDriver(TorchDDPDriver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = False, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        从 ``filepath`` 中加载权重并赋值到当前 driver 的模型上。
 
         :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重；在 ``DeepSpeedDriver`` 中该参数无效；
-        :param kwargs:
+        :param load_state_dict: 保存的内容是否只是权重；在 ``DeepSpeedDriver`` 中该参数无效。
+        :param kwargs: 需要传入 **deepspeed** 模型 :meth:`load_checkpoint` 的其它参数。
         :return:
         """
         if not only_state_dict:
@@ -411,6 +431,17 @@ class DeepSpeedDriver(TorchDDPDriver):
         self.model.load_checkpoint(filepath, **kwargs)
 
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``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 的保存。
+        """
         # deepspeed engine 要求在每个 rank 都调用 save_checkpoint，故去掉了 rank_zero_call 装饰器
         # 1. 保存 sampler 的状态
         num_consumed_batches = states.pop('num_consumed_batches')
@@ -425,6 +456,36 @@ class DeepSpeedDriver(TorchDDPDriver):
                                     client_state=states)
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``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)。
+        """
         # 1. 加载模型状态；
         if not should_load_model:
             logger.rank_zero_warning("Loading checkpoint without model is not allowed for `DeepSpeedDriver`, "
@@ -442,4 +503,7 @@ class DeepSpeedDriver(TorchDDPDriver):
 
     @property
     def stage_3(self) -> bool:
+        """
+        判断是否为第三阶段的 ZeRO 优化
+        """
         return self.config.get("zero_optimization") and self.config.get("zero_optimization").get("stage") == 3

fastNLP/core/drivers/torch_driver/dist_utils.py

@@ -22,6 +22,7 @@ if _NEED_IMPORT_TORCH:
 
 from fastNLP.core.utils import apply_to_collection
 
+__all__ = []
 
 def _validate_output_list_for_rank(my_rank, dst, gather_list):
     if dst == my_rank:
@@ -148,7 +149,7 @@ def send_recv_object(obj, src, cur_rank, device, group=None, tag=0):
     r"""
     pytorch 中的单点对多点的分发函数；
 
-    例如将进程 0 上的对象 object 分发到其它进程上；
+    例如将进程 0 上的对象 object 分发到其它进程上。
 
     Example::
 
@@ -158,11 +159,11 @@ def send_recv_object(obj, src, cur_rank, device, group=None, tag=0):
 
         send_recv_object(object, 0, cur_rank, local_device)
 
-    :param obj: 一个可以序列化的 python 对象；
-    :param src: 从哪一个 rank 上发送到其它 rank；
-    :param cur_rank: 当前的进程的 rank 序号；
-    :param device: 当前的进程所在的设备；
-    :param group: 通信组，默认为 None；
+    :param obj: 一个可以序列化的 python 对象。
+    :param src: 从哪一个 rank 上发送到其它 rank。
+    :param cur_rank: 当前的进程的 rank 序号。
+    :param device: 当前的进程所在的设备。
+    :param group: 通信组，默认为 None。
     :param tag: 将发送与远程接收匹配的标记；
     :return:
     """
@@ -198,18 +199,15 @@ def fastnlp_torch_all_gather(obj: Any, device=None, group=DEFAULT_TORCH_GROUP) -
 
     example::
 
-        obj = {
-            'a': [1, 1],
-            'b': [[1, 2], [1, 2]],
-            'c': {
-                'd': [1, 2]
-            }
-        }
-        ->
-        [
-            {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
-            {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
-        ]
+        >>> # rank 0
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 1}}
+        >>> # rank 1
+        >>> obj = {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+        >>> # after all_gather():
+        >>> result = [
+                {'a': 1, 'b':[1, 2], 'c':{'d': 1}},
+                {'a': 1, 'b':[1, 2], 'c':{'d': 2}}
+            ]
 
     :param obj: 任意结构的数据，如果为 tensor ，需要保证每个显卡上的 tensor 的形状是一样的。如果传入的是非 tensor 对象都将直接进行
         序列化之后进行传输。

fastNLP/core/drivers/torch_driver/fairscale.py

@@ -301,6 +301,6 @@ class FairScaleDriver(TorchDDPDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原本的模型，例如没有被 ``DataParallel`` 包裹；
+        :return: 原本的模型，例如没有被 ``DataParallel`` 包裹；
         """
         return self.model.module.model

fastNLP/core/drivers/torch_driver/initialize_torch_driver.py

@@ -21,11 +21,15 @@ def initialize_torch_driver(driver: str, device: Optional[Union[str, "torch.devi
     r"""
     用来根据参数 ``driver` 和 ``device`` 来确定并且初始化一个具体的 ``Driver`` 实例然后返回回去；
 
-    :param driver: 该参数的值应为以下之一：``["torch", "fairscale", "deepspeed"]``；
-    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致；
-    :param model: 训练或者评测的具体的模型；
+    :param driver: 该参数的值应为以下之一：``["torch", "fairscale", "deepspeed"]``
+    :param device: 该参数的格式与 ``Trainer`` 对参数 ``device`` 的要求一致
+    :param model: 训练或者评测的具体的模型
 
-    :return: 返回一个 :class:`~fastNLP.core.TorchSingleDriver` 或 :class:`~fastNLP.core.TorchDDPDriver` 实例；
+    :return: 下列类型之一的实例：
+        * :class:`~fastNLP.core.drivers.torch_driver.TorchSingleDriver` 
+        * :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver` 
+        * :class:`~fastNLP.core.drivers.torch_driver.DeepSpeedDriver`
+        * :class:`~fastNLP.core.drivers.torch_driver.FairScaleDriver`
     """
     if parse_version(torch.__version__) < parse_version('1.6'):
         raise RuntimeError(f"Pytorch(current version:{torch.__version__}) need to be older than 1.6.")

fastNLP/core/drivers/torch_driver/single_device.py

@@ -26,19 +26,25 @@ from fastNLP.core.log import logger
 
 class TorchSingleDriver(TorchDriver):
     r"""
-    ``TorchSingleDriver`` 是用于 cpu 和 单卡 gpu 运算的 ``driver``；
+    ``TorchSingleDriver`` 是用于 cpu 和 单卡 gpu 运算的 ``driver``。
 
     .. note::
 
-        如果您希望使用 ``DataParallel`` 来训练您的模型，您应当自己在 ``Trainer`` 初始化之前初始化好 ``DataParallel``，然后将其传入 ``Trainer`` 中；
+        如果您希望使用 ``DataParallel`` 来训练您的模型，您应当自己在 ``Trainer`` 初始化之前初始化好 ``DataParallel``，然后将其传入 ``Trainer`` 中。
 
-    :param model: 传入给 ``Trainer`` 的 ``model`` 参数；
-    :param device: torch.device，当前进程所使用的设备；
-    :param fp16: 是否开启 fp16；
+    :param model: 传入给 ``Trainer`` 的 ``model`` 参数
+    :param device: torch.device，当前进程所使用的设备
+    :param fp16: 是否开启 fp16
     :param torch_kwargs:
-        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 None；
-        * *non_blocking* -- 表示用于 pytorch 的 tensor 的 to 方法的参数 non_blocking；
-        * *gradscaler_kwargs* -- 用于 fp16=True 时，提供给 ``torch.amp.cuda.GradScaler`` 的参数;
+        * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
+        * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
+    :kwargs:
+        * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
+
+        .. note::
+
+            关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
     """
 
     def __init__(self, model, device: "torch.device", fp16: bool = False, torch_kwargs: Dict = None, **kwargs):
@@ -69,7 +75,7 @@ class TorchSingleDriver(TorchDriver):
 
     def setup(self):
         r"""
-        将模型迁移到相应的设备上；
+        将模型迁移到相应的设备上。
         """
         if self.model_device is not None:
             self.model.to(self.model_device)
@@ -153,7 +159,7 @@ class TorchSingleDriver(TorchDriver):
 
     def unwrap_model(self):
         r"""
-        :return: 返回原本的模型，例如没有被 ``DataParallel`` 包裹；
+        :return: 原本的模型，该函数可以取出被 ``DataParallel`` 包裹的模型
         """
         if isinstance(self.model, torch.nn.DataParallel) or \
                 isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
@@ -164,12 +170,12 @@ class TorchSingleDriver(TorchDriver):
     @property
     def data_device(self):
         r"""
-        注意单卡模式下使用 ``driver.data_device`` 等价于使用 ``driver.model_device``；
+        数据和模型所在的设备
         """
         return self.model_device
 
     def is_distributed(self):
         r"""
-        :return: 返回当前使用的 driver 是否是分布式的 driver，对于 ``TorchSingleDriver`` 来说直接返回 ``False``；
+        :return: 当前使用的 driver 是否是分布式的 driver，对于 ``TorchSingleDriver`` 来说直接返回 ``False``
         """
         return False

fastNLP/core/drivers/torch_driver/torch_driver.py

@@ -36,18 +36,18 @@ from fastNLP.core.dataloaders import OverfitDataLoader
 
 class TorchDriver(Driver):
     r"""
-    专属于 ``pytorch`` 的 ``driver``，是 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 的父类；
+    专属于 ``pytorch`` 的 ``driver``，是 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 的父类。
 
     .. warning::
 
         您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``TorchSingleDriver`` 和 ``TorchDDPDriver``，而不是
-        该类本身；
+        该类本身。
 
     .. note::
 
-        您可以在使用 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 时使用 ``TorchDriver`` 提供的接口；
+        您可以在使用 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 时使用 ``TorchDriver`` 提供的接口。
 
-    :param model: 训练时使用的 **pytorch** 模型；
+    :param model: 训练时使用的 **pytorch** 模型。
     :param fp16: 是否开启混合精度训练;
     :param torch_kwargs:
     """
@@ -70,6 +70,9 @@ class TorchDriver(Driver):
         self.wo_auto_param_call = kwargs.get("model_wo_auto_param_call", False)
 
     def zero_grad(self):
+        """
+        实现梯度置零的过程
+        """
         for optimizer in self.optimizers:
             self._clear_grad(optimizer, self.set_grad_to_none)
 
@@ -88,14 +91,25 @@ class TorchDriver(Driver):
                         p.grad.zero_()
 
     def backward(self, loss):
+        """
+        对 ``loss`` 进行反向传播
+        """
         self.grad_scaler.scale(loss).backward()
 
     def step(self):
+        r"""
+        实现参数的优化更新过程
+        """
         for optimizer in self.optimizers:
             self.grad_scaler.step(optimizer)
             self.grad_scaler.update()
 
     def check_dataloader_legality(self, dataloader):
+        """
+        检测 DataLoader 是否合法。支持的类型包括 :class:`~fastNLP.core.dataloaders.TorchDataLoader`、 :class:`torch.utils.data.DataLoader` 。
+
+        :param dataloder:
+        """
         if not isinstance(dataloader, DataLoader) and not isinstance(dataloader, OverfitDataLoader):
             raise TypeError(f"{DataLoader} is expected, instead of `{type(dataloader)}`")
         if len(dataloader) == 0:
@@ -112,11 +126,11 @@ class TorchDriver(Driver):
     @staticmethod
     def tensor_to_numeric(tensor, reduce: str = None):
         r"""
-        将 ``torch.Tensor`` 转换成 python 中的数值类型；
+        将 ``torch.Tensor`` 转换成 python 中的数值类型。
 
-        :param tensor: ``torch.Tensor``；
-        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``；
-        :return: 返回一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等；
+        :param tensor: ``torch.Tensor``。
+        :param reduce: 当 tensor 是一个多数值的张量时，应当使用何种归一化操作来转换成单一数值，应当为以下类型之一：``['max', 'min', 'sum', 'mean']``。
+        :return: 一个单一数值，其数值类型是 python 中的基本的数值类型，例如 ``int，float`` 等。
         """
 
         if tensor is None:
@@ -137,8 +151,9 @@ class TorchDriver(Driver):
 
     def set_model_mode(self, mode: str):
         r"""
-        设置模型的状态是 ``train`` 还是 ``eval``；
-        :param mode: ``train`` 或者 ``eval``；
+        设置模型为 ``train`` 或 ``eval`` 的模式；目的是为切换模型的训练和推理（会关闭 dropout 等）模式。
+
+        :param mode: 应为二者之一：``["train", "eval"]``
         """
         assert mode in {"train", "eval"}
         getattr(self.model, mode)()
@@ -146,10 +161,10 @@ class TorchDriver(Driver):
     @rank_zero_call
     def save_model(self, filepath: Union[str, Path], only_state_dict: bool = True, **kwargs):
         """
-        保存当前 driver 的模型到 folder 下。
+        保存当前 driver 的模型到 ``filepath``。
 
-        :param filepath: 保存到哪个文件夹；
-        :param only_state_dict: 是否只保存权重；
+        :param filepath: 保存文件的文件位置
+        :param only_state_dict: 是否只保存权重
         :return:
         """
         model = self.unwrap_model()
@@ -169,12 +184,10 @@ class TorchDriver(Driver):
 
     def load_model(self, filepath: Union[Path, str], only_state_dict: bool = True, **kwargs):
         """
-        从 folder 中加载权重并赋值到当前 driver 的模型上。
+        加载模型的函数；将 ``filepath`` 中的模型加载并赋值给当前 ``model`` 。
 
-        :param filepath: 加载权重或模型的路径
-        :param load_state_dict: 保存的内容是否只是权重。
-        :param kwargs:
-        :return:
+        :param filepath: 保存文件的文件位置
+        :param load_state_dict: 保存的内容是否只是权重
         """
         model = self.unwrap_model()
         res = torch.load(filepath, map_location='cpu')
@@ -190,6 +203,17 @@ class TorchDriver(Driver):
 
     @rank_zero_call
     def save_checkpoint(self, folder: Path, states: Dict, dataloader, only_state_dict: bool = True, should_save_model: bool = True, **kwargs):
+        r"""
+        断点重训的保存函数，该函数会负责保存 **优化器** 、 **sampler** 和 **fp16** 的状态，以及 **模型** （若 ``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 的状态，从而适配训练或者评测环境；
 
@@ -297,6 +321,36 @@ class TorchDriver(Driver):
         logger.debug("Load optimizer state dict.")
 
     def load_checkpoint(self, folder: Path, dataloader, only_state_dict: bool = True, should_load_model: bool = True, **kwargs) -> Dict:
+        r"""
+        断点重训的加载函数，该函数会负责读取数据，并且恢复 **优化器** 、**sampler** 、 **fp16** 的状态和 **模型** （如果 ``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 = torch.load(folder.joinpath(FASTNLP_CHECKPOINT_FILENAME))
 
         # 1. 加载 optimizers 的状态；
@@ -326,29 +380,33 @@ class TorchDriver(Driver):
 
     def get_evaluate_context(self):
         r"""
-        :return: 返回 ``torch.no_grad`` 这个 context；
+        返回一个不计算梯度的上下文环境用来对模型进行评测。
+
+        :return: 上下文环境 ``torch.no_grad`` 
         """
         return torch.no_grad
 
     @staticmethod
     def move_model_to_device(model: "torch.nn.Module", device: "torch.device"):
         r"""
-        将模型迁移到对应的设备上；
+        将模型迁移到对应的设备上
         """
         if device is not None:
             model.to(device)
 
     def move_data_to_device(self, batch):
         """
-        将一个 batch 的数据迁移到对应的设备上；
+        将一个 ``batch`` 的数据迁移到对应的设备上
 
-        :param batch: 一个 batch 的数据，可以是 ``list、dict`` 等；
-        :return:
+        :param batch: 包含 :class:`torch.Tensor` 的数据集合，可以是 **List**、**Dict** 等嵌套类型
+        :return: 移动到指定机器后的 ``batch``
         """
         return torch_move_data_to_device(batch, self.data_device, self.non_blocking)
 
     @staticmethod
     def worker_init_function(worker_id: int, rank: Optional[int] = None) -> None:  # pragma: no cover
+        """
+        """
         """The worker_init_fn that Lightning automatically adds to your dataloader if you previously set the seed
         with ``seed_everything(seed, workers=True)``.
 
@@ -371,11 +429,20 @@ class TorchDriver(Driver):
         random.seed(stdlib_seed)
 
     def set_deterministic_dataloader(self, dataloader: "DataLoader"):
+        """
+        为了确定性训练要对 ``dataloader`` 进行修改，保证在确定随机数种子后，每次重新训练得到的结果是一样的。 
+        """
         if dataloader.worker_init_fn is None:
             dataloader.worker_init_fn = partial(self.worker_init_function,
                                                 rank=int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)))
 
     def set_sampler_epoch(self, dataloader: "DataLoader", cur_epoch_idx: int):
+        r"""
+        对于分布式的 ``sampler``，需要在每一个 ``epoch`` 前设置随机数种子，来保证每一个进程上的 ``shuffle`` 是一样的。
+
+        :param dataloader: 需要设置 ``epoch`` 的 ``dataloader``
+        :param cur_epoch_idx: 当前是第几个 ``epoch``
+        """
         # 保证 ddp 训练时的 shuffle=True 时的正确性，因为需要保证每一个进程上的 sampler 的shuffle 的随机数种子是一样的；
         if callable(getattr(dataloader.sampler, "set_epoch", None)):
             dataloader.sampler.set_epoch(cur_epoch_idx)
@@ -383,9 +450,9 @@ class TorchDriver(Driver):
     @staticmethod
     def get_dataloader_args(dataloader: "DataLoader"):
         """
-        获取 dataloader 的 shuffle 和 drop_last 属性；
+        从 ``dataloader`` 中获取参数 ``dataset``, ``batch_sampler``, ``sampler``, ``batch_size``, ``shuffle`` 
+        和 ``drop_last`` 。
         """
-
         @dataclass
         class Res:
             dataset: Optional[Dataset] = None

fastNLP/core/drivers/torch_driver/utils.py

@@ -42,7 +42,7 @@ def torch_seed_everything(seed: int = None, add_global_rank_to_seed: bool = True
 
     :param seed: 全局随机状态的整数值种子。如果为 ``None`` 则会根据时间戳生成一个种子。
     :param add_global_rank_to_seed: 在分布式训练中，是否在不同 **rank** 中使用不同的随机数。
-        当设置为 ``True`` 时，**FastNLP** 会将种子加上当前的 ``global_rank``。
+        当设置为 ``True`` 时，**fastNLP** 会将种子加上当前的 ``global_rank``。
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -290,11 +290,11 @@ def replace_batch_sampler(dataloader, new_batch_sampler):
 
 def optimizer_state_to_device(state, device):
     r"""
-    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备；
+    将一个 ``optimizer`` 的 ``state_dict`` 迁移到对应的设备。
 
-    :param state: ``optimzier.state_dict()``；
-    :param device: 要迁移到的目的设备；
-    :return: 返回迁移后的新的 state_dict；
+    :param state: ``optimzier.state_dict()``。
+    :param device: 要迁移到的目的设备。
+    :return: 迁移后的新的 state_dict。
     """
     new_state = {}
     for name, param in state.items():

fastNLP/core/drivers/utils.py

@@ -1,18 +1,22 @@
 from typing import List
 import subprocess
 
+__all__ = []
 
 def distributed_open_proc(output_from_new_proc:str, command:List[str], env_copy:dict, rank:int=None):
     r"""
     使用 command 通过 subprocess.Popen 开启新的进程。
 
-    :param output_from_new_proc: 可选 ["ignore", "all", "only_error"]，以上三个为特殊关键字，分别表示完全忽略拉起进程的打印输出，
-        only_error 表示只打印错误输出流；all 表示子进程的所有输出都打印。如果不为以上的关键字，则表示一个文件夹，将在该文件夹下建立
-        两个文件，名称分别为 {rank}_std.log, {rank}_err.log 。原有的文件会被直接覆盖。
-    :param command: List[str] 启动的命令
+    :param output_from_new_proc: 可选 ``["ignore", "all", "only_error"]``，以上三个为特殊关键字，分别表示：
+        * ``"ignore:`` -- 完全忽略拉起进程的打印输出；
+        * ``"only_error"`` -- 表示只打印错误输出流；
+        * ``"all"`` -- 子进程的所有输出都打印。
+        * 如果不为以上的关键字，则表示一个文件夹，将在该文件夹下建立两个文件，名称分别为 {rank}_std.log, {rank}_err.log 。
+          原有的文件会被直接覆盖。
+    :param command: 启动的命令
     :param env_copy: 需要注入的环境变量。
     :param rank: global_rank；
-    :return: 返回使用 ``subprocess.Popen`` 打开的进程；
+    :return: 使用 ``subprocess.Popen`` 打开的进程；
     """
     if output_from_new_proc == "all":
         proc = subprocess.Popen(command, env=env_copy)

fastNLP/core/utils/paddle_utils.py

@@ -98,7 +98,7 @@ def get_paddle_gpu_str(device: Union[str, int]) -> str:
         'gpu:1'
 
     :param device: 设备编号或设备名；
-    :return: 返回对应的 ``gpu:x`` 格式的设备名；
+    :return: 对应的 ``gpu:x`` 格式的设备名；
     """
     if isinstance(device, str):
         return device.replace("cuda", "gpu")

fastNLP/core/utils/utils.py

@@ -83,7 +83,7 @@ def auto_param_call(fn: Callable, *args, signature_fn: Optional[Callable] = None
         然后通过该函数签名提取参数值后，再传给 ``fn`` 进行实际的运算；
     :param mapping: 一个字典，用来更改其前面的字典的键值；
 
-    :return: 返回 ``fn`` 运行的结果；
+    :return:  ``fn`` 运行的结果；
     """
 
     if signature_fn is not None:
@@ -233,7 +233,7 @@ def check_user_specific_params(user_params: Dict, fn: Callable, fn_name=None):
         ``value`` 为每一个参数的值；
     :param fn: 将要被调用的函数；
     :param fn_name: 在打印提示信息是如何显示函数名
-    :return: 返回一个字典，其中为在之后调用函数 ``fn`` 时真正会被传进去的参数的值；
+    :return: 一个字典，其中为在之后调用函数 ``fn`` 时真正会被传进去的参数的值；
     """
     if fn_name is None:
         fn_name = fn.__name__
@@ -285,7 +285,7 @@ def match_and_substitute_params(mapping: Optional[Union[Callable, Dict]] = None,
 
     :param mapping: 用于转换的字典或者函数；当 ``mapping`` 是函数时，返回值必须为字典类型；
     :param data: 需要被转换的对象；
-    :return: 返回转换后的结果；
+    :return: 转换后的结果；
     """
     if mapping is None:
         return data

fastNLP/io/loader/matching.py

@@ -164,7 +164,7 @@ class SNLILoader(JsonLoader):
         :param str paths: 传入一个目录, 将在该目录下寻找snli_1.0_train.jsonl, snli_1.0_dev.jsonl
             和snli_1.0_test.jsonl三个文件。
 
-        :return: 返回的 :class:`~fastNLP.io.DataBundle`
+        :return: :class:`~fastNLP.io.DataBundle`
         """
         _paths = {}
         if paths is None:

fastNLP/io/loader/summarization.py

@@ -42,7 +42,7 @@ class ExtCNNDMLoader(JsonLoader):
             test.label.jsonl三个文件(该目录还应该需要有一个名字为vocab的文件，在 :class:`~fastNLP.io.ExtCNNDMPipe`
             当中需要用到)。
 
-        :return: 返回 :class:`~fastNLP.io.DataBundle`
+        :return:  :class:`~fastNLP.io.DataBundle`
         """
         if paths is None:
             paths = self.download()

fastNLP/io/pipe/utils.py

@@ -71,7 +71,7 @@ def get_tokenizer(tokenize_method: str, lang='en'):
 
     :param str tokenize_method: 获取tokenzier方法
     :param str lang: 语言，当前仅支持en
-    :return: 返回tokenize函数
+    :return: tokenize函数
     """
     tokenizer_dict = {
         'spacy': None,

fastNLP/modules/torch/decoder/crf.py

@@ -290,7 +290,7 @@ class ConditionalRandomField(nn.Module):
         :param bool unpad: 是否将结果删去padding。False, 返回的是batch_size x max_len的tensor; True，返回的是
             List[List[int]], 内部的List[int]为每个sequence的label，已经除去pad部分，即每个List[int]的长度是这
             个sample的有效长度。
-        :return: 返回 (paths, scores)。
+        :return:  (paths, scores)。
                     paths: 是解码后的路径, 其值参照unpad参数.
                     scores: torch.FloatTensor, size为(batch_size,), 对应每个最优路径的分数。
 

tests/core/controllers/_test_trainer_fleet.py

@@ -32,7 +32,7 @@ from tests.helpers.callbacks.helper_callbacks import RecordMetricCallback
 
 @dataclass
 class MNISTTrainFleetConfig:
-    num_labels: int = 3
+    num_labels: int = 5
     feature_dimension: int = 5
 
     batch_size: int = 4
@@ -79,7 +79,7 @@ def test_trainer_fleet(
 
         n_epochs=n_epochs,
         callbacks=callbacks,
-        # output_from_new_proc="logs",
+        output_from_new_proc="logs",
     )
     trainer.run()