Make pipelined TP comm overlap available with mcore (#8005)

* Make pipelined TP comm overlap available with mcore Signed-off-by: Sangkug Lym <[email protected]> * remove unnecessary condition Signed-off-by: Sangkug Lym <[email protected]> --------- Signed-off-by: Sangkug Lym <[email protected]>
NVIDIA · Jan 10, 2024 · 8bdcf37 · 8bdcf37
1 parent 58d6bce
commit 8bdcf37
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -852,12 +852,14 @@ def _validate_and_override_config(self):
                 ) % vp_size == 0, 'Make sure the number of model chunks is the same across all pipeline stages.'
 
         if self.cfg.get('ub_tp_comm_overlap', False):
-            if not self.cfg.get('transformer_engine', False) or not self.cfg.get('sequence_parallel', False):
+            if not self.cfg.get('sequence_parallel', False):
                 logging.info(
-                    "Userbuffer tensor-parallel communication overlap is available with both Transformer Engine and sequence-parallelism."
+                    "Pipelined tensor-parallel communication overlap is available with sequence-parallelism."
+                    "Setting `ub_tp_comm_overlap` to False."
                 )
                 with open_dict(self.cfg):
                     self.cfg.ub_tp_comm_overlap = False
+
             if self.cfg.get('fsdp', False):
                 logging.info(
                     "Userbuffer tensor-parallel communication overlap is not available with FSDP."