f

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_export/multimodal/run.py

@@ -25,20 +25,25 @@
 
 import einops
 import numpy as np
-import tensorrt as trt
-import tensorrt_llm
-import tensorrt_llm.profiler as profiler
 import torch
 import yaml
 from PIL import Image
-from tensorrt_llm import logger
-from tensorrt_llm._utils import str_dtype_to_trt
-from tensorrt_llm.runtime import ModelRunner, Session, TensorInfo
 from torch.nn import functional as F
 from torchvision import transforms
 from transformers import AutoProcessor, CLIPImageProcessor
 
 from nemo_export.utils.constants import TRTLLM_ENGINE_DIR
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    import tensorrt as trt
+    import tensorrt_llm
+    import tensorrt_llm.profiler as profiler
+    from tensorrt_llm import logger
+    from tensorrt_llm._utils import str_dtype_to_trt
+    from tensorrt_llm.runtime import ModelRunner, Session, TensorInfo
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt_llm is not installed. Please install it with `pip install tensorrt-llm`.")
 
 
 def trt_dtype_to_torch(dtype):

nemo_export/onnx_llm_exporter.py

@@ -17,7 +17,6 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
 
-import modelopt.torch.quantization as mtq
 import numpy as np
 import tensorrt as trt
 import torch
@@ -35,6 +34,12 @@
     is_nemo2_checkpoint,
     validate_fp8_network,
 )
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    import modelopt.torch.quantization as mtq
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("modelopt is not installed. Please install it with `pip install nvidia-modelopt[torch]`.")
 
 QUANT_CFG_CHOICES = get_quant_cfg_choices()
 

nemo_export/quantize/quantizer.py

@@ -24,11 +24,17 @@
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import (
     MegatronGPTModel,
 )
-from nemo.collections.nlp.parts.utils_funcs import torch_dtype_from_precision
-from nemo.utils import logging
-from nemo.utils.distributed import temporary_directory
-from nemo.utils.model_utils import save_artifacts, unwrap_model
-from omegaconf.omegaconf import DictConfig, open_dict
+
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from nemo.collections.nlp.parts.utils_funcs import torch_dtype_from_precision
+    from nemo.utils import logging
+    from nemo.utils.distributed import temporary_directory
+    from nemo.utils.model_utils import save_artifacts, unwrap_model
+    from omegaconf.omegaconf import DictConfig, open_dict
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("nemo is not installed. Please install it with `pip install nemo`.")
 
 try:
     import modelopt.torch.quantization as mtq