check restore_config first

Signed-off-by: Alexandros Koumparoulis <[email protected]>

Author: akoumpa

nemo/lightning/resume.py

@@ -103,23 +103,7 @@ def setup(self, trainer: Union[pl.Trainer, fl.Fabric], model=None):
         if isinstance(trainer, fl.Fabric):
             raise NotImplementedError("Fabric is not supported yet.")
 
-        trainer_ckpt_path = self.get_trainer_ckpt_path(model)
-        if trainer_ckpt_path:
-            trainer.ckpt_path = trainer_ckpt_path
-            trainer.checkpoint_callback.last_model_path = trainer_ckpt_path
-            # Load artifacts
-            if getattr(self.restore_config, 'load_artifacts', False):
-                if isinstance(trainer_ckpt_path, AdapterPath):
-                    # load tokenizer from the base model during peft resume, in case the first peft checkpoint
-                    # is deleted before the current peft checkpoint is saved
-                    context_path = trainer_ckpt_path.base_model_path / "context"
-                    if not context_path.exists():
-                        context_path = trainer_ckpt_path.base_model_path
-                else:
-                    context_path = self.get_context_path(model)
-                model = _try_restore_tokenizer(model, context_path)
-
-        elif self.restore_config:
+        if self.restore_config:
             new_path = self._extract_path(
                 model=model,
                 path=self.restore_config.path,
@@ -139,6 +123,21 @@ def setup(self, trainer: Union[pl.Trainer, fl.Fabric], model=None):
 
                 _try_restore_tokenizer(model, context_path)
 
+        elif (trainer_ckpt_path := self.get_trainer_ckpt_path(model)) is not None:
+            trainer.ckpt_path = trainer_ckpt_path
+            trainer.checkpoint_callback.last_model_path = trainer_ckpt_path
+            # Load artifacts
+            if getattr(self.restore_config, 'load_artifacts', False):
+                if isinstance(trainer_ckpt_path, AdapterPath):
+                    # load tokenizer from the base model during peft resume, in case the first peft checkpoint
+                    # is deleted before the current peft checkpoint is saved
+                    context_path = trainer_ckpt_path.base_model_path / "context"
+                    if not context_path.exists():
+                        context_path = trainer_ckpt_path.base_model_path
+                else:
+                    context_path = self.get_context_path(model)
+                model = _try_restore_tokenizer(model, context_path)
+
     def _extract_path(
         self, model: Optional[io.ConnectorMixin], path: str, adapter_path: Optional[str] = None
     ) -> BasePath: