fix bert test

Signed-off-by: dimapihtar <[email protected]>
NVIDIA · May 25, 2024 · e9eedb1 · e9eedb1
1 parent b2e6b88
commit e9eedb1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
@@ -2798,7 +2798,7 @@ jobs:
             model.megatron_amp_O2=True \
             trainer.gradient_clip_val=1.0 \
             exp_manager.exp_dir=examples/nlp/language_modeling/bert_pretrain_results \
-            model.pipeline_model_parallel_size=2 \
+            model.tensor_model_parallel_size=2 \
             model.optim.name=fused_adam \
             model.optim.lr=2e-4 \
             model.optim.sched.warmup_steps=2 \
@@ -2829,7 +2829,7 @@ jobs:
             trainer.gradient_clip_val=1.0 \
             exp_manager.exp_dir=examples/nlp/language_modeling/bert_pretrain_results \
             exp_manager.resume_if_exists=True \
-            model.pipeline_model_parallel_size=2 \
+            model.tensor_model_parallel_size=2 \
             model.optim.name=fused_adam \
             model.optim.lr=2e-4 \
             model.optim.sched.warmup_steps=2 \