nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/onnx/llm_export_utils/surgeon_utils.py

@@ -0,0 +1,120 @@
+# 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 to surgeon ONNX graph after export."""
+
+import re
+import time
+
+import onnx
+import onnx_graphsurgeon as gs
+import torch
+from onnx_graphsurgeon.ir.tensor import LazyValues
+
+
+def clear_inputs(node: gs.Node | gs.Tensor):
+    """Clear all inputs for a node or tensor in ONNX."""
+    for i in node.inputs:
+        i.outputs.clear()
+    node.inputs.clear()
+    return node
+
+
+def clear_outputs(node: gs.Node | gs.Tensor):
+    """Clear all outputs for a node or tensor in ONNX."""
+    for o in node.outputs:
+        o.inputs.clear()
+    node.outputs.clear()
+    return node
+
+
+def extract_layer_id(name: str):
+    """Extract layer id from certain ONNX layer name.
+
+    Parameters:
+        name: str
+            The name of ONNX layer. e.g. /model/layer.0/q_proj/...
+
+    Returns:
+        The layer id for the layer as int. In the example above, it returns 0
+    """
+    match = re.search(r"layers\.(\d+)", name)
+    if match:
+        return int(match.group(1))
+    raise Exception(f"{name} does not contain layer info!")
+
+
+def no_none_elements(elements: list):
+    """Check if all elements in the list are not None."""
+    return all(i is not None for i in elements)
+
+
+def fold_fp8_qdq_to_dq(graph: gs.Graph):
+    """Convert FP32/FP16 weights of the given ONNX model to FP8 weights.
+
+    Even though modelopt supports FP8 onnx export, the weights are represented in fp32 + QDQ.
+    The storage is therefore very bad. In this function,
+    Q nodes will get removed from the weights and have only DQ nodes with those converted FP8
+    weights in the output model.
+
+    Parameters:
+        graph: gs.Graph.
+
+    Returns:
+        gs.Graph with only DQ nodes for weights and same QDQ nodes for activations.
+    """
+    start_time = time.time()
+    print("Replacing all (fp32 weights + fp8 QDQ) with (fp8 weights + DQ)...")
+    # Fold constants is required since the scale is not constant yet.
+    graph.cleanup().toposort().fold_constants().cleanup()
+
+    for node in graph.nodes:
+        if node.op == "TRT_FP8QuantizeLinear":
+            # Should not remove input QDQ
+            if not isinstance(node.inputs[0], gs.Constant):
+                continue
+
+            weights = node.inputs[0]
+            scale = node.inputs[1]
+            torch_weights = torch.from_numpy(weights.values)
+            torch_scale = torch.from_numpy(scale.values)
+            quantizer_name = scale.name.rsplit("/", 1)[0]
+            dq_op = node.outputs[0].outputs[0]
+            assert dq_op.op == "TRT_FP8DequantizeLinear", (
+                f"QDQ does not occur in pairs. You reached {dq_op.op}"
+            )
+
+            # Replace it with Dequantize with FP8 weights. This is a WAR because numpy does not support fp8.
+            numpy_weights = (
+                (torch_weights / torch_scale).to(torch.float8_e4m3fn).view(torch.uint8).numpy()
+            )
+            tensor = onnx.TensorProto()
+            tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN
+            tensor.dims.extend(numpy_weights.shape)
+            tensor.raw_data = numpy_weights.tobytes()
+            values = LazyValues(tensor)
+            onnx_weights_fp8 = gs.Constant(quantizer_name + "/fp8_weights", values)
+
+            node.outputs.clear()
+            # DQ Op is separated out
+            dq_op.inputs[0] = onnx_weights_fp8
+            dq_op.op = "DequantizeLinear"
+            dq_op.outputs[0].dtype = dq_op.inputs[1].dtype
+
+    graph.cleanup().toposort()
+    end_time = time.time()
+    print(f"fp8 qdq replaced with only dq completed in {end_time - start_time}s.")
+
+    return graph

modelopt/onnx/logging_config.py

