Add ability to enable/disable act ckpt and seq parallelism in GPT

nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py

@@ -159,6 +159,19 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
 
         self.get_attention_mask_from_fusion = self.cfg.get('get_attention_mask_from_fusion', True)
 
+        # Save activations checkpointing and sequence parallelism parameters to be able to restore them later.
+        if self.cfg.get('megatron_amp_O2', False):
+            base_module = self.model.module
+        else:
+            base_module = self.model
+        self.original_checkpointing_granularity = base_module.language_model.encoder.activations_checkpoint_granularity
+        self.original_checkpointing_num_layers = base_module.language_model.encoder.activations_checkpoint_num_layers
+        self.original_checkpointing_method = base_module.language_model.encoder.activations_checkpoint_method
+        self.original_activations_checkpoint_layers_per_pipeline = (
+            base_module.language_model.encoder.activations_checkpoint_layers_per_pipeline
+        )
+        self.original_sequence_parallel = base_module.language_model.encoder.sequence_parallel
+
     def set_inference_config(self, inference_config):
         self._inference_config = inference_config
 
@@ -1069,3 +1082,42 @@ def on_train_batch_end(self, outputs, dataloader_iter: Any, batch_idx: int, unus
                     # Reset the optimizer update skipped to `None` - this is to prevent scheduler no-ops during
                     # accumulated gradient updates.
                     grad_scaler.optimizer_update_skipped = None
+
+    def _reset_activation_checkpointing_args(self):
+        if self.cfg.get('megatron_amp_O2', False):
+            base_module = self.model.module
+        else:
+            base_module = self.model
+
+        base_module.language_model.encoder.activations_checkpoint_granularity = None
+        base_module.language_model.encoder.activations_checkpoint_method = None
+        base_module.language_model.encoder.activations_checkpoint_num_layers = None
+        base_module.language_model.encoder.activations_checkpoint_layers_per_pipeline = None
+
+    def _restore_activation_checkpointing_args(self):
+        if self.cfg.get('megatron_amp_O2', False):
+            base_module = self.model.module
+        else:
+            base_module = self.model
+        base_module.language_model.encoder.activations_checkpoint_granularity = self.original_checkpointing_granularity
+        base_module.language_model.encoder.activations_checkpoint_method = self.original_checkpointing_method
+        base_module.language_model.encoder.activations_checkpoint_num_layers = self.original_checkpointing_num_layers
+        base_module.language_model.encoder.activations_checkpoint_layers_per_pipeline = (
+            self.original_activations_checkpoint_layers_per_pipeline
+        )
+
+    def _reset_sequence_parallelism_args(self):
+        if self.cfg.get('megatron_amp_O2', False):
+            base_module = self.model.module
+        else:
+            base_module = self.model
+
+        base_module.language_model.encoder.sequence_parallel = None
+
+    def _restore_sequence_parallelism_args(self):
+        if self.cfg.get('megatron_amp_O2', False):
+            base_module = self.model.module
+        else:
+            base_module = self.model
+
+        base_module.language_model.encoder.sequence_parallel = self.original_sequence_parallel