add call to p2p overlap (NVIDIA#6779) (NVIDIA#6786)

* add call to p2p overlap * update Jenkins for test --------- Signed-off-by: Abhinav Khattar <[email protected]> Signed-off-by: Eric Harper <[email protected]> Co-authored-by: Abhinav Khattar <[email protected]> Co-authored-by: Eric Harper <[email protected]>
CGranger-sorenson · Jun 5, 2023 · aa1986f · aa1986f
1 parent 76fc488
commit aa1986f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -61,7 +61,7 @@ pipeline {
       steps {
         sh 'git clone https://github.com/NVIDIA/Megatron-LM.git && \
             cd Megatron-LM && \
-            git checkout cd2537d444792b487b1ab5a6fa685e09c9957409 && \
+            git checkout e6d7e09845590d0a36bc7f29eb28db974fb8da4e && \
             pip install -e .'
       }
     }

diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -468,6 +468,8 @@ def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only):
             no_sync_func=no_sync_func,
             grad_sync_func=grad_sync_func,
             param_sync_func=param_sync_func,
+            overlap_p2p_comm=self.cfg.get('overlap_p2p_comm', False),
+            batch_p2p_comm=self.cfg.get('batch_p2p_comm', True),
         )
 
         # only the last stages of the pipeline return losses