fix mcore version issue

Signed-off-by: ashors1 <[email protected]>
NVIDIA · Dec 2, 2024 · 0ba86c4 · 0ba86c4
1 parent 30954ab
commit 0ba86c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/megatron_init.py b/nemo/collections/nlp/modules/common/megatron/megatron_init.py
@@ -446,7 +446,7 @@ def generator_wrapper(group_type, **kwargs):
     # EP rank
     expert_model_parallel_rank = 0
     if expert_model_parallel_size_ is not None and expert_model_parallel_size_ > 1:
-        for ranks in generator_wrapper('ep', independent_ep=True):
+        for ranks in generator_wrapper('ep'):
             if rank in ranks:
                 expert_model_parallel_rank = list(ranks).index(rank)
 

diff --git a/nemo/lightning/megatron_init.py b/nemo/lightning/megatron_init.py
@@ -446,7 +446,7 @@ def generator_wrapper(group_type, **kwargs):
     # EP rank
     expert_model_parallel_rank = 0
     if expert_model_parallel_size_ is not None and expert_model_parallel_size_ > 1:
-        for ranks in generator_wrapper('ep', independent_ep=True):
+        for ranks in generator_wrapper('ep'):
             if rank in ranks:
                 expert_model_parallel_rank = list(ranks).index(rank)