add deallocate pipeline output optimization

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
JimmyZhang12 · Jan 30, 2024 · 788ed1b · 788ed1b
1 parent 7b2415a
commit 788ed1b
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