Add world-size getter in Engine

Signed-off-by: WoosungMyung <dntjd517@naver.com>

Author: WoosungMyung

deepspeed/runtime/engine.py

@@ -748,6 +748,10 @@ def get_tensor_parallel_rank(self):
     def get_model_parallel_rank(self):
         return groups.get_model_parallel_rank()
 
+    def get_parallel_world_sizes(self):
+        """Return a dict of parallel world sizes for data/tensor parallelism."""
+        return {"dp": groups.get_data_parallel_world_size(), "tp": groups.get_tensor_model_parallel_world_size()}
+
     def get_sequence_parallel_group(self):
         return self.seq_parallel_group
 

deepspeed/runtime/pipe/engine.py

@@ -537,6 +537,12 @@ def is_last_stage(self):
     def get_pipeline_parallel_rank(self):
         return self.stage_id
 
+    def get_parallel_world_sizes(self):
+        """Return a dict of parallel world sizes for data/tensor/pipeline parallelism."""
+        sizes = super().get_parallel_world_sizes()
+        sizes["pp"] = self.num_stages
+        return sizes
+
     def _reduce_outputs(self, outputs, reduce='avg', reduce_dp=True, micro_batches=None):
         if reduce is None:
             return outputs