nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/onnx/autocast/referencerunner.py

@@ -0,0 +1,159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Reference runner module for ONNX model execution.
+
+This module provides functionality for running ONNX models using ONNXRuntime as a reference
+implementation. It supports both random input generation and user-provided inputs through
+NPZ or Polygraphy JSON files. The runner is used to analyze model behavior and validate
+outputs during precision conversion.
+"""
+
+import copy
+import io
+import sys
+from collections import OrderedDict
+
+import numpy as np
+import onnx
+
+from modelopt.onnx.autocast.logging_config import configure_logging, logger
+from modelopt.onnx.quantization.ort_utils import _prepare_ep_list
+
+configure_logging()
+
+
+class ReferenceRunner:
+    """A class to run ONNX models with ONNXRuntime for reference inference."""
+
+    def __init__(
+        self, model: onnx.ModelProto, providers: list[str] = ["cpu"], trt_plugins: list[str] = []
+    ):
+        """Initialize with ONNX model path."""
+        self.model = model
+        self.input_names = [input.name for input in self.model.graph.input]
+        self.providers = self._prepare_ep_list_with_trt_plugin_path(providers, trt_plugins)
+
+    def _prepare_ep_list_with_trt_plugin_path(self, providers, trt_plugins):
+        providers = _prepare_ep_list(providers) or providers
+        if "TensorrtExecutionProvider" in providers:
+            providers.remove("TensorrtExecutionProvider")
+            # Ensure that the TRT EP is the first in the providers list to avoid fallback issues
+            trt_ep_options = (
+                {"trt_extra_plugin_lib_paths": ";".join(trt_plugins)} if trt_plugins else {}
+            )
+            providers.insert(0, ("TensorrtExecutionProvider", trt_ep_options))
+            logger.info(f"Successfully updated EPs for ORT: {providers}")
+        return providers
+
+    def _load_inputs_from_json(self, input_data_path):
+        """Load inputs from Polygraphy JSON format."""
+        from polygraphy.json import load_json
+
+        return load_json(input_data_path, description="input data")
+
+    def _load_inputs_from_npz(self, input_data_path):
+        """Load inputs from NPZ format."""
+        return [np.load(input_data_path)]
+
+    def _validate_inputs(self, data_loader):
+        """Validate that input names match the model."""
+        if isinstance(data_loader, list) and (
+            isinstance(data_loader[0], (dict, np.lib.npyio.NpzFile))
+        ):
+            if sorted(self.input_names) != sorted(data_loader[0].keys()):
+                raise ValueError("Input names from ONNX model do not match provided input names.")
+        else:
+            raise ValueError("Invalid input file.")
+
+    def _load_inputs(self, inputs):
+        """Get data loader from inputs or create random data loader if no inputs provided."""
+        from polygraphy.comparator import DataLoader
+
+        # If no inputs are provided, use random inputs
+        data_loader = DataLoader(val_range={"": (-1, 1)})
+
+        if inputs is not None:
+            if isinstance(inputs, str):
+                if inputs.endswith(".json"):
+                    data_loader = self._load_inputs_from_json(inputs)
+                elif inputs.endswith(".npz"):
+                    data_loader = self._load_inputs_from_npz(inputs)
+                else:
+                    raise ValueError(
+                        f"Invalid input file: {inputs}. Supported input file types: .json (Polygraphy JSON format), "
+                        ".npz (Numpy)"
+                    )
+            elif isinstance(inputs, (dict, OrderedDict)):
+                data_loader = [inputs]
+            else:
+                raise ValueError(
+                    f"Invalid input type: {type(inputs)}. Supported input types: dict, OrderedDict, or a path to a "
+                    "JSON or NPZ file."
+                )
+            self._validate_inputs(data_loader)
+
+        return data_loader
+
+    def run(self, inputs=None):
+        """Run FP32 inference with provided or random inputs."""
+        import onnxruntime as ort
+        from polygraphy import constants
+        from polygraphy.backend.onnx import BytesFromOnnx
+        from polygraphy.backend.onnx import ModifyOutputs as ModifyOnnxOutputs
+        from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+        from polygraphy.comparator import Comparator
+
+        logger.info("Running ONNX Runtime to obtain reference outputs (this may take a while)...")
+        # Set ONNX Runtime log level to ERROR to suppress warnings
+        ort.set_default_logger_severity(3)
+
+        model_copy = copy.deepcopy(self.model)
+        modify_outputs = ModifyOnnxOutputs(model_copy, outputs=constants.MARK_ALL)
+        serialize_onnx = BytesFromOnnx(modify_outputs)
+        build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers)
+        runners = [OnnxrtRunner(build_onnxrt_session)]
+
+        # Comparator is used despite the fact that we are using ONNXRuntime
+        # because it provides the ability to generate random inputs using DataLoader
+        data_loader = self._load_inputs(inputs)
+
+        # Temporarily redirect stdout to suppress Comparator.run() output
+        stdout = sys.stdout
+        string_buffer = io.StringIO()
+        sys.stdout = string_buffer
+        try:
+            results = Comparator.run(runners, data_loader=data_loader)
+        finally:
+            # Capture the output before restoring stdout
+            captured_output = string_buffer.getvalue()
+            sys.stdout = stdout
+
+        if not results:
+            logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}")
+            raise Exception("ONNXRuntime failed to run, see logs for details")
+
+        # Get the output results
+        output_dict = OrderedDict(results[0][1][0])
+
+        # Include input data for completeness
+        input_data = next(iter(data_loader))
+
+        # Combine inputs and outputs in the returned dictionary
+        combined_dict = OrderedDict()
+        combined_dict.update(input_data)
+        combined_dict.update(output_dict)
+
+        return combined_dict

modelopt/onnx/autocast/utils.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Utility functions for AutoCast.
+
+This module provides common utility functions used across the AutoCast package.
+It includes functions for graph traversal, tensor type inference, model validation,
+and mapping setup between nodes, initializers, and value info. These utilities
+support the core functionality of model precision conversion.
+"""
+
+import onnx
+
+
+def setup_mappings(model: onnx.ModelProto) -> tuple[dict, dict, dict]:
+    """Setup and return mappings for model components.
+
+    Args:
+        model: ONNX model to create mappings for.
+
+    Returns:
+        Tuple containing:
+        - value_info_map: Mapping of names to value infos.
+        - initializer_map: Mapping of names to initializers.
+        - node_to_init_map: Mapping of node names to their initializer inputs.
+    """
+    value_info_map = {}
+    for container in (model.graph.value_info, model.graph.input, model.graph.output):
+        value_info_map.update((vi.name, vi) for vi in container)
+
+    initializer_map = {init.name: init for init in model.graph.initializer}
+
+    node_to_init_map = {
+        node.name: [
+            initializer_map[input_name]
+            for input_name in node.input
+            if input_name in initializer_map
+        ]
+        for node in model.graph.node
+    }
+
+    return value_info_map, initializer_map, node_to_init_map
+
+
+def get_consumer_nodes(model: onnx.ModelProto, tensor_name: str) -> list[onnx.NodeProto]:
+    """Get all consumer nodes for a given tensor name.
+
+    Args:
+        model: The ONNX model to search.
+        tensor_name: Name of the tensor to find consumers for.
+
+    Returns:
+        list[onnx.NodeProto]: List of nodes that consume the tensor.
+    """
+    return [n for n in model.graph.node if tensor_name in n.input]
+
+
+def get_producer_nodes(model: onnx.ModelProto, tensor_name: str) -> list[onnx.NodeProto]:
+    """Get all producer nodes for a given tensor name.
+
+    Args:
+        model: The ONNX model to search.
+        tensor_name: Name of the tensor to find producers for.
+
+    Returns:
+        list[onnx.NodeProto]: List of nodes that produce the tensor.
+    """
+    return [n for n in model.graph.node if tensor_name in n.output]
+
+
+def get_unique_consumer_node(model: onnx.ModelProto, tensor_name: str) -> onnx.NodeProto:
+    """Get a single consumer node and raise exception if there are multiple consumers.
+
+    Args:
+        model: The ONNX model to search.
+        tensor_name: Name of the tensor to find consumer for.
+
+    Returns:
+        onnx.NodeProto: The single consumer node.
+
+    Raises:
+        Exception: If there is not exactly one consumer node.
+    """
+    consumers = get_consumer_nodes(model, tensor_name)
+    if len(consumers) != 1:
+        raise Exception(f"Expected single consumer for {tensor_name}, found {len(consumers)}")
+    return consumers[0]
+
+
+def get_cast_to_type(cast_node: onnx.NodeProto) -> int:
+    """Get the target type from a Cast node.
+
+    Args:
+        cast_node: The Cast node to extract type from.
+
+    Returns:
+        int: The target type value from the Cast node's 'to' attribute.
+
+    Raises:
+        ValueError: If the Cast node does not have a 'to' attribute.
+    """
+    for attr in cast_node.attribute:
+        if attr.name == "to":
+            return attr.i
+    raise ValueError("Cast node does not have 'to' attribute")

modelopt/onnx/llm_export_utils/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Utilities for exporting LLM models to ONNX."""

modelopt/onnx/llm_export_utils/export_utils.py

@@ -0,0 +1,162 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Utilities for exporting LLM models to ONNX."""
+
+import json
+import os
+import time
+from enum import Enum
+
+import torch
+from transformers import DynamicCache
+
+
+class RopeType(Enum):
+    """Rope type enum."""
+
+    K_NONE = 0
+    K_ROPE_ROTATE_GPTJ = 1
+    K_ROPE_ROTATE_NEOX = 2
+    K_MROPE = 3
+
+
+class ModelLoader:
+    """A class to handle HuggingFace model loading and configuration."""
+
+    def __init__(self, torch_dir, config_path):
+        """Initialize the ModelLoader."""
+        self.config_path = config_path
+        self.torch_dir = torch_dir
+        self.model_type = self.get_model_type()
+        self.hf_model = None
+        self.rope_type = RopeType.K_ROPE_ROTATE_NEOX
+
+    def get_model_type(self):
+        """Get model type from config file."""
+        with open(self.config_path) as f:
+            return json.load(f).get("model_type")
+
+    def load_model(self):
+        """Load HuggingFace model based on model type."""
+        print(f"Loading HF model from {self.torch_dir} with model type {self.model_type}")
+        from transformers import AutoModelForCausalLM
+
+        self.hf_model = AutoModelForCausalLM.from_pretrained(
+            self.torch_dir, torch_dtype=torch.float16, trust_remote_code=True
+        )
+
+        return self.hf_model.eval().cuda()
+
+    def get_rope_type(self):
+        """Get rope type."""
+        return self.rope_type
+
+
+class WrapperModelForCausalLM(torch.nn.Module):
+    """Wrapper Model to ensure all models have the same I/O."""
+
+    def __init__(self, model):
+        """Initialize the WrapperModelForCausalLM."""
+        super().__init__()
+        try:
+            self.model = model.model
+        except Exception:
+            self.model = model
+        self.lm_head = model.lm_head
+        self.config = model.config
+
+    def forward(
+        self,
+        input_ids,
+        past_key_values,
+    ):
+        """Forward pass."""
+        past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+        outputs = self.model(input_ids=input_ids, past_key_values=past_key_values, use_cache=True)
+        hidden_states = outputs[0]
+        past_key_values = outputs.past_key_values.to_legacy_cache()
+        logits = self.lm_head(hidden_states)
+        return logits, past_key_values
+
+
+def llm_to_onnx(model, output_dir, extra_inputs={}, extra_dyn_axes={}):
+    """Export the WrapperModelForCausalLM to ONNX with fixed I/O names and shape definitions and save to `output_dir`.
+
+    Parameters:
+        model: torch.Module
+        output_dir: str, the output_dir of the original ONNX.
+        extra_inputs: dict, append additional inputs after kv_cache. Usually for VL models
+        extra_dyn_axes: dict. Usually for VL models
+    """
+    start_time = time.time()
+    config = model.config
+    num_layers = config.num_hidden_layers
+    num_attention_heads = config.num_attention_heads
+    num_key_value_heads = config.num_key_value_heads
+    hidden_size = config.hidden_size
+    hidden_size_per_layer = hidden_size // num_attention_heads
+
+    dummy_bs = 1
+    dummy_len = 10
+    dummy_input_ids = torch.randint(100, (dummy_bs, dummy_len), dtype=torch.int64).cuda()
+    input_names = ["input_ids"]
+    output_names = ["logits"]
+    dynamic_axes = {"input_ids": {0: "batch_size", 1: "seq_len"}}
+    dummy_kv_cache = ()
+    for i in range(num_layers):
+        dummy_k = torch.rand(
+            (dummy_bs, num_key_value_heads, dummy_len, hidden_size_per_layer), dtype=torch.float16
+        ).cuda()
+        dummy_v = torch.rand(
+            (dummy_bs, num_key_value_heads, dummy_len, hidden_size_per_layer), dtype=torch.float16
+        ).cuda()
+        dummy_kv_cache = (*dummy_kv_cache, (dummy_k, dummy_v))
+        input_names.extend([f"past_key_values.{i}.key", f"past_key_values.{i}.value"])
+        output_names.extend([f"present_key_values.{i}.key", f"present_key_values.{i}.value"])
+        input_dynamic_axes = {0: "batch_size", 2: "past_len"}
+        dynamic_axes[f"past_key_values.{i}.key"] = input_dynamic_axes
+        dynamic_axes[f"past_key_values.{i}.value"] = input_dynamic_axes
+
+    torch_to_onnx(
+        model,
+        (dummy_input_ids, {"past_key_values": dummy_kv_cache, **extra_inputs}),
+        output_dir,
+        "model.onnx",
+        input_names=input_names + list(extra_inputs.keys()),
+        output_names=output_names,
+        dynamic_axes=dynamic_axes | extra_dyn_axes,
+    )
+
+    end_time = time.time()
+    print(
+        f"Native ONNX Export from torch completed in {end_time - start_time}s. ONNX file is saved to {output_dir}."
+    )
+
+
+def torch_to_onnx(model, inputs, onnx_dir, onnx_name, input_names, output_names, dynamic_axes):
+    """Export the model to ONNX."""
+    os.makedirs(onnx_dir, exist_ok=True)
+    with torch.inference_mode():
+        torch.onnx.export(
+            model,
+            inputs,
+            f"{onnx_dir}/{onnx_name}",
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+            opset_version=19,
+            do_constant_folding=True,
+        )

modelopt/onnx/llm_export_utils/quantization_utils.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Quantization utilities for LLM models."""
+
+import time
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.utils.dataset_utils import get_dataset_dataloader
+
+
+def _quantize_model(model, quant_config, calib_dataloader=None):
+    """The calibration loop for the model can be setup using the modelopt API.
+
+    Example usage:
+    from modelopt.torch.utils.dataset_utils import create_forward_loop
+    model = ...  # Initialize the model
+    tokenizer = ...  # Initialize the tokenizer
+    quant_cfg = ...  # Setup quantization configuration
+    forward_loop = create_forward_loop(model=model, dataset_name="cnn_dailymail", tokenizer=tokenizer)
+    mtq.quantize(model, quant_cfg, forward_loop=forward_loop)
+    """
+
+    def calibrate_loop(model):
+        """Adjusts weights and scaling factors based on selected algorithms."""
+        for idx, data in enumerate(calib_dataloader):
+            if idx % 10 == 0:
+                print(f"Calibrating batch {idx}...")
+            if isinstance(data, dict):
+                data = {k: v.to(model.device) for k, v in data.items()}
+                model(**data)
+            else:
+                data = data.to(model.device)
+                model(data)
+
+    print("Starting quantization...")
+    start_time = time.time()
+    mtq.quantize(model, quant_config, forward_loop=calibrate_loop)
+    end_time = time.time()
+    print(f"Quantization finishes in {end_time - start_time}s.")
+
+    return model
+
+
+def get_quant_config(precision, lm_head_precision="fp16"):
+    """Get the quantization configuration."""
+    if precision == "fp8":
+        quant_cfg = mtq.FP8_DEFAULT_CFG
+
+    elif precision == "nvfp4":
+        quant_cfg = mtq.NVFP4_DEFAULT_CFG
+
+    elif precision == "int4_awq":
+        quant_cfg = mtq.INT4_AWQ_CFG
+
+    else:
+        raise ValueError(f"Unsupported precision: {precision}")
+
+    config_dict = quant_cfg["quant_cfg"]  # type: dict
+
+    if lm_head_precision == "fp8":
+        config_dict["*lm_head.input_quantizer"] = {"num_bits": (4, 3), "axis": None}
+        config_dict["*lm_head.weight_quantizer"] = {"num_bits": (4, 3), "axis": None}
+    elif lm_head_precision == "nvfp4":
+        config_dict["*lm_head.input_quantizer"] = {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        }
+        config_dict["*lm_head.weight_quantizer"] = {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        }
+    return quant_cfg
+
+
+def quantize(
+    model, tokenizer, precision, lm_head_precision="fp16", dataset_dir=None, calib_size=512
+):
+    """Quantize the PyTorch model to fp8 or int4_awq."""
+    assert precision in [
+        "fp8",
+        "int4_awq",
+        "nvfp4",
+    ], (
+        f"Only fp8(W8A8), int4_awq(W4A16), nvfp4(W4A4) is supported. You passed an unsupported precision: {precision}."
+    )
+
+    assert lm_head_precision in ["fp16"], (
+        f"Only fp16(unquantized) is supported for lm_head. You passed an unsupported precision: {lm_head_precision}."
+    )
+
+    if tokenizer.pad_token != "<unk>":  # nosec B105
+        tokenizer.pad_token = tokenizer.eos_token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if not dataset_dir:
+        dataset_dir = "cnn_dailymail"
+
+    batch_size = 1
+    data_loader = get_dataset_dataloader(
+        dataset_name=dataset_dir, tokenizer=tokenizer, batch_size=batch_size, num_samples=calib_size
+    )
+    quant_config = get_quant_config(precision, lm_head_precision)
+    quantized_model = _quantize_model(model, quant_config, data_loader)
+    mtq.print_quant_summary(quantized_model)
+    return quantized_model