Fixed the tp overlap switch (#8195)

Signed-off-by: Selvaraj Anandaraj <[email protected]> Co-authored-by: Selvaraj Anandaraj <[email protected]> Co-authored-by: Layali R <[email protected]>
NVIDIA · Jan 31, 2024 · a4ef3d0 · a4ef3d0
1 parent e440c0c
commit a4ef3d0
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()