modify model parallel config as well

Signed-off-by: Jason Wang <[email protected]>
NVIDIA · Aug 18, 2023 · d2f3742 · d2f3742
1 parent b51d6f6
commit d2f3742
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -1348,6 +1348,7 @@ def _reset_sequence_parallelism_args(self):
 
         # Reset config values. Needed for calling generate.
         self.cfg.sequence_parallel = False
+        self.model_parallel_config.sequence_parallel = False
         self.transformer_config.sequence_parallel = False
 
         # Reset model parameters.
@@ -1363,6 +1364,7 @@ def _restore_sequence_parallelism_args(self):
         """
         # Restore config values.
         self.cfg.sequence_parallel = self.last_sequence_parallel
+        self.model_parallel_config.sequence_parallel = self.last_sequence_parallel
         self.transformer_config.sequence_parallel = self.last_sequence_parallel
 
         # Restore model parameters.