Implementation changes for get_rank(group) for the 'mpi' and 'nccl' backends

[RadhaGulhane13]

Problem :

mcr-dl's support for get_rank(group) doesn't support returning rank with given group, instead it will return world rank.

Description :

• Take a look at mcr-dl's get_rank(group) implementation for :

1. mpi backend
   

   MCR-DL/mcr_dl/mpi.py

   Line 82 in 870751f

   def get_rank(self, group=None):

   
   

   MCR-DL/mcr_dl/ops/csrc/comm/mpi.cpp

   Line 76 in 870751f

   int get_rank(int group = 0)

2. nccl backend
   

   MCR-DL/mcr_dl/nccl.py

   Line 85 in 870751f

   def get_rank(self, group=None):

   
   

   MCR-DL/mcr_dl/ops/csrc/comm/nccl.cpp

   Line 79 in 870751f

   int get_rank(int group = 0)

It take's argument as group but doesn't consider it while retrieving rank w.r.t. given group, instead returns world rank.Implementation support is required to return rank for given group.

• Diving deeper,
  This requirement is necessary to support MCR-DL with Megatron-LM as Megatron-LM also relies on such calls. (refer : https://github.com/OSU-Nowlab/Megatron-LM/blob/5f9c870f9f24b482509699d206a9dbb00958f6fc/megatron/model/transformer.py#L1563)

Testing :

• Add new benchmarks at https://github.com/OSU-Nowlab/MCR-DL/tree/main/tests to test these changes
• Test changes for 'V23.05 Megatron-LM with mcr-dl support Megatron-LM#1' with backend='mpi'/'nccl' and distributed_engine='mcr-dl'