[Auto Parallel] Add doc for DistModel (#60002)

pkuzyc · web-flow · commit 22b617bca79c · 2023-12-18T19:42:10.000+08:00
* modify the default arg 'stop_gradient' in shard_tensor api

* 1. add doc for DistModel; 2. modify the default setting of 'stop_gradient' in shard_tensor

* bug fix

* fix the codestyle and typos of the docstring

* update docs; test=docs_preciew

* update docs

* polish as comments
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -86,6 +86,7 @@
     shard_optimizer,
     to_static,
     Strategy,
+    DistModel,
 )
 
 from .fleet import BoxPSDataset  # noqa: F401
@@ -167,4 +168,5 @@
     "shard_optimizer",
     "to_static",
     "Strategy",
+    "DistModel",
 ]
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -113,29 +113,40 @@ def sharding_specs(self):
 
 class DistModel:
     """
-    DistModel is generated by ``paddle.distributed.to_static``. It contains the
-    static graph converted from a ``paddle.nn.layer`` whose parameters are
-    distributed tensors (constructed from ``paddle.distributed.shard_tensor``),
-    and provides the APIs for training, evaluation and prediction with the
-    static graph.
-
-    Please first set the DistModel to "train", "eval" or "predict" mode and
-    then use the __call__ method for training, evaluation and prediction
-    respectively.
-
-    In "train" mode, executing ``__call__`` will update the parameters
-    of the model and return the loss. In "eval" mode, executing ``__call__``
-    will return the loss. In "predict" mode, executing ``__call__`` returns a
-    dict that contains the outputs of the model, where the value of "out0" is
-    the first output.
-
-    If loss and optimizer are both given, DistModel will be set to "train" mode
-    in default. If loss is given but optimizer is None, DistModel will be set
-    to "eval" mode in default. If loss and optimizer are both None, DistModel
-    will be set to "predict" mode in default.
+    `DistModel` is the model converted from a ``paddle.nn.layer`` with distributed
+    tensors as its parameters. It contains the static graph converted from a
+    ``paddle.nn.layer`` whose parameters are distributed tensors (constructed
+    from ``paddle.distributed.shard_tensor``), and provides the APIs for training,
+    evaluation and prediction with the static graph.
+
+    It is suggested to generate DistModel by ``paddle.distributed.to_static``,
+    not directly by ``paddle.distributed.DistModel``.
+
+    Please first set the DistModel to "train", "eval" or "predict" mode with
+    ``train()/eval()/predict()`` method and then use the ``__call__`` method for
+    training, evaluation and prediction respectively.
 
     For more details of the usage, please refer to the sample code in
     ``paddle.distributed.to_static``.
+
+    Args:
+        layer(paddle.nn.Layer): The layer in dygraph mode, whose parameters
+            are distributed tensors generated by ``shard_tensor``.
+        loader(paddle.io.DataLoader): The data loader used in dygraph mode,
+            used to generate DistributedDataloader for training with the
+            static graph.
+        loss(Loss|Callable|None, optional): The loss function for training
+            or evaluating the model. Can be a `paddle.nn.Layer` instance or
+            any callable function. If loss is not None, DistModel will be set
+            to "train" (when the optimizer is also not None) or "eval" mode
+            (when optimizer is None) in default. If it is None, DistModel will
+            be set to "predict" mode in default. Default: None.
+        optimizer(paddle.optimizer.Optimizer|None, optional): The optimizer
+            for training. If both optimizer and loss are set, DistModel will
+            be set to "train" mode in default. Default: None.
+        strategy(paddle.distributed.Strategy|None, optional): Configs for
+            parallel strategies and optimization settings (e.g. sharding,
+            pipeline parallelism). Default: None.
     """
 
     def __init__(
@@ -192,27 +203,34 @@ def __init__(
 
     def train(self):
         """
-        Set the mode of DistModel to "train".
+        Set the DistModel to "train" mode. In "train" mode,
+        executing ``__call__`` method will update the
+        parameters of the model and return the loss.
         """
         if not self._engine._has_prepared["train"]:
-            raise RuntimeError("The model for training has not been prepared.")
+            raise RuntimeError(
+                "The model for training has not been prepared, please set 'loss' and 'optimizer' when constructing DistModel."
+            )
         self._mode = "train"
         self._engine.to_mode("train")
 
     def eval(self):
         """
-        Set the mode of DistModel to "eval".
+        Set the mode of DistModel to "eval". In "eval" mode,
+        executing ``__call__`` will return the loss.
         """
         if not self._engine._has_prepared["eval"]:
             raise RuntimeError(
-                "The model for evaluation has not been prepared."
+                "The model for evaluation has not been prepared, please set 'loss' when constructing DistModel."
             )
         self._mode = "eval"
         self._engine.to_mode("eval")
 
     def predict(self):
         """
-        Set the mode of DistModel to "predict".
+        Set the mode of DistModel to "predict". In "predict" mode,
+        executing ``__call__`` returns a dict that contains the
+        outputs of the model.
         """
         if not self._engine._has_prepared["predict"]:
             raise RuntimeError(
@@ -234,32 +252,76 @@ def __validate_mode(self, mode):
 
     def dist_main_program(self, mode=None):
         """
-        Get the distributed main program of ``mode``. The ``mode``
-        arg can be "train", "eval" or "predict", if it is not set,
-        ``self._mode`` will be used.
+        Get the distributed main program of the specified ``mode``. Each
+        'mode' has its own distributed main program, ``dist_main_program``
+        returns the corresponding distributed main program of ``mode``.
+
+        Args:
+            mode (str|None, optional): Can be 'train' , 'eval' , 'predict' or None.
+                'train' : Return the distributed main program for training.
+                'eval' : Return the distributed main program for evaluation.
+                'predict' : Return the distributed main program for prediction.
+                None : The current mode of the DistModel will be used.
+                Default : None.
+
+        Returns:
+            The distributed main program of ``mode``.
         """
         mode = self.__validate_mode(mode)
         return self._engine.get_dist_main_program(mode)
 
     def dist_startup_program(self, mode=None):
         """
-        Get the distributed startup program of ``mode``. The ``mode``
-        arg can be "train", "eval" or "predict", if it is not set,
-        ``self._mode`` will be used.
+        Get the corresponding distributed startup program of ``mode``,
+        which is used for initializing the parameters.
+
+        Args:
+            mode (str|None, optional): Can be 'train' , 'eval' , 'predict' or None.
+                'train' : Return the distributed startup program for training.
+                'eval' : Return the distributed startup program for evaluation.
+                'predict' : Return the distributed startup program for prediction.
+                None: The current mode of the DistModel will be used.
+                Default : None.
+
+        Returns:
+            The distributed startup program of ``mode``.
         """
         mode = self.__validate_mode(mode)
         return self._engine.get_dist_startup_program(mode)
 
     def serial_main_program(self, mode=None):
         """
-        Get the serial main program of ``mode``.
+        Get the corresponding serial main program of ``mode``, containing
+        the whole variables and operators of the given ``layer``.
+
+        Args:
+            mode (str|None, optional): Can be 'train', 'eval', 'predict' or None.
+                'train' : Return the main program for training.
+                'eval' : Return the main program for evaluation.
+                'predict' : Return the main program for prediction.
+                None : The current mode of the DistModel will be used.
+                Default : None.
+
+        Returns:
+            The serial main program of ``mode``.
         """
         mode = self.__validate_mode(mode)
         return self._engine.get_serial_main_program(mode)
 
     def serial_startup_program(self, mode=None):
         """
-        Get the serial startup program of ``mode``.
+        Get the corresponding serial startup program of ``mode``.
+
+        Args:
+            mode (str|None, optional): Can be 'train' , 'eval' , 'predict' or None.
+                'train' : Return the serial startup program for training.
+                'eval' : Return the serial startup program for evaluation.
+                'predict' : Return the serial startup program for prediction.
+                None : The current mode of the DistModel will be used.
+                Default : None.
+
+        Returns:
+            The serial startup program of ``mode``.
         """
         mode = self.__validate_mode(mode)
         return self._engine.get_serial_startup_program(mode)
@@ -444,11 +506,11 @@ def __init__(self, config_dict=None):
 class Strategy(auto_strategy.BaseConfig):
     """
     The `Strategy` object is used to configure the parallelization
-    and optimization strategies for static graph. Currently contains
+    and optimization strategies for static graph. Currently supports
     configuring ``sharding``, ``fused_passes``, ``gradient_merge``
-    and ``pipline``. More strategies will be supported in the future.
+    and ``pipeline``. More strategies will be supported in the future.
 
-    ``sharding`` is used to cnofigure the sharding states of the optimizer,
+    ``sharding`` is used to configure the sharding states of the optimizer,
     for saving the GPU memory.
 
     ``fused_passes`` is used to configure the fusion of the computation in
@@ -460,10 +522,10 @@ class Strategy(auto_strategy.BaseConfig):
     ``pipeline`` is used to configure the pipeline parallelism strategy.
 
     Args:
-        config (dict|None, optional): If ``config`` is None, the default
-        configurations will be set. If it is a dict, the itmes inside
-        the dict will be used to set the configurations, the others remain
-        the default values.
+        config(dict|None, optional): The user-defined configurations.
+            If ``config`` is None, use default configurations. If it is
+            a dict, the items inside the dict will be used to set the
+            configurations, and the others remain the default values.
 
     Examples:
         .. code-block:: python
@@ -528,7 +590,7 @@ def __init__(self, config=None):
     @property
     def sharding(self):
         """
-        ``sharding`` is used to cnofigure the sharding states of the optimizer,
+        ``sharding`` is used to configure the sharding states of the optimizer,
         containing following configs:
 
             ``enable`` (bool): whether to enable sharding. Default: False.
@@ -609,7 +671,7 @@ def fused_passes(self):
     @property
     def pipeline(self):
         """
-        ``pipeline`` is used to configure the pipeline parallelism in training,
+        ``pipeline`` is used to configure the pipeline parallelism,
         containing following configs:
 
             ``enable`` (bool): whether to enable pipeline parallelism. Default: False.
@@ -647,7 +709,7 @@ def to_static(
 ):
     """
     Converts the ``layer`` with distributed tensor (constructed from
-    ``paddle.distributed.shard_tensor``) to a static graph. to_static
+    ``paddle.distributed.shard_tensor``) to a static graph. ``to_static``
     returns a DistModel instance containing the static graph for
     distributed training, evaluation and prediction, and an object of
     DistributedDataLoader to generate data.
@@ -667,11 +729,8 @@ def to_static(
             pipeline parallelism). Default: None.
 
     Returns:
-        DistModel: A DistModel tha contains corresponding computational graph
-            for the input ``layer`` and provides APIs for training, evaluation
-            and prediction.
-        DistributedDataLoader: An optimized data loader that can be used
-            to generate data.
+        DistModel: A ``DistModel`` instance converted the input ``layer``.
+        DistributedDataLoader: An optimized data loader that can be used to generate data, it's usage is the same as ``paddle.io.DataLoader`` .
 
     Examples:
         .. code-block:: python
@@ -762,7 +821,7 @@ def to_static(
             ...     # where the value of "out0" is the first output.
             ...     outs = dist_model(image)
 
-            >>> # This case need to be excuted in multi-card environment
+            >>> # This case need to be executed in multi-card environment
             >>> # export CUDA_VISIBLE_DEVICES=0,1
             >>> # python -m paddle.distributed.launch {test_case}.py
     """
@@ -773,7 +832,7 @@ def to_static(
 
 
 def shard_tensor(
-    data, mesh, placements, dtype=None, place=None, stop_gradient=True
+    data, mesh, placements, dtype=None, place=None, stop_gradient=None
 ):
     """
     Constructs a ``paddle.Tensor`` with distributed attributes from ``data``,
@@ -794,7 +853,10 @@ def shard_tensor(
         place(CPUPlace|CUDAPinnedPlace|CUDAPlace|str, optional): The place to allocate Tensor. Can be
             CPUPlace, CUDAPinnedPlace, CUDAPlace. Default: None, means global place. If ``place`` is
             string, It can be ``cpu``, ``gpu:x`` and ``gpu_pinned``, where ``x`` is the index of the GPUs.
-        stop_gradient(bool, optional): Whether to block the gradient propagation of Autograd. Default: True.
+        stop_gradient(bool, optional): Whether to block the gradient propagation of Autograd. If
+            ``stop_gradient`` is None, set the returned Tensor's ``stop_gradient`` identical as the
+            ``data.stop_gradient`` when ``data`` has ``stop_gradient`` attribute and True otherwise.
+            Default: None.
 
     Returns:
         Tensor: A Tensor constructed from ``data`` with distributed attributes.
@@ -824,6 +886,8 @@ def shard_tensor(
 
     # 1. create dense tensor
     # `paddle.to_tensor` supports both dynamic and static mode
+    if stop_gradient is None:
+        stop_gradient = getattr(data, "stop_gradient", True)
     tensor = paddle.to_tensor(
         data, dtype=dtype, place=place, stop_gradient=stop_gradient
     )

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,7 @@`
`86`	`86`	`shard_optimizer,`
`87`	`87`	`to_static,`
`88`	`88`	`Strategy,`
	`89`	`+ DistModel,`
`89`	`90`	`)`
`90`	`91`
`91`	`92`	`from .fleet import BoxPSDataset # noqa: F401`
`@@ -167,4 +168,5 @@`
`167`	`168`	`"shard_optimizer",`
`168`	`169`	`"to_static",`
`169`	`170`	`"Strategy",`
	`171`	`+ "DistModel",`
`170`	`172`	`]`