Fix missing parallelisms (NVIDIA#9725) (NVIDIA#9758)

* pass cp and ep cfg to init mp * update test --------- Signed-off-by: Maanu Grover <[email protected]> Signed-off-by: Alexandros Koumparoulis <[email protected]> Co-authored-by: Maanu Grover <[email protected]> Signed-off-by: Vivian Chen <[email protected]>
xuanzic · Aug 1, 2024 · 2eda68f · 2eda68f
1 parent 1b10e42
commit 2eda68f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/nemo/lightning/_strategy_lib.py b/nemo/lightning/_strategy_lib.py
@@ -61,12 +61,14 @@ def init_parallel_ranks(
         global_rank=init_global_rank,
         local_rank=init_local_rank,
         tensor_model_parallel_size=parallel_config.tensor_model_parallel_size,
+        expert_model_parallel_size=parallel_config.expert_model_parallel_size,
         pipeline_model_parallel_size=parallel_config.pipeline_model_parallel_size,
         virtual_pipeline_model_parallel_size=parallel_config.virtual_pipeline_model_parallel_size,
+        context_parallel_size=parallel_config.context_parallel_size,
         seed=seed,
         pipeline_model_parallel_split_rank=getattr(parallel_config, "pipeline_model_parallel_split_rank", None),
         use_fp8=fp8,
-        init_mpi_proc_group=getattr(parallel_config, "ub_tp_comm_overlap", False),
+        init_mpi_proc_group=getattr(parallel_config, "tp_comm_overlap", False),
         # apex_transformer_log_level=self.cfg.get('apex_transformer_log_level', 30),
     )
 
@@ -92,6 +94,8 @@ def init_model_parallel(model: Optional[nn.Module] = None) -> None:
                 pipeline_model_parallel_size=app_state.pipeline_model_parallel_size,
                 virtual_pipeline_model_parallel_size=app_state.virtual_pipeline_model_parallel_size,
                 pipeline_model_parallel_split_rank=app_state.pipeline_model_parallel_split_rank,
+                context_parallel_size=app_state.context_parallel_size,
+                expert_model_parallel_size=app_state.expert_model_parallel_size,
             )
 
             # assert that fake tp and pp rank match after model parallel init

diff --git a/tests/lightning/test_strategy_lib.py b/tests/lightning/test_strategy_lib.py
@@ -51,14 +51,18 @@ def test_init_parallel_ranks(mock_initialize_model_parallel) -> None:
 
     app_state.tensor_model_parallel_size = 2
     app_state.pipeline_model_parallel_size = 3
+    app_state.context_parallel_size = 2
+    app_state.expert_model_parallel_size = 2
     app_state.global_rank = 1
     app_state.local_rank = 0
 
     mock_parallel_config = MagicMock()
     mock_parallel_config.tensor_model_parallel_size = 2
     mock_parallel_config.pipeline_model_parallel_size = 3
     mock_parallel_config.virtual_pipeline_model_parallel_size = 4
-    mock_parallel_config.ub_tp_comm_overlap = False
+    mock_parallel_config.context_parallel_size = 2
+    mock_parallel_config.expert_model_parallel_size = 2
+    mock_parallel_config.tp_comm_overlap = False
     mock_parallel_config.pipeline_model_parallel_split_rank = None
 
     _strategy_lib.init_parallel_ranks(
@@ -76,6 +80,8 @@ def test_init_parallel_ranks(mock_initialize_model_parallel) -> None:
         tensor_model_parallel_size=2,
         pipeline_model_parallel_size=3,
         virtual_pipeline_model_parallel_size=4,
+        context_parallel_size=2,
+        expert_model_parallel_size=2,
         seed=1234,
         pipeline_model_parallel_split_rank=None,
         use_fp8=False,
@@ -93,6 +99,8 @@ def test_init_model_parallel(mock_mpu, *args):
     app_state.tensor_model_parallel_size = 2
     app_state.pipeline_model_parallel_size = 1
     app_state.pipeline_model_parallel_split_rank = None
+    app_state.context_parallel_size = 2
+    app_state.expert_model_parallel_size = 2
     app_state.init_mpi_proc_group = False
     app_state.tensor_model_parallel_rank = 2
     app_state.pipeline_model_parallel_rank = 0
@@ -105,6 +113,8 @@ def test_init_model_parallel(mock_mpu, *args):
         pipeline_model_parallel_size=1,
         virtual_pipeline_model_parallel_size=None,
         pipeline_model_parallel_split_rank=None,
+        context_parallel_size=2,
+        expert_model_parallel_size=2,
     )