修改 Trainer torch_kwargs paddle_kwargs fairscale_kwargs 的描述

fastNLP/core/controllers/trainer.py

@@ -286,20 +286,11 @@ class Trainer(TrainerEventTrigger):
             第一个 ``Trainer`` 实例，即使该 ``Trainer`` 实例的 marker 不为 None；这一点详见 :meth:`~fastNLP.core.controllers.Trainer.on`
 
     :kwargs:
-        * *torch_kwargs* -- 用于在指定 ``driver`` 为 'torch' 时设定具体 driver 实例的一些参数：
-            
-            * 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`` 的参数。
-        * *paddle_kwargs* -- 用于在指定 ``driver`` 为 'paddle' 时设定具体 driver 实例的一些参数：
-
-            * fleet_kwargs -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括：
-
-                * is_collective -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况；
-                * role_maker -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``
-                * 其它用于初始化 ``DataParallel`` 的参数；
+        * *torch_kwargs* -- ``TorchDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.TorchSingleDriver` 和
+          :class:`~fastNLP.core.drivers.torch_driver.TorchDDPDriver`；
+        * *paddle_kwargs* -- ``PaddleDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver` 和
+          :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver`；
+        * *fairscale_kwargs* -- ``FairScaleDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.FairScaleDriver`；
         * *data_device* -- 一个具体的 driver 实例中，有 ``model_device`` 和 ``data_device``，前者表示模型所在的设备，后者表示
          当 ``model_device`` 为 None 时应当将数据迁移到哪个设备；
 
@@ -313,23 +304,23 @@ class Trainer(TrainerEventTrigger):
                 3. 对于 paddle，该参数无效；
 
         * *use_dist_sampler* -- 表示是否使用分布式的 ``sampler``。在多卡时，分布式 ``sampler`` 将自动决定每张卡上读取的 sample ，使得一个 epoch
-         内所有卡的 sample 加起来为一整个数据集的 sample，同时为了保证所有卡上拥有相同数量的 sample ，有的卡上可能会有重复的 sample ，例如
-         8卡训练，只有9个sample，如果batch_size为1，那么第二个batch时，有7张卡将没有 sample 可用，因此只有重复使用 sample 来 pad 到第二个
-         batch 中。如果不希望 fastNLP 对 dataloader 的sampler 做特殊设置，请将该值设置为 False ，若确实需要分布式的训练，请在 Trainer 外
-         对 train_dataloader 做的数据做特殊处理使得其在不同的卡之间 sample 是
+          内所有卡的 sample 加起来为一整个数据集的 sample，同时为了保证所有卡上拥有相同数量的 sample ，有的卡上可能会有重复的 sample ，例如
+          8卡训练，只有9个sample，如果batch_size为1，那么第二个batch时，有7张卡将没有 sample 可用，因此只有重复使用 sample 来 pad 到第二个
+          batch 中。如果不希望 fastNLP 对 dataloader 的sampler 做特殊设置，请将该值设置为 False ，若确实需要分布式的训练，请在 Trainer 外
+          对 train_dataloader 做的数据做特殊处理使得其在不同的卡之间 sample 是
         * *evaluate_use_dist_sampler* -- 表示在 ``Evaluator`` 中在使用分布式的时候是否将保证 dataloader 的 ``sampler`` 替换为
-         evaluate 时使用的分布式的 ``sampler``，其特点是每个卡上的数据之间不重叠，所有卡上数据的加起来是整个数据集。若传入的 dataloader
-         的 sampler 为 (a) 深度学习框架自带的默认 sampler ; (b) fastNLP 的 Sampler 等，则将替换为
+          evaluate 时使用的分布式的 ``sampler``，其特点是每个卡上的数据之间不重叠，所有卡上数据的加起来是整个数据集。若传入的 dataloader
+          的 sampler 为 (a) 深度学习框架自带的默认 sampler ; (b) fastNLP 的 Sampler 等，则将替换为
           :class:`~fastNLP.UnrepeatedSequentialSampler`，如果这个行为不是期待的，请本参数设置为 ``False``，并针对每个卡控制其可以
           用到的数据。
         * *output_from_new_proc* -- 应当为一个字符串，表示在多进程的 driver 中其它进程的输出流应当被做如何处理；其值应当为以下之一：
-         ["all", "ignore", "only_error"]；当该参数的值不是以上值时，该值应当表示一个文件夹的名字，我们会将其他 rank 的输出流重定向到
-         log 文件中，然后将 log 文件保存在通过该参数值设定的文件夹中；默认为 "only_error"；
+          ["all", "ignore", "only_error"]；当该参数的值不是以上值时，该值应当表示一个文件夹的名字，我们会将其他 rank 的输出流重定向到
+          log 文件中，然后将 log 文件保存在通过该参数值设定的文件夹中；默认为 "only_error"；
 
             注意该参数仅当使用分布式的 ``driver`` 时才有效，例如 ``TorchDDPDriver``；
         * *progress_bar* -- 以哪种方式显示 progress ，目前支持[None, 'raw', 'rich', 'auto', 'tqdm'] 或者 :class:`~fastNLP.RichCallback`, :class:`~fastNLP.RawTextCallback`等对象，
