support DDP training with no_sync and fused_allreduce_gradients (#332)

* support DDP training with no_sync and fused_allreduce_gradients

* fix

* enhance run_check and sanity check

* remove autocast context manager in visu.py

* fix logger

Author: HydrogenSulfate

ppsci/solver/eval.py

@@ -55,12 +55,13 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
                 for key in solver.eval_time_info:
                     solver.eval_time_info[key].reset()
             reader_cost = time.perf_counter() - reader_tic
-
             for v in input_dict.values():
                 v.stop_gradient = False
 
             # forward
-            with solver.autocast_context_manager(), solver.no_grad_context_manager():
+            with solver.autocast_context_manager(
+                solver.use_amp, solver.amp_level
+            ), solver.no_grad_context_manager(solver.eval_with_no_grad):
                 output_dict, validator_loss = solver.forward_helper.eval_forward(
                     _validator.output_expr,
                     input_dict,
@@ -179,11 +180,13 @@ def _eval_by_batch(solver, epoch_id: int, log_freq: int) -> float:
                     solver.eval_time_info[key].reset()
             reader_cost = time.perf_counter() - reader_tic
             batch_size = next(iter(input_dict.values())).shape[0]
-
             for v in input_dict.values():
                 v.stop_gradient = False
+
             # forward
-            with solver.autocast_context_manager(), solver.no_grad_context_manager():
+            with solver.autocast_context_manager(
+                solver.use_amp, solver.amp_level
+            ), solver.no_grad_context_manager(solver.eval_with_no_grad):
                 output_dict, validator_loss = solver.forward_helper.eval_forward(
                     _validator.output_expr,
                     input_dict,

ppsci/solver/solver.py

@@ -16,6 +16,7 @@
 
 import contextlib
 import copy
+import itertools
 import os
 import sys
 from typing import Any
@@ -203,6 +204,16 @@ def __init__(
 
         # whether calculate metrics after each batch during evaluate
         self.compute_metric_by_batch = compute_metric_by_batch
+        if validator is not None:
+            for metric in itertools.chain(
+                *[_v.metric.values() for _v in self.validator.values()]
+            ):
+                if metric.keep_batch ^ compute_metric_by_batch:
+                    raise ValueError(
+                        f"{misc.typename(metric)}.keep_batch should be "
+                        f"{compute_metric_by_batch} when compute_metric_by_batch="
+                        f"{compute_metric_by_batch}."
+                    )
         # whether set `stop_gradient=True` for every Tensor if no differentiation involved during computation
         self.eval_with_no_grad = eval_with_no_grad
 
@@ -247,6 +258,11 @@ def __init__(
             self.model = fleet.distributed_model(self.model)
             if self.optimizer is not None:
                 self.optimizer = fleet.distributed_optimizer(self.optimizer)
+            logger.warning(
+                f"Detected world_size({self.world_size}) > 1, it is recommended to "
+                "scale up the learning rate and reduce the epochs or "
+                "iters_per_epoch according to the world_size number both linearly."
+            )
 
         self.global_step = 0
 
@@ -493,7 +509,7 @@ def predict(
                 batch_input_dict[key].stop_gradient = False
 
             # forward
-            with self.autocast_context_manager():
+            with self.autocast_context_manager(self.use_amp, self.amp_level):
                 batch_output_dict = self.model(batch_input_dict)
 
             # collect batch data
@@ -522,36 +538,74 @@ def export(self):
         jit.save(static_model, save_path)
         logger.info(f"The inference model has been exported to {export_dir}")
 
-    def autocast_context_manager(self) -> contextlib.AbstractContextManager:
-        """Autocast context manager for Auto Mix Precision.
+    def autocast_context_manager(
+        self, enable: bool, level: Literal["O0", "O1", "O2"] = "O1"
+    ) -> contextlib.AbstractContextManager:
+        """Smart autocast context manager for Auto Mix Precision.
+
+        Args:
+            enable (bool): Enable autocast.
+            level (Literal["O0", "O1", "O2"]): Autocast level.
 
         Returns:
-            Union[contextlib.AbstractContextManager]: Context manager.
+            contextlib.AbstractContextManager: Smart autocast context manager.
         """
-        if self.use_amp:
-            ctx_manager = amp.auto_cast(level=self.amp_level)
+        if enable:
+            ctx_manager = amp.auto_cast(level=level)
         else:
             ctx_manager = (
                 contextlib.nullcontext()
                 if sys.version_info >= (3, 7)
                 else contextlib.suppress()
             )
-
         return ctx_manager
 
-    def no_grad_context_manager(self) -> contextlib.AbstractContextManager:
-        """No grad manager.
+    def no_grad_context_manager(
+        self, enable: bool
+    ) -> contextlib.AbstractContextManager:
+        """Smart no_grad context manager.
+
+        Args:
+            enable (bool): Enable no_grad.
 
         Returns:
-            Union[contextlib.AbstractContextManager]: Context manager.
+            contextlib.AbstractContextManager: Smart no_grad context manager.
         """
-        if self.eval_with_no_grad:
+        if enable:
             ctx_manager = paddle.no_grad()
         else:
             ctx_manager = (
                 contextlib.nullcontext()
                 if sys.version_info >= (3, 7)
                 else contextlib.suppress()
             )
+        return ctx_manager
+
+    def no_sync_context_manager(
+        self,
+        enable: bool,
+        ddp_model: paddle.DataParallel,
+    ) -> contextlib.AbstractContextManager:
+        """Smart no_sync context manager for given model.
+        NOTE: Only `paddle.DataParallel` object has `no_sync` interface.
 
+        Args:
+            enable (bool): Enable no_sync.
+
+        Returns:
+            contextlib.AbstractContextManager: Smart no_sync context manager.
+        """
+        if enable:
+            if not isinstance(ddp_model, paddle.DataParallel):
+                raise TypeError(
+                    "no_sync interface is only for model with type paddle.DataParallel, "
+                    f"but got type {type(ddp_model)}"
+                )
+            ctx_manager = ddp_model.no_sync()
+        else:
+            ctx_manager = (
+                contextlib.nullcontext()
+                if sys.version_info >= (3, 7)
+                else contextlib.suppress()
+            )
         return ctx_manager

ppsci/solver/train.py

@@ -14,13 +14,15 @@
 
 import time
 
+from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu
+
+from ppsci import solver
 from ppsci.solver import printer
-from ppsci.utils import expression
 from ppsci.utils import misc
 from ppsci.utils import profiler
 
 
-def train_epoch_func(solver, epoch_id: int, log_freq: int):
+def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
     """Train program for one epoch
 
     Args:
@@ -61,38 +63,48 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
             total_batch_size += next(iter(input_dict.values())).shape[0]
             reader_tic = time.perf_counter()
 
-        # forward for every constraint, including model and equation expression
-        with solver.autocast_context_manager():
-            constraint_losses = solver.forward_helper.train_forward(
-                [_constraint.output_expr for _constraint in solver.constraint.values()],
-                input_dicts,
-                solver.model,
-                solver.constraint,
-                label_dicts,
-                weight_dicts,
-            )
-
-        # compute loss for each constraint according to its' own output, label and weight
-        for i, _constraint in enumerate(solver.constraint.values()):
-            total_loss += constraint_losses[i]
-            loss_dict[_constraint.name] += float(constraint_losses[i])
-
-        if solver.update_freq > 1:
-            total_loss = total_loss / solver.update_freq
-        loss_dict["loss"] = float(total_loss)
-
-        # backward
-        if solver.use_amp:
-            total_loss_scaled = solver.scaler.scale(total_loss)
-            total_loss_scaled.backward()
-            if iter_id % solver.update_freq == 0:
+        with solver.no_sync_context_manager(solver.world_size > 1, solver.model):
+            # forward for every constraint, including model and equation expression
+            with solver.autocast_context_manager(solver.use_amp, solver.amp_level):
+                constraint_losses = solver.forward_helper.train_forward(
+                    [
+                        _constraint.output_expr
+                        for _constraint in solver.constraint.values()
+                    ],
+                    input_dicts,
+                    solver.model,
+                    solver.constraint,
+                    label_dicts,
+                    weight_dicts,
+                )
+                # accumulate all losses
+                for i, _constraint in enumerate(solver.constraint.values()):
+                    total_loss += constraint_losses[i]
+                    loss_dict[_constraint.name] += (
+                        float(constraint_losses[i]) / solver.update_freq
+                    )
+                if solver.update_freq > 1:
+                    total_loss = total_loss / solver.update_freq
+                loss_dict["loss"] = float(total_loss)
+
+            # backward
+            if solver.use_amp:
+                total_loss_scaled = solver.scaler.scale(total_loss)
+                total_loss_scaled.backward()
+            else:
+                total_loss.backward()
+
+        # update parameters
+        if iter_id % solver.update_freq == 0 or iter_id == solver.iters_per_epoch:
+            if solver.world_size > 1:
+                # fuse + allreduce manually before optimization if use DDP + no_sync
+                # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622
+                hpu.fused_allreduce_gradients(list(solver.model.parameters()), None)
+            if solver.use_amp:
                 solver.scaler.minimize(solver.optimizer, total_loss_scaled)
-                solver.optimizer.clear_grad()
-        else:
-            total_loss.backward()
-            if iter_id % solver.update_freq == 0:
+            else:
                 solver.optimizer.step()
-                solver.optimizer.clear_grad()
+            solver.optimizer.clear_grad()
 
         # update learning rate by step
         if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch:
@@ -111,7 +123,7 @@ def train_epoch_func(solver, epoch_id: int, log_freq: int):
         batch_tic = time.perf_counter()
 
 
-def train_LBFGS_epoch_func(solver, epoch_id: int, log_freq: int):
+def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
     """Train function for one epoch with L-BFGS optimizer.
 
     Args:
@@ -152,30 +164,38 @@ def closure():
                 Tensor: Computed loss.
             """
             total_loss = 0
-            for i, _constraint in enumerate(solver.constraint.values()):
-                evaluator = expression.ExpressionSolver(
-                    _constraint.input_keys, _constraint.output_keys, solver.model
-                )
-                for output_name, output_formula in _constraint.output_expr.items():
-                    if output_name in label_dict:
-                        evaluator.add_target_expr(output_formula, output_name)
-
-                # forward for every batched data dict
-                with solver.autocast_context_manager():
-                    output_dict_i = evaluator(input_dicts[i])
-                    constraint_loss = _constraint.loss(
-                        output_dict_i, label_dicts[i], weight_dicts[i]
+            with solver.no_sync_context_manager(solver.world_size > 1, solver.model):
+                with solver.autocast_context_manager(solver.use_amp, solver.amp_level):
+                    # forward for every constraint, including model and equation expression
+                    constraint_losses = solver.forward_helper.train_forward(
+                        [
+                            _constraint.output_expr
+                            for _constraint in solver.constraint.values()
+                        ],
+                        input_dicts,
+                        solver.model,
+                        solver.constraint,
+                        label_dicts,
+                        weight_dicts,
                     )
-                    total_loss += constraint_loss
+                    # accumulate all losses
+                    for i, _constraint in enumerate(solver.constraint.values()):
+                        total_loss += constraint_losses[i]
+                        loss_dict[_constraint.name] = float(constraint_losses[i])
+                    loss_dict["loss"] = float(total_loss)
 
-                loss_dict[_constraint.name] = float(constraint_loss)
+                # backward
+                solver.optimizer.clear_grad()
+                total_loss.backward()
 
-            solver.optimizer.clear_grad()
-            total_loss.backward()
-            loss_dict["loss"] = float(total_loss)
+            if solver.world_size > 1:
+                # fuse + allreduce manually before optimization if use DDP model
+                # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622
+                hpu.fused_allreduce_gradients(list(solver.model.parameters()), None)
 
             return total_loss
 
+        # update parameters
         solver.optimizer.step(closure)
 
         # update learning rate by step

ppsci/solver/visu.py

@@ -52,7 +52,7 @@ def visualize_func(solver, epoch_id: int):
                 batch_input_dict[key].stop_gradient = False
 
             # forward
-            with solver.no_grad_context_manager():
+            with solver.no_grad_context_manager(solver.eval_with_no_grad):
                 batch_output_dict = solver.forward_helper.visu_forward(
                     _visualizer.output_expr, batch_input_dict, solver.model
                 )

ppsci/utils/save_load.py

@@ -127,6 +127,11 @@ def save_checkpoint(
     """
     if paddle.distributed.get_rank() != 0:
         return
+    if model_dir is None:
+        logger.warning(
+            f"model_dir({model_dir}) is set to None, skip save_checkpoint..."
+        )
+        return
     model_dir = os.path.join(model_dir, "checkpoints")
     os.makedirs(model_dir, exist_ok=True)
     model_path = os.path.join(model_dir, prefix)