feat: print expert groups on megatron init (#13874)

Signed-off-by: Alexander Zhipa <[email protected]>
Co-authored-by: Alexander Zhipa <[email protected]>

Author: clumsy

nemo/lightning/megatron_init.py

@@ -499,16 +499,24 @@ def generator_wrapper(group_type, is_expert=False, **kwargs):
     # EP rank
     expert_model_parallel_rank = 0
     if expert_model_parallel_size_ is not None and expert_model_parallel_size_ > 1:
+        all_expert_model_parallel_ranks = []
         for ranks in generator_wrapper('ep', is_expert=True):
+            all_expert_model_parallel_ranks.append(ranks)
             if rank in ranks:
                 expert_model_parallel_rank = list(ranks).index(rank)
+        logging.info(f'All expert model parallel group ranks: {all_expert_model_parallel_ranks}')
+        logging.info(f'Rank {rank} has expert model parallel rank: {expert_model_parallel_rank}')
 
     # ETP
     expert_tensor_parallel_rank = 0
     if expert_tensor_parallel_size_ is not None and expert_tensor_parallel_size_ > 1:
+        all_expert_tensor_parallel_ranks = []
         for ranks in generator_wrapper('tp', is_expert=True):
+            all_expert_tensor_parallel_ranks.append(ranks)
             if rank in ranks:
                 expert_tensor_parallel_rank = list(ranks).index(rank)
+        logging.info(f'All expert tensor parallel group ranks: {all_expert_tensor_parallel_ranks}')
+        logging.info(f'Rank {rank} has expert tensor parallel rank: {expert_tensor_parallel_rank}')
 
     # Build the pipeline model-parallel groups and embedding groups
     # (first and last rank in each pipeline model-parallel group).