f

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

Author: ko3n1g

nemo_export/onnx_llm_exporter.py

@@ -18,7 +18,6 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
-import tensorrt as trt
 import torch
 import wrapt
 from nemo.collections.llm.modelopt.quantization.quant_cfg_choices import (
@@ -36,6 +35,11 @@
 )
 from nemo_export_deploy_common.import_utils import UnavailableError
 
+try:
+    import tensorrt as trt
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt is not installed. Please install it with `pip install nvidia-tensorrt`.")
+
 try:
     import modelopt.torch.quantization as mtq
 except (ImportError, ModuleNotFoundError):

nemo_export/quantize/quantizer.py

@@ -21,13 +21,11 @@
 import torch.distributed as dist
 from megatron.core import parallel_state
 from megatron.core.transformer.module import Float16Module
-from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import (
-    MegatronGPTModel,
-)
 
 from nemo_export_deploy_common.import_utils import UnavailableError
 
 try:
+    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

nemo_export/tensorrt_lazy_compiler.py

@@ -22,8 +22,14 @@
 from typing import Any, Dict, List, Sequence, Tuple, Union
 
 import torch
-from nemo.utils.export_utils import add_casts_around_norms, replace_for_export
-from nemo.utils.import_utils import safe_import
+
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from nemo.utils.export_utils import add_casts_around_norms, replace_for_export
+    from nemo.utils.import_utils import safe_import
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("nemo is not installed. Please install it with `pip install nemo`.")
 
 polygraphy, polygraphy_imported = safe_import("polygraphy")
 if polygraphy_imported:

nemo_export/tensorrt_mm_exporter.py

@@ -21,17 +21,26 @@
 
 import numpy as np
 import wrapt
-from tensorrt_llm.runtime import MultimodalModelRunner as TRTLLMRunner
-
-from nemo_deploy import ITritonDeployable
-from nemo_export.multimodal.build import (
-    build_mllama_engine,
-    build_trtllm_engine,
-    build_visual_engine,
-    extract_lora_ckpt,
-)
-from nemo_export.multimodal.run import MultimodalModelRunner
-from nemo_export.tarutils import unpack_tarball
+
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from tensorrt_llm.runtime import MultimodalModelRunner as TRTLLMRunner
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt_llm is not installed. Please install it with `pip install tensorrt-llm`.")
+
+try:
+    from nemo_deploy import ITritonDeployable
+    from nemo_export.multimodal.build import (
+        build_mllama_engine,
+        build_trtllm_engine,
+        build_visual_engine,
+        extract_lora_ckpt,
+    )
+    from nemo_export.multimodal.run import MultimodalModelRunner
+    from nemo_export.tarutils import unpack_tarball
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("nemo is not installed. Please install it with `pip install nemo`.")
 
 use_deploy = True
 try:

nemo_export/trt_llm/converter/model_converter.py

@@ -18,16 +18,22 @@
 import numpy as np
 import tensorrt_llm
 import torch
-from tensorrt_llm._utils import pad_vocab_size
-from tensorrt_llm.functional import non_gated_version
-from tensorrt_llm.layers import MoeConfig
-from tensorrt_llm.models.modeling_utils import PretrainedConfig
 
 from nemo_export.trt_llm.converter.model_to_trt_llm_ckpt import (
     convert_model_to_trt_llm_ckpt,
     dist_model_to_trt_llm_ckpt,
 )
 from nemo_export.trt_llm.converter.utils import DECODER_MODEL_TYPE, split
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from tensorrt_llm._utils import pad_vocab_size
+    from tensorrt_llm.functional import non_gated_version
+    from tensorrt_llm.layers import MoeConfig
+    from tensorrt_llm.models.modeling_utils import PretrainedConfig
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt_llm is not installed. Please install it with `pip install tensorrt-llm`.")
+
 
 LOGGER = logging.getLogger("NeMo")
 

nemo_export/trt_llm/converter/model_to_trt_llm_ckpt.py

@@ -19,7 +19,6 @@
 from pathlib import Path
 
 import torch
-from tensorrt_llm._utils import pad_vocab_size, str_dtype_to_torch
 from tqdm import tqdm
 
 from nemo_export.trt_llm.converter.utils import (
@@ -29,6 +28,12 @@
     weights_dict,
 )
 from nemo_export.utils import torch_dtype_from_precision
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from tensorrt_llm._utils import pad_vocab_size, str_dtype_to_torch
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt_llm is not installed. Please install it with `pip install tensorrt-llm`.")
 
 LOGGER = logging.getLogger("NeMo")
 

nemo_export/trt_llm/qnemo/qnemo_to_tensorrt_llm.py

@@ -19,9 +19,13 @@
 import warnings
 from typing import List, Optional
 
-from tensorrt_llm.models import PretrainedConfig
-
 from nemo_export.trt_llm.qnemo.utils import CONFIG_NAME, WEIGHTS_NAME
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from tensorrt_llm.models import PretrainedConfig
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("tensorrt_llm is not installed. Please install it with `pip install tensorrt-llm`.")
 
 
 def qnemo_to_tensorrt_llm(