Fix "Safely import optional python packages (NVIDIA-NeMo#13936)" (NVIDIA-NeMo#14198)

* Revert "Revert "Safely import optional python packages (NVIDIA-NeMo#13936)" (NVIDIA-NeMo#14197)"

This reverts commit 808845b.

Signed-off-by: Charlie Truong <[email protected]>

* Fix LocalCheckpointCallback safe import

Signed-off-by: Charlie Truong <[email protected]>

---------

Signed-off-by: Charlie Truong <[email protected]>
Signed-off-by: Guyue Huang <[email protected]>

Author: chtruong814

nemo/collections/llm/modelopt/speculative/model_transform.py

@@ -16,17 +16,19 @@
 
 from nemo.collections.llm import GPTModel
 from nemo.utils import logging
-from nemo.utils.import_utils import safe_import
+from nemo.utils.import_utils import UnavailableError, safe_import
 from nemo.utils.model_utils import unwrap_model
 
 mto, HAVE_MODELOPT = safe_import("modelopt.torch.opt")
 mtsp, _ = safe_import("modelopt.torch.speculative")
 
-
-ALGORITHMS = {
-    "eagle3": mtsp.EAGLE3_DEFAULT_CFG,
-    # more TBD
-}
+try:
+    ALGORITHMS = {
+        "eagle3": mtsp.EAGLE3_DEFAULT_CFG,
+        # more TBD
+    }
+except UnavailableError:
+    ALGORITHMS = {}
 
 
 def apply_speculative_decoding(model: nn.Module, algorithm: str = "eagle3") -> nn.Module:

nemo/lightning/nemo_logger.py

@@ -22,12 +22,14 @@
 import lightning.pytorch as pl
 from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint as PTLModelCheckpoint
 from lightning.pytorch.loggers import Logger, TensorBoardLogger, WandbLogger
-from nvidia_resiliency_ext.ptl_resiliency.local_checkpoint_callback import LocalCheckpointCallback
 
 from nemo.lightning.io.mixin import IOMixin
 from nemo.lightning.pytorch.callbacks import ModelCheckpoint
 from nemo.utils import logging
 from nemo.utils.app_state import AppState
+from nemo.utils.import_utils import safe_import
+
+lcp, HAVE_RES = safe_import('nvidia_resiliency_ext.ptl_resiliency.local_checkpoint_callback')
 
 
 @dataclass
@@ -203,7 +205,7 @@ def _setup_trainer_model_checkpoint(self, trainer, log_dir, ckpt=None):
         if ckpt:
             _overwrite_i = None
             for i, callback in enumerate(trainer.callbacks):
-                if isinstance(callback, PTLModelCheckpoint) and not isinstance(callback, LocalCheckpointCallback):
+                if isinstance(callback, PTLModelCheckpoint) and not isinstance(callback, lcp.LocalCheckpointCallback):
                     logging.warning(
                         "The Trainer already contains a ModelCheckpoint callback. " "This will be overwritten."
                     )
@@ -248,7 +250,7 @@ def _setup_trainer_model_checkpoint(self, trainer, log_dir, ckpt=None):
         from nemo.lightning import MegatronStrategy
 
         for callback in trainer.callbacks:
-            if isinstance(callback, PTLModelCheckpoint) and not isinstance(callback, LocalCheckpointCallback):
+            if isinstance(callback, PTLModelCheckpoint) and not isinstance(callback, lcp.LocalCheckpointCallback):
                 if callback.dirpath is None:
                     callback.dirpath = Path(log_dir / "checkpoints")
                 if callback.filename is None: