Fix nemo2 interface for ucc

Signed-off-by: Guyue Huang <guyueh@login-preos01.a51.clusters.nvidia.com>
guyueh1 · Jan 4, 2025 · 90a99d3 · 90a99d3
1 parent af35a8c
commit 90a99d3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/nemo/lightning/_strategy_lib.py b/nemo/lightning/_strategy_lib.py
@@ -82,7 +82,7 @@ def init_parallel_ranks(
         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,
-        pipeline_model_parallel_comm_backend=parallel_config.get('pipeline_model_parallel_comm_backend', 'nccl'),
+        pipeline_model_parallel_comm_backend=parallel_config.pipeline_model_parallel_comm_backend,
         virtual_pipeline_model_parallel_size=parallel_config.virtual_pipeline_model_parallel_size,
         context_parallel_size=parallel_config.context_parallel_size,
         encoder_tensor_model_parallel_size=getattr(parallel_config, "encoder_tensor_model_parallel_size", 0),

diff --git a/nemo/lightning/pytorch/strategies/megatron_strategy.py b/nemo/lightning/pytorch/strategies/megatron_strategy.py
@@ -905,6 +905,7 @@ def parallelism(self) -> ParallelismConfig:
         return ParallelismConfig(
             tensor_model_parallel_size=self.tensor_model_parallel_size,
             pipeline_model_parallel_size=self.pipeline_model_parallel_size,
+            pipeline_model_parallel_comm_backend=self.pipeline_model_parallel_comm_backend,
             virtual_pipeline_model_parallel_size=self.virtual_pipeline_model_parallel_size,
             microbatch_group_size_per_vp_stage=self.microbatch_group_size_per_vp_stage,
             context_parallel_size=self.context_parallel_size,