clean up distributed training automatically

betty/engine.py

@@ -118,6 +118,8 @@ def run(self):
                         self.logger.info("Early stopping is executed!")
                         break
 
+        self.cleanup()
+
     def initialize(self):
         """
         Initialize dependencies (computational graph) between problems.
@@ -339,6 +341,15 @@ def is_rank_zero(self):
         """
         return self._rank == 0
 
+    def cleanup(self):
+        """
+        Clean up distributed training
+        """
+        if self._strategy in ["distributed", "zero", "fsdp"]:
+            dist.destroy_process_group()
+        if self.is_rank_zero():
+            self.logger.info("Multilevel optimization finished!")
+
     def is_implemented(self, fn_name):
         """
         Check whether ``fn_name`` method is implemented in the class.