@@ -0,0 +1,77 @@
+# 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.
+
+"""Logging configuration for modelopt ONNX quantization."""
+
+import logging
+import os
+import sys
+
+# Create a parent logger for all ONNX components
+logger = logging.getLogger("modelopt.onnx")
+
+
+def configure_logging(level=logging.INFO, log_file=None):
+    """Configure logging for all ONNX components.
+
+    Args:
+        level: The logging level to use (default: logging.INFO)
+        log_file: Optional path to a log file. If provided, logs will be written to this file
+                 in addition to stdout (default: None)
+    """
+    # Set level for the parent logger and all child loggers
+    logger.setLevel(level)
+
+    # Remove any existing handlers to ensure clean configuration
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+
+    formatter = logging.Formatter("[modelopt][onnx] - %(levelname)s - %(message)s")
+
+    # Add file handler if log_file is specified
+    if log_file:
+        try:
+            # Create directory if it doesn't exist
+            log_dir = os.path.dirname(log_file)
+            if log_dir:
+                os.makedirs(log_dir, exist_ok=True)
+
+            file_handler = logging.FileHandler(log_file)
+            file_handler.setFormatter(formatter)
+            logger.addHandler(file_handler)
+            logger.info(f"Logging configured to write to file: {log_file}")
+        except Exception as e:
+            print(
+                f"[modelopt][onnx] - ERROR - Failed to setup file logging to {log_file}: {e!s}",
+                file=sys.stderr,
+            )
+            print("[modelopt][onnx] - INFO - Falling back to console logging.", file=sys.stderr)
+
+    # Setup handler to print log in stdout
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    # Prevent log messages from propagating to the root logger
+    logger.propagate = False
+
+    # Ensure all child loggers inherit the level setting
+    for name in logging.root.manager.loggerDict:
+        if name.startswith("modelopt.onnx"):
+            logging.getLogger(name).setLevel(level)
+
+
+# Configure with default settings if not already configured
+configure_logging()

modelopt/onnx/op_types.py

