Rachitg/ag (#9083)

* 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>

* Apply isort and black reformatting

Signed-off-by: michal2409 <[email protected]>

---------

Signed-off-by: Rachit Garg <[email protected]>
Signed-off-by: Rachit Garg <[email protected]>
Signed-off-by: michal2409 <[email protected]>
Co-authored-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>
Co-authored-by: Michal Futrega <[email protected]>
Co-authored-by: michal2409 <[email protected]>
Signed-off-by: Jan Lasek <[email protected]>

nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py

@@ -256,7 +256,9 @@ def __init__(
             from pkg_resources import packaging
 
             te_version = packaging.version.Version(version("transformer-engine"))
-            if te_version >= packaging.version.Version("1.5.0dev") and not model_parallel_config.tp_comm_overlap:
+            if te_version >= packaging.version.Version("1.5.0dev") and (
+                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
                 # unless TP communication overlap is used