Added knob for ub_tp_comm_overlap for the MCORE pass (NVIDIA#7902)

* Added knob fot user buffer/tensor parallel communication overlap for MCORE pass Signed-off-by: Selvaraj Anandaraj <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Selvaraj Anandaraj <[email protected]> Co-authored-by: Selvaraj Anandaraj <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Piotr Żelasko <[email protected]>
pzelasko · Jan 3, 2024 · fe1478e · fe1478e
1 parent 333a222
commit fe1478e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -1555,6 +1555,8 @@ def build_transformer_config(self) -> TransformerConfig:
         recompute_method = self.cfg.get('activations_checkpoint_method', None)
         recompute_num_layers = self.cfg.get('activations_checkpoint_num_layers', None)
 
+        ub_tp_comm_overlap = self.cfg.get('ub_tp_comm_overlap', False)
+
         if not self.cfg.get('fp8', False):
             fp8 = None
         elif self.cfg.get('fp8_e4m3', False):
@@ -1581,6 +1583,7 @@ def build_transformer_config(self) -> TransformerConfig:
             'recompute_method': recompute_method,
             'recompute_num_layers': recompute_num_layers,
             'distribute_saved_activations': False,  # not currently used in NeMo
+            'ub_tp_comm_overlap': ub_tp_comm_overlap,
             'fp8': fp8,
         }