Add getter APIs for TP/PP/DP ranks in DeepSpeedEngine (#7427)

Thanks again for giving opportunity for improving this Community!
This PR is from Issue #7423.

1) Motivation

To improve compatibility with low-level profiling tools (e.g., NVIDIA
CUPTI or DCGM), it can be useful to expose parallelism-specific rank
(tensor/pipeline/data) at the engine level.

2) Changes

I Added three getter methods to DeepSpeedEngine:
  - get_tensor_parallel_rank()
  - get_pipeline_parallel_rank()
  - get_data_parallel_rank()


Thank you for reviewing this contribution!

---------

Signed-off-by: WoosungMyung <dntjd517@naver.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>

Author: WoosungMyung

deepspeed/runtime/engine.py

@@ -730,6 +730,15 @@ def random_ltd_initialize(self):
             raise ValueError(f'not yet support')
             #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler)
 
+    def get_data_parallel_rank(self):
+        return groups.get_data_parallel_rank()
+
+    def get_tensor_parallel_rank(self):
+        return groups.get_tensor_model_parallel_rank()
+
+    def get_model_parallel_rank(self):
+        return groups.get_model_parallel_rank()
+
     def get_sequence_parallel_group(self):
         return self.seq_parallel_group
 

deepspeed/runtime/pipe/engine.py

@@ -535,6 +535,9 @@ def is_last_stage(self):
         """True if this process is in the last stage in the pipeline."""
         return self.stage_id == self.num_stages - 1
 
+    def get_pipeline_parallel_rank(self):
+        return self.stage_id
+
     def _reduce_outputs(self, outputs, reduce='avg', reduce_dp=True, micro_batches=None):
         if reduce is None:
             return outputs