EP: add expert_* to app_state in init_model_parallel.

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Jan 26, 2024 · 34fb998 · 34fb998
1 parent 38dee6c
commit 34fb998
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py
@@ -134,7 +134,8 @@ def init_model_parallel(sharp: bool, nccl_communicator_config_path: str = None)
             assert app_state.tensor_model_parallel_rank == parallel_state.get_tensor_model_parallel_rank()
             assert app_state.pipeline_model_parallel_rank == parallel_state.get_pipeline_model_parallel_rank()
 
-            # TODO(akoumparouli): expert parallelism
+            app_state.expert_model_parallel_rank = parallel_state.get_expert_model_parallel_rank()
+            app_state.expert_model_parallel_size = parallel_state.get_expert_model_parallel_world_size()
             app_state.tensor_model_parallel_group = parallel_state.get_tensor_model_parallel_group()
             app_state.data_parallel_group = parallel_state.get_data_parallel_group()
             app_state.data_parallel_rank = parallel_state.get_data_parallel_rank()