fix tp_overlap config var name (NVIDIA#7928)

Signed-off-by: Xiaowei Ren <xren@nvidia.com>
rohitrango · Nov 22, 2023 · 1801f64 · 1801f64
1 parent 1a5ee38
commit 1801f64
Showing 1 changed file with 1 addition and 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
@@ -1585,7 +1585,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,
+            'tp_comm_overlap': ub_tp_comm_overlap,
             'fp8': fp8,
         }