Rachitg/ag (#9081)

* disable overlap for qkv Signed-off-by: Rachit Garg <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * bug fix * bugfix --------- Signed-off-by: Rachit Garg <[email protected]> Signed-off-by: Rachit Garg <[email protected]> Co-authored-by: Rachit Garg <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
NVIDIA · May 1, 2024 · 7135609 · 7135609
1 parent 5ecc903
commit 7135609
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py b/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py
@@ -251,7 +251,7 @@ def __init__(
 
             te_version = packaging.version.Version(version("transformer-engine"))
             if te_version >= packaging.version.Version("1.5.0dev") and (
-                not self.input_is_parallel and not model_parallel_config.tp_comm_atomic_ag
+                not self.input_is_parallel and model_parallel_config.tp_comm_disable_qkv
             ):
                 # TE 1.5 introduces the option `return_layernorm_output_gathered`, so the all gather
                 # in the forward method is not needed, so set self._sequence_parallel to False