add deallocate pipeline output optimization (NVIDIA#8279) (NVIDIA#8318)

* add deallocate pipeline output optimization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jimmy Zhang <[email protected]> Co-authored-by: JimmyZhang12 <[email protected]> Co-authored-by: Jimmy Zhang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Zeeshan Patel <[email protected]>
zpx01 · Mar 8, 2024 · 62452bd · 62452bd
1 parent 39de66b
commit 62452bd
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -475,6 +475,7 @@ def build_transformer_config(self) -> TransformerConfig:
             'recompute_num_layers': recompute_num_layers,
             'distribute_saved_activations': False,  # not currently used in NeMo
             'fp8': None,
+            'deallocate_pipeline_outputs': True,
         }
 
         # populate the transformer config dict