Fix missing func for T5 model (NVIDIA#9141)

* add conditions Signed-off-by: Gao Deng <[email protected]> * condition check Signed-off-by: Gao Deng <[email protected]> --------- Signed-off-by: Gao Deng <[email protected]> Signed-off-by: Boxiang Wang <[email protected]>
BoxiangW · Jun 5, 2024 · 8775abe · 8775abe
1 parent 2a090bb
commit 8775abe
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py
@@ -238,7 +238,8 @@ def configure_ddp(self):
             hasattr(self.model, 'with_distributed_adam') and self.model.with_distributed_adam
         ):
             # do not use DDP if using megatron amp O2 or distributed optimizer
-            self.model.setup_mcore_distributed_parallel()
+            if self.model.use_mcore_dist_optim:
+                self.model.setup_mcore_distributed_parallel()
             self._model = self.model
         else:
             app_state = AppState()