# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union
import numpy as np
import torch
import wrapt
from transformers import AutoModel, AutoTokenizer
from nemo_deploy import ITritonDeployable
from nemo_export.utils import (
get_example_inputs,
get_model_device_type,
is_nemo2_checkpoint,
validate_fp8_network,
)
from nemo_export_deploy_common.import_utils import (
MISSING_MODELOPT_MSG,
MISSING_NEMO_MSG,
MISSING_TENSORRT_MSG,
UnavailableError,
)
try:
from nemo.utils import logging
except (ImportError, ModuleNotFoundError):
import logging
logging = logging.getLogger(__name__)
try:
import modelopt.torch.quantization as mtq
HAVE_MODELOPT = True
except (ImportError, ModuleNotFoundError):
from unittest.mock import MagicMock
mtq = MagicMock()
HAVE_MODELOPT = False
try:
import tensorrt as trt
HAVE_TENSORRT = True
except (ImportError, ModuleNotFoundError):
from unittest.mock import MagicMock
trt = MagicMock()
HAVE_TENSORRT = False
try:
from nemo.collections.llm.modelopt.quantization.quant_cfg_choices import (
get_quant_cfg_choices,
)
QUANT_CFG_CHOICES = get_quant_cfg_choices()
HAVE_NEMO = True
except (ImportError, ModuleNotFoundError):
HAVE_NEMO = False
[docs]
@wrapt.decorator
def noop_decorator(func):
"""No op decorator."""
def wrapper(*args, **kwargs):
return func(*args, **kwargs)
return wrapper
use_pytriton = True
batch = noop_decorator
try:
from pytriton.decorators import batch
except Exception:
logging.warning("PyTriton is not available.")
use_pytriton = False
use_onnxruntime = True
try:
import onnxruntime
except Exception:
logging.warning("onnxruntime is not available.")
use_onnxruntime = False
# pylint: disable=line-too-long
[docs]
class OnnxLLMExporter(ITritonDeployable):
"""Exports models to ONNX and run fast inference.
Example:
from nemo_export.onnx_llm_exporter import OnnxLLMExporter
onnx_llm_exporter = OnnxLLMExporter(
onnx_model_dir="/path/for/onnx_model/files",
model_name_or_path="/path/for/model/files",
)
onnx_llm_exporter.export(
input_names=["input_ids", "attention_mask", "dimensions"],
output_names=["embeddings"],
)
output = onnx_llm_exporter.forward(["Hi, how are you?", "I am good, thanks, how about you?"])
print("output: ", output)
"""
def __init__(
self,
onnx_model_dir: str,
model: Optional[torch.nn.Module] = None,
tokenizer=None,
model_name_or_path: str = None,
load_runtime: bool = True,
):
"""Initializes the ONNX Exporter.
Args:
onnx_model_dir (str): path for storing the ONNX model files.
model (Optional[torch.nn.Module]): torch model.
tokenizer (HF or NeMo tokenizer): tokenizer class.
model_name_or_path (str): a path for ckpt or HF model ID
load_runtime (bool): load ONNX runtime if there is any exported model available in
the onnx_model_dir folder.
"""
self.onnx_model_dir = onnx_model_dir
self.model_name_or_path = model_name_or_path
self.onnx_model_path = str(Path(onnx_model_dir) / "model.onnx")
self.model = model
self.tokenizer = tokenizer
self.model_input_names = None
self.model_output_names = None
self.onnx_runtime_session = None
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if self.model_name_or_path is not None:
if model is not None:
raise ValueError("A model was also passed but it will be overridden.")
if Path(self.model_name_or_path).is_dir():
if is_nemo2_checkpoint(self.model_name_or_path):
raise NotImplementedError("NeMo 2.0 checkpoint will be supported later.")
else:
self._load_hf_model()
self.model.to(self.device)
if load_runtime:
self._load_runtime()
[docs]
def _load_runtime(self):
if use_onnxruntime:
if Path(self.onnx_model_path).exists():
self.onnx_runtime_session = onnxruntime.InferenceSession(self.onnx_model_path)
self.model_input_names = [input.name for input in self.onnx_runtime_session.get_inputs()]
self.model_output_names = [output.name for output in self.onnx_runtime_session.get_outputs()]
self.tokenizer = AutoTokenizer.from_pretrained(
Path(self.onnx_model_dir) / "tokenizer", trust_remote_code=True
)
[docs]
def _load_hf_model(self):
self.model = AutoModel.from_pretrained(
self.model_name_or_path,
trust_remote_code=True,
).eval()
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, trust_remote_code=True)
[docs]
def export(
self,
input_names: list,
output_names: list,
example_inputs: dict = None,
opset: int = 20,
dynamic_axes_input: Optional[dict] = None,
dynamic_axes_output: Optional[dict] = None,
export_dtype: str = "fp32",
verbose: bool = False,
):
"""Performs ONNX conversion from a PyTorch model.
Args:
input_names (list): input parameter names of the model that ONNX will export will use.
output_names (list): output parameter names of the model that ONNX will export will use.
example_inputs (dict): example input for the model to build the engine.
opset (int): ONNX opset version. Default is 20.
dynamic_axes_input (dict): Variable length axes for the input.
dynamic_axes_output (dict): Variable length axes for the output.
export_dtype (str): Export dtype, fp16 or fp32.
verbose (bool): Enable verbose or not.
"""
self._export_to_onnx(
input_names=input_names,
example_inputs=example_inputs,
output_names=output_names,
opset=opset,
dynamic_axes_input=dynamic_axes_input,
dynamic_axes_output=dynamic_axes_output,
export_dtype=export_dtype,
verbose=verbose,
)
self._load_runtime()
[docs]
def _export_to_onnx(
self,
input_names: list,
output_names: list,
example_inputs: dict = None,
opset: int = 20,
dynamic_axes_input: Optional[dict] = None,
dynamic_axes_output: Optional[dict] = None,
export_dtype: Union[torch.dtype, str] = "fp16",
verbose: bool = False,
):
if example_inputs is None:
example_inputs = get_example_inputs(self.tokenizer, self.device)
if "dimensions" in input_names:
example_inputs["dimensions"] = torch.ones(
len(example_inputs["input_ids"]), dtype=torch.int64, device=self.device
)
if isinstance(export_dtype, str):
export_dtype = {"fp16": torch.float16, "fp32": torch.float32}[export_dtype]
self.model.to(export_dtype)
Path(self.onnx_model_dir).mkdir(parents=True, exist_ok=True)
with torch.autocast(device_type=get_model_device_type(self.model), dtype=export_dtype):
torch.onnx.export(
model=self.model,
args=(example_inputs,),
f=self.onnx_model_path,
input_names=input_names,
output_names=output_names,
dynamic_axes={**dynamic_axes_input, **dynamic_axes_output},
verbose=verbose,
opset_version=opset,
)
logging.info(f"Successfully exported PyTorch model to ONNX model {self.onnx_model_path}")
existing_directory_path = Path(self.onnx_model_dir) / "tokenizer"
existing_directory_path.mkdir(exist_ok=True)
self.tokenizer.save_pretrained(existing_directory_path)
[docs]
def export_onnx_to_trt(
self,
trt_model_dir: str,
profiles=None,
override_layernorm_precision_to_fp32: bool = False,
override_layers_to_fp32: List = None,
trt_dtype: str = "fp16",
profiling_verbosity: str = "layer_names_only",
trt_builder_flags: List[trt.BuilderFlag] = None,
) -> None:
"""Performs TensorRT conversion from an ONNX model.
Args:
trt_model_dir: path to store the TensorRT model.
profiles: TensorRT profiles.
override_layernorm_precision_to_fp32 (bool): whether to convert layers to fp32 or not.
override_layers_to_fp32 (List): Layer names to be converted to fp32.
trt_dtype (str): "fp16" or "fp32".
profiling_verbosity (str): Profiling verbosity. Default is "layer_names_only".
trt_builder_flags (List[trt.BuilderFlag]): TRT specific flags.
"""
if not HAVE_TENSORRT:
raise UnavailableError(MISSING_TENSORRT_MSG)
logging.info(f"Building TRT engine from ONNX model ({self.onnx_model_path})")
trt_logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(trt_logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()
parser = trt.OnnxParser(network, trt_logger)
# we use parse_from_file() instead of parse() because it can be used for both single
# file models as well as externally stored models (required when model >2GiB)
if not parser.parse_from_file(self.onnx_model_path):
logging.warning("ONNX model could not be parsed")
for error in range(parser.num_errors):
logging.error(parser.get_error(error))
return
if profiles:
for profile in profiles:
optimization_profile = builder.create_optimization_profile()
for i in range(network.num_inputs):
in_tensor = network.get_input(i)
optimization_profile.set_shape(
in_tensor.name,
min=profile[in_tensor.name][0],
opt=profile[in_tensor.name][1],
max=profile[in_tensor.name][2],
)
config.add_optimization_profile(optimization_profile)
if trt_dtype == "fp16":
logging.info("Setting Build Flag FP16")
config.set_flag(trt.BuilderFlag.FP16)
elif trt_dtype == "fp8":
# With FP8 export we want to also enable FP16 layers as a fallback instead of FP32
logging.info("Setting Build Flag FP8 and FP16")
config.set_flag(trt.BuilderFlag.FP8)
config.set_flag(trt.BuilderFlag.FP16)
validate_fp8_network(network)
# patch network
if override_layernorm_precision_to_fp32:
logging.info("Overriding TensorRT network LayerNorm precision to float32.")
self._override_layernorm_precision_to_fp32(network)
if override_layers_to_fp32:
logging.info("Overriding some layers to float32.")
self._override_layers_to_fp32(network, override_layers_to_fp32)
try:
config.profiling_verbosity = {
"detailed": trt.ProfilingVerbosity.DETAILED,
"layer_names_only": trt.ProfilingVerbosity.LAYER_NAMES_ONLY,
"none": trt.ProfilingVerbosity.NONE,
}[profiling_verbosity]
except KeyError:
error_msg = "Unknown profiling verbosity value."
raise ValueError(error_msg)
logging.info(f"Setting Profiling Verbosity to {config.profiling_verbosity}")
if trt_builder_flags is not None:
for flag in trt_builder_flags:
config.set_flag(flag)
engine_string = builder.build_serialized_network(network, config)
if engine_string is None:
raise Exception("Failed to serialize the TensorRT Engine. Please check the TensorRT logs for details")
trt_model_path = Path(trt_model_dir)
trt_model_path.mkdir(parents=True, exist_ok=True)
trt_model_path = trt_model_path / "model.plan"
trt_model_path.write_bytes(engine_string)
logging.info(f"Successfully exported ONNX model ({self.onnx_model_path}) to TRT engine ({trt_model_path})")
[docs]
def _override_layer_precision_to_fp32(self, layer: trt.ILayer) -> None:
if not HAVE_TENSORRT:
raise UnavailableError(MISSING_TENSORRT_MSG)
layer.precision = trt.float32
layer.set_output_type(0, trt.float32)
[docs]
def _override_layers_to_fp32(self, network: trt.INetworkDefinition, fp32_layer_patterns: list[str]) -> None:
if not HAVE_TENSORRT:
raise UnavailableError(MISSING_TENSORRT_MSG)
for i in range(network.num_layers):
layer = network.get_layer(i)
layer_name = layer.name
if any(layer_name.startswith(pattern) for pattern in fp32_layer_patterns) and layer.precision in {
trt.float32,
trt.float16,
}:
if layer.type in {trt.LayerType.CAST}:
logging.info(f"Skipping overriding {layer.type} layer {i} {layer_name} dtype")
continue
if any(
layer.get_input(input_idx).dtype in {trt.float32, trt.float16}
for input_idx in range(layer.num_inputs)
):
# Note: Assigning to layer.precision (even the same value) sets precision_is_set=True,
# which prevents TensorRT from changing this layer's precision
layer.precision = trt.float32
logging.info(f"Setting layer {i} {layer_name} (type: {layer.type}) precision to FP32")
for j in range(layer.num_outputs):
if layer.get_output_type(j) in {trt.float32, trt.float16}:
layer.set_output_type(j, trt.float32)
logging.info(f"Setting layer {i} {layer_name} (type: {layer.type}) output type {j} to FP32")
[docs]
def _override_layernorm_precision_to_fp32(self, network: trt.INetworkDefinition) -> None:
"""Set the precision of LayerNorm subgraphs to FP32 to preserve accuracy.
- https://nvbugs/4478448 (Mistral)
- https://nvbugs/3802112 (T5)
Args:
network: tensorrt.INetworkDefinition
"""
# Logic originally from OSS T5 HF export script:
# https://gitlab-master.nvidia.com/TensorRT/Public/oss/-/blob/77495ec/demo/HuggingFace/T5/export.py
if not HAVE_TENSORRT:
raise UnavailableError(MISSING_TENSORRT_MSG)
pow_ops = {}
for layer_index, layer in enumerate(network):
if layer.type == trt.LayerType.IDENTITY:
all_fp32 = all(
[
layer.output_type_is_set(o) and layer.get_output_type(o) == trt.float32
for o in range(layer.num_outputs)
]
)
if all_fp32:
if layer.get_input(0).dtype == trt.float32:
layer.precision = trt.float32
if layer.type == trt.LayerType.ELEMENTWISE:
layer.__class__ = getattr(trt, "IElementWiseLayer")
if layer.op == trt.ElementWiseOperation.POW:
pow_ops[layer] = layer_index
self._override_layer_precision_to_fp32(layer)
for _, index in pow_ops.items():
# Iterate from few layers before pow to include residual add and cast op.
# Iterate till 10 layers after pow op to include all
# operations included in layer norm.
START_OFFSET = 4
END_OFFSET = 12
for i in range(index - START_OFFSET, index + END_OFFSET):
layer = network.get_layer(i)
if layer.type == trt.LayerType.REDUCE:
self._override_layer_precision_to_fp32(layer)
if layer.type == trt.LayerType.ELEMENTWISE:
layer.__class__ = getattr(trt, "IElementWiseLayer")
if layer.op == trt.ElementWiseOperation.SUM:
self._override_layer_precision_to_fp32(layer)
if layer.type == trt.LayerType.UNARY:
layer.__class__ = getattr(trt, "IUnaryLayer")
if layer.op == trt.UnaryOperation.SQRT:
self._override_layer_precision_to_fp32(layer)
if layer.type == trt.LayerType.ELEMENTWISE:
layer.__class__ = getattr(trt, "IElementWiseLayer")
if layer.op == trt.ElementWiseOperation.DIV:
self._override_layer_precision_to_fp32(layer)
if layer.type == trt.LayerType.ELEMENTWISE:
layer.__class__ = getattr(trt, "IElementWiseLayer")
if layer.op == trt.ElementWiseOperation.PROD:
self._override_layer_precision_to_fp32(layer)
[docs]
def forward(self, inputs: Union[List, Dict], dimensions: Optional[List] = None):
"""Run inference for a given input.
Args:
inputs (Union[List, Dict]): Input for the model. If list, it should be a list of strings.
If dict, it should be a dictionary with keys as the model input names.
dimensions (Optional[List]): The dimensions parameter of the model. Required if the model
was exported to accept dimensions parameter and inputs is given as a list of strings.
Returns:
np.ndarray: Model output.
"""
if self.onnx_runtime_session is None:
warnings.warn("ONNX Runtime is not available. Please install the onnxruntime-gpu and try again.")
return None
if isinstance(inputs, List):
if "dimensions" in self.model_input_names and dimensions is None:
raise ValueError("Dimensions should be provided for list input.")
inputs = dict(self.tokenizer(inputs))
inputs["dimensions"] = dimensions
output = self.onnx_runtime_session.run(self.model_output_names, inputs)
return output[0]
[docs]
def quantize(
self,
quant_cfg: Union[str, Dict[str, Any]],
forward_loop: Optional[Callable],
) -> None:
"""Quantize the model by calibrating it using a given forward loop.
Args:
quant_cfg (str or dict): The quantization configuration to use.
forward_loop (callable): A function that accepts the model as a single parameter
and runs sample data through it. This is used for calibration during quantization.
"""
if not HAVE_NEMO:
raise UnavailableError(MISSING_NEMO_MSG)
if not HAVE_MODELOPT:
raise UnavailableError(MISSING_MODELOPT_MSG)
if isinstance(quant_cfg, str):
assert quant_cfg in QUANT_CFG_CHOICES, (
f"Quantization config {quant_cfg} is not supported. Supported configs: {list(QUANT_CFG_CHOICES)}"
)
quant_cfg = QUANT_CFG_CHOICES[quant_cfg]
logging.info("Starting quantization...")
mtq.quantize(self.model, quant_cfg, forward_loop=forward_loop)
logging.info("Quantization is completed.")
@property
def get_model(self):
"""Returns the model."""
return self.model
@property
def get_tokenizer(self):
"""Returns the tokenizer."""
return self.tokenizer
@property
def get_model_input_names(self):
"""Returns the model input names."""
return self.model_input_names
@property
def get_triton_input(self):
"""Get triton input."""
raise NotImplementedError("This function will be implemented later.")
@property
def get_triton_output(self):
"""Get triton output."""
raise NotImplementedError("This function will be implemented later.")
[docs]
@batch
def triton_infer_fn(self, **inputs: np.ndarray):
"""PyTriton inference function."""
raise NotImplementedError("This function will be implemented later.")
