From 23befa1f6a0f15a7023a91786dd8f6704b2f668d Mon Sep 17 00:00:00 2001 From: Guyue Huang Date: Tue, 3 Dec 2024 10:30:14 -0800 Subject: [PATCH] Fix expert rank generator Signed-off-by: Guyue Huang --- .../nlp/modules/common/megatron/megatron_init.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/nemo/collections/nlp/modules/common/megatron/megatron_init.py b/nemo/collections/nlp/modules/common/megatron/megatron_init.py index d58b97e3dd10..ff16d36fb4bb 100644 --- a/nemo/collections/nlp/modules/common/megatron/megatron_init.py +++ b/nemo/collections/nlp/modules/common/megatron/megatron_init.py @@ -361,12 +361,9 @@ def fake_initialize_model_parallel( expert_decoder_rank_generator = RankGenerator( tp=tensor_model_parallel_size, # the same as Attention part ep=expert_model_parallel_size_, - dp=( - decoder_world_size - // (expert_model_parallel_size_ * tensor_model_parallel_size * pipeline_model_parallel_size) - ), + dp=(data_parallel_size // expert_model_parallel_size_), pp=pipeline_model_parallel_size, - cp=1, + cp=context_parallel_size, order='tp-pp-dp' if use_tp_pp_dp_mapping else 'tp-cp-ep-dp-pp', rank_offset=encoder_world_size, )