Fix bug in MegatronParallel

NVIDIA · Jun 10, 2024 · cd8d333 · cd8d333
1 parent 8c522c3
commit cd8d333
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/lightning/megatron_parallel.py b/nemo/lightning/megatron_parallel.py
@@ -132,13 +132,13 @@ def __init__(
                         _model.configure_model()
                     _pipeline.append(_model)
 
-            if self.ddp_config is not None:
+            if isinstance(ddp_config, DistributedDataParallelConfig):
                 from megatron.core.distributed import DistributedDataParallel as McoreDDP
 
                 _pipeline = [
                     McoreDDP(
                         model_chunk.config,
-                        self.ddp_config,
+                        ddp_config,
                         model_chunk,
                         data_parallel_group=parallel_state.get_data_parallel_group(with_context_parallel=True),
                         expert_data_parallel_group=parallel_state.get_data_modulo_expert_parallel_group(),