-         默认为 auto , auto 表示如果检测到当前 terminal 为交互型则使用 :class:`~fastNLP.RichCallback`，否则使用 :class:`~fastNLP.RawTextCallback` 对象。如果
-         需要定制 progress bar 的参数，例如打印频率等，可以传入 :class:`~fastNLP.RichCallback`, :class:`~fastNLP.RawTextCallback` 等对象。
+          默认为 auto , auto 表示如果检测到当前 terminal 为交互型则使用 :class:`~fastNLP.RichCallback`，否则使用 :class:`~fastNLP.RawTextCallback` 对象。如果
+          需要定制 progress bar 的参数，例如打印频率等，可以传入 :class:`~fastNLP.RichCallback`, :class:`~fastNLP.RawTextCallback` 等对象。
         * *train_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Trainer`` 中。与 input_mapping 互斥。
         * *train_output_mapping* -- 与 output_mapping 一致，但是只用于 ``Trainer`` 中。与 output_mapping 互斥。
         * *evaluate_input_mapping* -- 与 input_mapping 一致，但是只用于 ``Evaluator`` 中。与 input_mapping 互斥。

fastNLP/core/drivers/paddle_driver/fleet.py

@@ -130,15 +130,15 @@ class PaddleFleetDriver(PaddleDriver):
     :param is_pull_by_paddle_run: 标记当前进程是否为通过 ``python -m paddle.distributed.launch`` 启动的。
         这个参数仅在 :class:`~fastNLP.core.Trainer` 中初始化 driver 时使用
     :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` 的参数;
+    
     :kwargs:
-        * *paddle_kwargs* -- 用于在指定 ``driver`` 为 'paddle' 时设定具体 driver 实例的一些参数：
-
-            * fleet_kwargs -- 用于在使用 ``PaddleFleetDriver`` 时指定 ``DataParallel`` 和 ``fleet`` 初始化时的参数，包括：
-
-                * is_collective -- 是否使用 paddle 集群式的分布式训练方法，目前仅支持为 ``True`` 的情况；
-                * role_maker -- 初始化 ``fleet`` 分布式训练 API 时使用的 ``RoleMaker``
-                * 其它用于初始化 ``DataParallel`` 的参数；
-
         * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
 
         .. note::
@@ -152,6 +152,7 @@ class PaddleFleetDriver(PaddleDriver):
             parallel_device: Optional[Union[List[str], str]],
             is_pull_by_paddle_run: bool = False,
             fp16: bool = False,
+            paddle_kwrags: Dict = {},
             **kwargs
     ):
         if USER_CUDA_VISIBLE_DEVICES not in os.environ:

fastNLP/core/drivers/paddle_driver/paddle_driver.py

@@ -56,8 +56,18 @@ class PaddleDriver(Driver):
     1. :class:`~fastNLP.core.drivers.PaddleSingleDriver`：实现了使用单卡和 ``cpu`` 训练的具体功能；
     2. :class:`~fastNLP.core.drivers.PaddleFleetDriver`：实现了使用 ``fleet`` 分布式训练 API 进行集群式分布式训练的具体功能；
 
+    .. warning::
+
+        您不应当直接初始化该类，然后传入给 ``Trainer``，换句话说，您应当使用该类的子类 ``PaddleSingleDriver`` 和 ``PaddleDDPDriver``，而不是
+        该类本身；
+
+    .. note::
+
+        您可以在使用 ``PaddleSingleDriver`` 和 ``PaddleFleetDriver`` 时使用 ``PaddleDriver`` 提供的接口；
+
     :param model: 训练时使用的 **PaddlePaddle** 模型；
     :param fp16: 是否开启混合精度训练；
+    :param paddle_kwargs:
     :kwargs:
         * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
 
@@ -66,7 +76,7 @@ class PaddleDriver(Driver):
             关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
 
     """
-    def __init__(self, model: "paddle.nn.Layer", fp16: Optional[bool] = False, **kwargs):
+    def __init__(self, model: "paddle.nn.Layer", fp16: Optional[bool] = False, paddle_kwrags: Dict = {}, **kwargs):
         if not isinstance(model, paddle.nn.Layer):
             raise ValueError(f"Parameter `model` can not be `{type(model)}` in `PaddleDriver`, it should be exactly "
                             f"`paddle.nn.Layer` type.")

fastNLP/core/drivers/paddle_driver/single_device.py

@@ -43,6 +43,8 @@ class PaddleSingleDriver(PaddleDriver):
     :param model: 训练时使用的 **PaddlePaddle** 模型；
     :param device: 训练使用的设备；
     :param fp16: 是否开启混合精度训练；
+    :param paddle_kwargs:
+        * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`paddle.amp.GradScaler` 的参数;
     :kwargs:
         * wo_auto_param_call (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为；
 
@@ -51,7 +53,7 @@ class PaddleSingleDriver(PaddleDriver):
             关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
 
     """
-    def __init__(self, model: "paddle.nn.Layer", device: Union[str, int], fp16: Optional[bool] = False, **kwargs):
+    def __init__(self, model: "paddle.nn.Layer", device: Union[str, int], fp16: Optional[bool] = False, paddle_kwrags: Dict = {}, **kwargs):
         if isinstance(model, DataParallel):
             raise ValueError("`paddle.DataParallel` is not supported in `PaddleSingleDriver`")
 
@@ -61,7 +63,7 @@ class PaddleSingleDriver(PaddleDriver):
             logger.info("You have set `CUDA_VISIBLE_DEVICES` to '' in system environment variable, and we are gonna to"
                         "use `cpu` instead of `gpu` device.")
 
-        super(PaddleSingleDriver, self).__init__(model, fp16=fp16, **kwargs)
+        super(PaddleSingleDriver, self).__init__(model, fp16=fp16, paddle_kwrags=paddle_kwrags, **kwargs)
 
         if device is None:
             raise ValueError("Parameter `device` can not be None in `PaddleSingleDriver`.")

fastNLP/core/drivers/torch_driver/ddp.py

@@ -235,7 +235,12 @@ class TorchDDPDriver(TorchDriver):
     :param parallel_device: 用于分布式训练的 ``gpu`` 设备；
     :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的；
     :param fp16: 是否开启 fp16 训练；
-    :param kwargs: 其余的一些用于设定 ddp 训练的参数；
+    :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`` 的参数;
     """
 
     def __init__(
@@ -244,11 +249,12 @@ class TorchDDPDriver(TorchDriver):
             parallel_device: Optional[Union[List["torch.device"], "torch.device"]],
             is_pull_by_torch_run: bool = False,
             fp16: bool = False,
+            torch_kwargs: Dict = {},
             **kwargs
     ):
 
         # 在加入很多东西后，需要注意这里调用 super 函数的位置；
-        super(TorchDDPDriver, self).__init__(model, fp16=fp16, **kwargs)
+        super(TorchDDPDriver, self).__init__(model, fp16=fp16, torch_kwargs=torch_kwargs, **kwargs)
 
         if isinstance(model, torch.nn.DataParallel):
             raise ValueError(f"Parameter `model` can not be `DataParallel` in `TorchDDPDriver`, it should be "

fastNLP/core/drivers/torch_driver/single_device.py

@@ -35,9 +35,13 @@ class TorchSingleDriver(TorchDriver):
     :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`` 的参数;
     """
 
-    def __init__(self, model, device: "torch.device", fp16: bool = False, **kwargs):
+    def __init__(self, model, device: "torch.device", fp16: bool = False, torch_kwargs: Dict = {}, **kwargs):
         if isinstance(model, DistributedDataParallel):
             raise ValueError("`DistributedDataParallel` is not supported in `TorchSingleDriver`")
 
@@ -47,7 +51,7 @@ class TorchSingleDriver(TorchDriver):
             logger.info("You have set `CUDA_VISIBLE_DEVICES` to '' in system environment variable, and we are gonna to"
                         "use `cpu` instead of `gpu` device.")
 
-        super(TorchSingleDriver, self).__init__(model, fp16=fp16, **kwargs)
+        super(TorchSingleDriver, self).__init__(model, fp16=fp16, torch_kwargs=torch_kwargs, **kwargs)
 
         if device is None:
             logger.debug("device is not set, fastNLP will try to automatically get it.")

fastNLP/core/drivers/torch_driver/torch_driver.py

@@ -47,12 +47,15 @@ class TorchDriver(Driver):
 
         您可以在使用 ``TorchSingleDriver`` 和 ``TorchDDPDriver`` 时使用 ``TorchDriver`` 提供的接口；
 
+    :param model: 训练时使用的 **pytorch** 模型；
+    :param fp16: 是否开启混合精度训练;
+    :param torch_kwargs:
     """
-    def __init__(self, model, fp16: Optional[bool] = False, **kwargs):
+    def __init__(self, model, fp16: Optional[bool] = False, torch_kwargs: Dict = {}, **kwargs):
         super(TorchDriver, self).__init__(model)
 
         """ 进行 fp16 的设置 """
-        self._torch_kwargs = kwargs.get("torch_kwargs", {})
+        self._torch_kwargs = torch_kwargs
 
         # 因为 ddp 和 single_device 的混合精度训练的设置是一样的，因此可以统一抽象到这里；
         self.fp16 = fp16