@@ -0,0 +1,300 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 to categorize onnx ops."""
+
+
+def is_unary_op(op_type: str):
+    """Returns whether the given op is a unary operator or not."""
+    return op_type in [
+        "Neg",
+        "Sqrt",
+        "Abs",
+        "Log",
+        "Exp",
+        "Not",
+        "Cast",
+        "Floor",
+        "Ceil",
+        "Round",
+        "Erf",
+        "Gelu",
+        "Sin",
+        "Cos",
+        "Atan",
+        "Sign",
+        "IsNaN",
+        "IsInf",
+        "Log",
+        "LeakyRelu",
+        # "Relu",
+        "Elu",
+        "Tanh",
+        "Sigmoid",
+        # "BatchNormalization",
+        "Softmax",
+        "Softplus",
+        "InstanceNormalization",
+        "CumSum",
+    ]
+
+
+def is_binary_op(op_type: str):
+    """Returns whether the given op is a binary operator or not."""
+    return op_type in [
+        "Add",
+        "Sub",
+        "Mul",
+        "Pow",
+        "Div",
+        "Min",
+        "Max",
+        "Greater",
+        "GreaterOrEqual",
+        "Less",
+        "LessOrEqual",
+        "Equal",
+        "BitwiseOr",
+        "BitwiseAnd",
+        "BitwiseXor",
+        "BitShift",
+    ]
+
+
+def is_fusible_reduction_op(op_type: str):
+    """Returns whether the given op type is of reduction category and fusible by the compiler."""
+    return op_type in [
+        "ReduceMax",
+        "ReduceMin",
+        "ReduceMean",
+        "ReduceProd",
+        "ReduceSum",
+        "TopK",  # Transformed to BottomK based on `largest` param
+    ]
+
+
+def is_fusible_scaling_op(op_type: str):
+    """Returns whether the given op type is of scaling category and fusible with input."""
+    return op_type in [
+        "SkipSimplifiedLayerNormalization",
+        "SimplifiedLayerNormalization",
+        "MatMul",
+        "Gemm",
+        "Mul",
+    ]
+
+
+def is_copy_op(op_type: str):
+    """Returns whether the given op is a copy operator or not."""
+    return op_type in [
+        "Flatten",
+        "Transpose",
+        "Concat",
+        "Split",
+        "Squeeze",
+        "Expand",
+        "ReverseSequence",
+        "Reshape",
+        "Tile",
+        "Gather",
+        "Slice",
+        "GatherElements",
+        "GatherND",
+        "ScatterElements",
+        "ScatterND",
+        "OneHot",
+    ]
+
+
+def is_linear_op(op_type: str):
+    """Returns whether the given op type is of Linear category or not."""
+    return op_type in ["Conv", "ConvTranspose", "Gemm", "MatMul"]
+
+
+def is_pointwise_or_elementwise_op(op_type: str):
+    """Returns whether the given op type is of Pointwise or Elementwise category or not.
+
+    This considers only the fusible types.
+    """
+    return is_unary_op(op_type) or is_binary_op(op_type)
+
+
+def is_pooling_or_window_op(op_type: str):
+    """Returns whether the given op type is of Pooling/Window category or not."""
+    return op_type in [
+        "AveragePool",
+        "GlobalAveragePool",
+        "MaxPool",
+        "GlobalMaxPool",
+        "GlobalLpPool",
+        "LpPool",
+        "MaxPoolGridSample",
+        "HammingWindow",
+        "BlackmanWindow",
+        "HannWindow",
+    ]
+
+
+def is_normalization_op(op_type: str):
+    """Returns whether the given op type is of Normalization category or not."""
+    return op_type in [
+        "BatchNormalization",
+        "InstanceNormalization",
+        "LRN",
+        "LpNormalization",
+        "GroupNormalization",
+        "LayerNormalization",
+    ]
+
+
+def is_conversion_op(op_type: str):
+    """Returns whether the given op type is of Conversion category or not."""
+    return op_type in ["Cast", "QuantizeLinear", "DequantizeLinear"]
+
+
+def is_non_reshape_copy_op(op_type: str):
+    """Returns whether the given op is a non-reshape copy op or not."""
+    return is_copy_op(op_type) and (op_type != "Reshape")
+
+
+def is_irregular_mem_access_op(op_type: str):
+    """Returns whether the given op type is of Irreggular mem access category or not."""
+    return op_type in [
+        "Gather",
+        "GatherElements",
+        "GatherND",
+        "MaxRoiPool",
+        "RoiAlign",
+        "ScatterND",
+        "ScatterElements",
+        "NonMaxSuppression",
+    ]
+
+
+def is_generator_op(op_type: str):
+    """Returns whether the given op type is of Generator category or not."""
+    return op_type in [
+        "Const",
+        "ConstOfShape",
+        "EyeLike",
+        "OneHot",
+        "Multinomial",
+        "RandomNormal",
+        "RandomUniform",
+        "Bernoulli",
+    ]
+
+
+def is_modifier_op(op_type: str):
+    """Returns whether the given op type is of Modifier category or not."""
+    return op_type in [
+        "Identity",
+        "Trilu",
+        "Expand",
+        "Pad",
+        "Dropout",
+        "TileDropout",
+        "Col2Im",
+        "MaxUnpool",
+    ]
+
+
+def is_sequence_op(op_type: str):
+    """Returns whether the given op type is of Sequence category or not."""
+    return op_type in [
+        "SequenceAt",
+        "SequenceConstruct",
+        "SequenceEmpty",
+        "SequenceErase",
+        "SequenceInsert",
+        "SequenceLength",
+    ]
+
+
+def is_selection_op(op_type: str):
+    """Returns whether the given op type is of Selection category or not."""
+    return op_type in ["Where", "Compress"]
+
+
+def is_control_flow_op(op_type: str):
+    """Returns whether the given op type is of Control Flow category or not."""
+    return op_type in ["If", "Loop"]
+
+
+def is_multiclass_op(op_type: str):
+    """Returns whether the given op type is of Multiclass category or not."""
+    return op_type in ["Einsum"]
+
+
+def is_recurrent_op(op_type: str):
+    """Returns whether the given op type is of Recurrent category or not."""
+    return op_type in ["LSTM", "RNN", "GRU"]
+
+
+def is_shape_op(op_type: str):
+    """Returns whether the given op type is of Shape category or not."""
+    return op_type in ["Shape", "Size"]
+
+
+def is_default_quantizable_op_by_ort(op_type: str):
+    """Returns if ORT quantizes the op type by default.
+
+    Note. Subject to change with different ORT versions.
+    Note. Users can use nodes_to_quantize and/or op_types_to_quantize arguments to quantize
+    non-default operations.
+    Reference: https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/quantization/registry.py
+    """
+    return op_type in [
+        "Conv",
+        "Gemm",
+        "ArgMax",
+        "Relu",
+        "Split",
+        "MaxPool",
+        "InstanceNormalization",
+        "Softmax",
+        "Where",
+        "Squeeze",
+        "GlobalAveragePool",
+        "Pad",
+        "Resize",
+        "ConvTranspose",
+        "Gather",
+        "Sigmoid",
+        "EmbedLayerNormalization",
+        "Reshape",
+        "Unsqueeze",
+        "Transpose",
+        "MatMul",
+        "Concat",
+        "Mul",
+        "Clip",
+        "Add",
+        "LeakyRelu",
+        "AveragePool",
+    ]
+
+
+def is_data_dependent_shape_op(op_type: str):
+    """Returns whether the op type has data-dependent shapes (DDS).
+
+    DDS ops have output shapes that are only determined at runtime based on its input data.
+    Source: https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#data-dependant-shape-dds-ops
+    """
+    return op_type in [
+        "NonMaxSuppression",
+        "NonZero",
+        "RoiAlign",
+    ]

modelopt/onnx/quantization/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Model optimization subpackage for onnx quantization."""
+
+# TODO: Use high-level quantize instead of quantize_int4
+from .int4 import quantize as quantize_int4
+from .quantize import quantize