add call to p2p overlap (#6779)

* add call to p2p overlap Signed-off-by: Abhinav Khattar <[email protected]> * update Jenkins for test Signed-off-by: Abhinav Khattar <[email protected]> --------- Signed-off-by: Abhinav Khattar <[email protected]>
NVIDIA · Jun 1, 2023 · 6b5a9e7 · 6b5a9e7
1 parent 8672af6
commit 6b5a9e7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -57,6 +57,15 @@ pipeline {
       }
     }
 
+    stage('Megatron Core installation') {
+      steps {
+        sh 'git clone https://github.com/NVIDIA/Megatron-LM.git && \
+            cd Megatron-LM && \
+            git checkout e6d7e09845590d0a36bc7f29eb28db974fb8da4e && \
+            pip install -e .'
+      }
+    }
+
     stage('PyTorch Lightning version') {
       steps {
         sh 'python -c "import pytorch_lightning; print(pytorch_lightning.__version__)"'

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