From db3461df2950e237f9c2d660d9029e8f14e75bc0 Mon Sep 17 00:00:00 2001 From: Selvaraj Anandaraj Date: Wed, 11 Oct 2023 14:51:54 -0700 Subject: [PATCH 1/2] Added knob fot user buffer/tensor parallel communication overlap for MCORE pass Signed-off-by: Selvaraj Anandaraj --- .../nlp/models/language_modeling/megatron_gpt_model.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py index b23fe0bbf886..33411b97c323 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py @@ -1554,6 +1554,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): @@ -1580,6 +1582,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, } From abff82c76dd8f76f6fecbd8e3b06a73849b3c17f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 11 Oct 2023 22:04:53 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../nlp/models/language_modeling/megatron_gpt_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py index 33411b97c323..b9ce5b3da1f1 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py @@ -1554,7 +1554,7 @@ 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) + ub_tp_comm_overlap = self.cfg.get('ub_tp_comm_overlap', False) if not self.cfg.get('fp8', False): fp8 = None