Bug fix to reset sequence parallelism (#6756)

* Bug fix to reset sequence parallelism Signed-off-by: Markel Sanz Ausin <[email protected]> * Update seq par reset/restore Signed-off-by: Markel Sanz Ausin <[email protected]> * Add nested loop Signed-off-by: Markel Sanz Ausin <[email protected]> --------- Signed-off-by: Markel Sanz Ausin <[email protected]>
NVIDIA · May 31, 2023 · a0f757e · a0f757e
1 parent dbd6a56
commit a0f757e
Showing 1 changed file with 7 additions and 4 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
@@ -1170,12 +1170,13 @@ def _reset_sequence_parallelism_args(self):
         self.last_sequence_parallel = self.cfg.sequence_parallel
 
         # Reset config values. Needed for calling generate.
-        self.cfg.sequence_parallel = None
+        self.cfg.sequence_parallel = False
 
         # Reset model parameters.
-
         for module in self.get_gpt_module_list():
-            module.language_model.encoder.sequence_parallel = None
+            for mod in module.modules():
+                if hasattr(mod, "sequence_parallel"):
+                    mod.sequence_parallel = self.last_sequence_parallel
 
     def _restore_sequence_parallelism_args(self):
         """ Restores the sequence parallelism parameters using the values saved by
@@ -1187,4 +1188,6 @@ def _restore_sequence_parallelism_args(self):
 
         # Restore model parameters.
         for module in self.get_gpt_module_list():
-            module.language_model.encoder.sequence_parallel = self.last_sequence_parallel
+            for mod in module.modules():
+                if hasattr(mod, "sequence_parallel"):
+                    mod.sequence_parallel = self.last_sequence_parallel