update bagua

Author: ananthsub

pytorch_lightning/strategies/bagua.py

@@ -12,9 +12,11 @@
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.strategies.strategy import TBroadcast
+from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.distributed import ReduceOp
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _BAGUA_AVAILABLE
+from pytorch_lightning.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.utilities.seed import reset_seed
 
 if _BAGUA_AVAILABLE:
@@ -148,6 +150,35 @@ def _set_node_environment_variables(self) -> None:
         os.environ["WORLD_SIZE"] = str(self.world_size)
         os.environ["LOCAL_RANK"] = str(self.local_rank)
 
+    def setup(self, trainer: "pl.Trainer") -> None:
+        self._rank_0_will_call_children_scripts = self.broadcast(self._rank_0_will_call_children_scripts)
+        if self._should_run_deadlock_detection():
+            self._share_information_to_prevent_deadlock()
+
+        self.accelerator.setup(trainer)
+
+        # move the model to the correct device
+        self.model_to_device()
+
+        if self._layer_sync:
+            self.model = self._layer_sync.apply(self.model)
+
+        # skip wrapping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = trainer.state.fn
+
+        # set up optimizers after the module has been moved to the device
+        # but before the module has been wrapped
+        self.setup_optimizers(trainer)
+        optimizers_to_device(self.optimizers, self.root_device)
+
+        if trainer_fn == TrainerFn.FITTING:
+            self._configure_bagua_model(trainer)
+
+        self.setup_precision_plugin()
+        self._rank_0_will_call_children_scripts = self.broadcast(self._rank_0_will_call_children_scripts)
+        if self._should_run_deadlock_detection():
+            self._share_information_to_prevent_deadlock()
+
     def _check_qadam_optimizer(self) -> None:
         has_qadam_optimizer = any([isinstance(opt, QAdamOptimizer) for opt in self.optimizers])
 
@@ -156,12 +187,12 @@ def _check_qadam_optimizer(self) -> None:
 
         self._bagua_kwargs["q_adam_optimizer"] = self.optimizers[0]
 
-    def configure_ddp(self) -> None:
+    def _configure_bagua_model(self, trainer: "pl.Trainer") -> None:
         model = LightningBaguaModule(self.model)  # type: ignore[arg-type]
         self._model = self._setup_model(model)
 
         # start the background communication for async algorithm
-        if self.lightning_module.trainer.training and self._bagua_algorithm == "async":
+        if trainer.training and self._bagua_algorithm == "async":
             self.model.bagua_algorithm.resume(self.model)  # type: ignore
 
     def _setup_model(self, model: Module) -> BaguaDistributedDataParallel:

tests/strategies/test_bagua_strategy.py

@@ -85,9 +85,9 @@ def test_configuration(algorithm, tmpdir):
     ), mock.patch("bagua.torch_api.communication.is_initialized", return_value=True):
         if algorithm == "qadam":
             with pytest.raises(MisconfigurationException, match="Bagua QAdam can only accept one QAdamOptimizer"):
-                trainer.strategy.configure_ddp()
+                trainer.strategy._configure_bagua_model()
         else:
-            trainer.strategy.configure_ddp()
+            trainer.strategy._configure_bagua_model()
 
 
 @RunIf(bagua=True, min_gpus=1)