Fixed the tp overlap switch (#8195)

Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com> Co-authored-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com> Co-authored-by: Layali R <31741533+layalir@users.noreply.github.com> Signed-off-by: Pablo Garay <pagaray@nvidia.com>
NVIDIA · Mar 19, 2024 · 56e54bf · 56e54bf
1 parent d9a15df
commit 56e54bf
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
@@ -1615,7 +1615,7 @@ def build_transformer_config(self) -> TransformerConfig:
             'layernorm_zero_centered_gamma': layernorm_zero_centered_gamma,
             'normalization': normalization,
             'fp8': fp8,
-            'ub_tp_comm_overlap': ub_tp_comm_overlap,
+            'tp_comm_overlap': ub_tp_comm_overlap,
         }
 
         transformer_config = super().build_transformer_config()