Migrate bnb4, io_datatype_converter, dynamic_to_fixed_shape, and static_llm passes to onnx_ir

olive/constants.py

@@ -96,6 +96,7 @@ class OpType(StrEnumBase):
     Gather = "Gather"
     GatherBlockQuantized = "GatherBlockQuantized"
     MatMulNBits = "MatMulNBits"
+    MatMulBnb4 = "MatMulBnb4"
     MatMul = "MatMul"
     QuickGelu = "QuickGelu"
     Sigmoid = "Sigmoid"
@@ -114,6 +115,7 @@ class OpType(StrEnumBase):
     PackedMultiHeadAttention = "PackedMultiHeadAttention"
     MultiHeadAttention = "MultiHeadAttention"
     Loop = "Loop"
+    Cast = "Cast"
 
 
 class AccuracyLevel(IntEnum):

olive/passes/onnx/bnb_quantization.py

@@ -7,19 +7,30 @@
 from pathlib import Path
 from typing import Optional
 
-import onnx
-from packaging import version
+import numpy as np
+import numpy.typing as npt
+import onnx_ir as ir
 
-from olive.constants import Precision
+from olive.constants import MSFT_DOMAIN, OpType, Precision
 from olive.hardware import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
-from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.common import get_external_data_config, ir_model_to_olive_model
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
 
+# 4-bit quantization types, must be consistent with the native ORT MatMulBnb4 kernel
+# (Bnb_DataType_t defined in blockwise_quant_block_bnb4.h).
+_BNB4_QUANT_TYPES = {
+    # 4b floating point with bias of 3
+    "fp4": 0,
+    # 4b NormalFloat
+    "nf4": 1,
+}
+_BNB4_BLOCK_SIZE = 64
+
 
 class OnnxBnb4Quantization(Pass):
     """Quantize MatMul nodes in ONNX model using 4bit FP4/NF4 quantization."""
@@ -63,14 +74,6 @@ def validate_config(
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        from onnxruntime import __version__ as OrtVersion
-
-        assert version.parse(OrtVersion) >= version.parse("1.16.2"), (
-            "MatMulBnb4Quantizer is only supported in onnxruntime >= 1.16.2"
-        )
-
-        from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
-
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
         precision = config.precision.value if config.precision else None
@@ -93,50 +96,157 @@ def _run_for_config(
                     logger.info(
                         "bnb_4bit_use_double_quant is set to True but double quantization is not supported. Ignoring."
                     )
-        assert precision in {"fp4", "nf4"}, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
-        quant_type_enum = getattr(MatMulBnb4Quantizer, precision.upper())
+        assert precision in _BNB4_QUANT_TYPES, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
+        quant_type = _BNB4_QUANT_TYPES[precision]
 
         # load the model
-        onnx_model = model.load_model()
+        ir_model = model.load_ir_model()
+        ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
 
         # get nodes to exclude from quantization
         nodes_to_exclude = config.nodes_to_exclude or []
 
         # find all MatMul nodes in the graph
-        matmul_nodes = self._find_matmul_nodes(onnx_model.graph)
+        matmul_nodes = self._find_matmul_nodes(ir_model.graph)
         # filter based on quantized_modules
         quantized_modules = set(quantized_modules or [])
-        nodes_to_exclude = nodes_to_exclude + [
+        nodes_to_exclude = set(nodes_to_exclude) | {
             node
             for node in matmul_nodes
             if quantized_modules and not any(re.match(f".*[./]{key}[./]MatMul$", node) for key in quantized_modules)
-        ]
+        }
 
         # quantize the model
-        quantizer = MatMulBnb4Quantizer(
-            onnx_model, quant_type=quant_type_enum, block_size=64, nodes_to_exclude=nodes_to_exclude
-        )
-        quantizer.process()
-        # topologically sort the graph at the end since previous optimizations may have broken it
-        quantizer.model.topological_sort()
+        self._quantize_model(ir_model, quant_type, nodes_to_exclude)
 
         # save the model to the output path and return the model
-        return model_proto_to_olive_model(onnx_model, output_model_path, config)
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
+
+    def _quantize_model(self, ir_model: ir.Model, quant_type: int, nodes_to_exclude: set[str]) -> None:
+        """Replace eligible MatMul nodes with MatMulBnb4 nodes carrying 4-bit quantized weights."""
+        ir_model.graph.sort()
+        for node in ir_model.graph.all_nodes():
+            if node.op_type != str(OpType.MatMul):
+                continue
+
+            if node.name in nodes_to_exclude:
+                logger.debug("exclude to quantize %s as specified by nodes_to_exclude...", node.name)
+                continue
+
+            quantized_node = self._quantize_matmul(node, quant_type)
+            if quantized_node is node:
+                # nothing changed (non-const or non-2D weight)
+                continue
+
+            weight_graph = node.inputs[1].graph
+            for input_value in quantized_node.inputs:
+                if input_value is not None and input_value.const_value is not None:
+                    weight_graph.register_initializer(input_value)
+
+            ir.convenience.replace_nodes_and_values(
+                node.graph, node, [node], [quantized_node], node.outputs, quantized_node.outputs
+            )
+
+        self._remove_unused_initializers(ir_model)
+
+    def _quantize_matmul(self, node: ir.Node, quant_type: int) -> ir.Node:
+        """Quantize weight B of a MatMul node to 4-bit and return the new MatMulBnb4 node.
+
+        Returns the original node unchanged if the weight is not a 2D constant.
+        """
+        logger.debug("start to quantize %s ...", node.name)
+
+        node_initializer = node.inputs[1]
+        if node_initializer is None or node_initializer.const_value is None:
+            logger.debug("MatMul doesn't have const weight. Skip to quantize")
+            return node
+
+        b_ndarray = node_initializer.const_value.numpy()
+        if len(b_ndarray.shape) != 2:
+            logger.debug("MatMul weight is not 2D. Skip to quantize")
+            return node
+
+        packed, absmax = self._bnb4_block_quant(b_ndarray, quant_type)
+
+        b_quant = ir.Value(name=node_initializer.name + "_Bnb4", const_value=ir.tensor(packed))
+        absmax_value = ir.Value(name=node_initializer.name + "_absmax", const_value=ir.tensor(absmax))
+
+        rows, cols = b_ndarray.shape
+        kwargs = {
+            "K": rows,
+            "N": cols,
+            "block_size": _BNB4_BLOCK_SIZE,
+            "quant_type": quant_type,
+        }
+
+        self._rename_output_unless_graph_output(node)
+
+        logger.debug("complete quantization of %s ...", node.name)
+
+        return ir.node(
+            domain=MSFT_DOMAIN,
+            op_type=str(OpType.MatMulBnb4),
+            inputs=[node.inputs[0], b_quant, absmax_value],
+            name=node.name + "_Bnb4" if node.name else "",
+            attributes=kwargs,
+        )
+
+    @staticmethod
+    def _bnb4_block_quant(fpweight: npt.ArrayLike, quant_type: int) -> tuple[np.ndarray, np.ndarray]:
+        """4b quantize fp32/fp16 weight using the native ORT bnb4 kernel."""
+        from onnxruntime.capi._pybind_state import quantize_matmul_bnb4
+
+        if len(fpweight.shape) != 2:
+            raise ValueError("Current bnb4 block quantization only supports 2D tensors!")
+        # need to copy since the transposed weight still has the original memory layout
+        # Linear4bit quantizes its weight data which is the transposed weight
+        fpweight_t = fpweight.transpose().copy()
+
+        rows, cols = fpweight.shape
+        numel = rows * cols
+        block_size = _BNB4_BLOCK_SIZE
+        num_blocks = (numel + block_size - 1) // block_size
+        quantized_numel = (numel + 1) // 2
+
+        packed = np.zeros(quantized_numel, dtype="uint8")
+        absmax = np.zeros(num_blocks, dtype=fpweight.dtype)
+        # block wise quantization, fpweight_t is flattened and divided into blocks
+        quantize_matmul_bnb4(packed, fpweight_t, absmax, block_size, quant_type, cols, rows)
+
+        return packed, absmax
+
+    @staticmethod
+    def _rename_output_unless_graph_output(node: ir.Node) -> None:
+        """Append a `_Bnb4` suffix to the node's output name for renaming after quantization.
+
+        Skips values that are graph outputs so external consumers relying on those names are
+        not broken. Internal tensors are safe to rename because their consumers are rewired by
+        replace_nodes_and_values.
+        """
+        graph = node.graph
+        is_graph_output = graph is not None and node.outputs[0] in graph.outputs
+        if not is_graph_output:
+            node.outputs[0].name = node.outputs[0].name + "_Bnb4"
+
+    @staticmethod
+    def _remove_unused_initializers(ir_model: ir.Model) -> None:
+        """Remove initializers that are no longer referenced by any node after quantization."""
+        used_names: set[str] = set()
+        for node in ir_model.graph.all_nodes():
+            for inp in node.inputs:
+                if inp is not None and inp.name:
+                    used_names.add(inp.name)
+        for out in ir_model.graph.outputs:
+            if out is not None and out.name:
+                used_names.add(out.name)
+
+        unused = [name for name in ir_model.graph.initializers if name not in used_names]
+        for name in unused:
+            del ir_model.graph.initializers[name]
+        if unused:
+            logger.debug("Removed %d unused initializers after quantization.", len(unused))
 
     @classmethod
-    def _find_matmul_nodes(cls, graph: onnx.GraphProto) -> list[str]:
-        """Find all MatMul nodes in the graph and return their names."""
-        matmul_nodes = []
-        for node in graph.node:
-            for attr in node.attribute:
-                if attr.type == onnx.AttributeProto.GRAPH:
-                    # recursive call to take care of sub-graph
-                    matmul_nodes += cls._find_matmul_nodes(attr.g)
-                elif attr.type == onnx.AttributeProto.GRAPHS:
-                    for subgraph in attr.graphs:
-                        # recursive call to take care of sub-graph
-                        matmul_nodes += cls._find_matmul_nodes(subgraph)
-            if node.op_type == "MatMul":
-                matmul_nodes.append(node.name)
-
-        return matmul_nodes
+    def _find_matmul_nodes(cls, graph: ir.Graph) -> list[str]:
+        """Find all MatMul nodes in the graph (including subgraphs) and return their names."""
+        return [node.name for node in graph.all_nodes() if node.op_type == str(OpType.MatMul)]

olive/passes/onnx/common.py

@@ -103,6 +103,19 @@ def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelPro
     return model
 
 
+def add_version_metadata_to_ir_model(model: ir.Model) -> ir.Model:
+    try:
+        import olive
+
+        olive_version = getattr(olive, "__version__", "unknown")
+    except Exception:
+        olive_version = "unknown"
+
+    model.metadata_props["olive_version"] = olive_version
+
+    return model
+
+
 def model_proto_to_file(
     model: onnx.ModelProto,
     output_path: Union[str, Path],
@@ -578,57 +591,6 @@ def model_has_adapters(model_path: Union[str, Path], adapter_type: AdapterType =
     )
 
 
-def _fix_output_shapes(model_proto: onnx.ModelProto):
-    """Run shape inference on the model and update the output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
-
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-
-    for idx, o in enumerate(model_proto.graph.output):
-        if not is_fixed_size_tensor(o):
-            new_o = inferred_proto.graph.output[idx]
-            if is_fixed_size_tensor(new_o):
-                o.type.tensor_type.shape.CopyFrom(new_o.type.tensor_type.shape)
-
-
-def fix_dim_params(model_proto: onnx.ModelProto, dim_params: list[str], dim_values: list[int]):
-    """Fix the dimension parameters in the model.
-
-    :param dim_params: The dimension parameters to fix.
-    :param dim_values: The values to set for the dimension parameters.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_dim_param_fixed
-
-    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
-    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
-
-    for param, value in zip(dim_params, dim_values):
-        make_dim_param_fixed(model_proto.graph, param, value)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
-def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input_shapes: list[list[int]]):
-    """Fix the input shapes in the model.
-
-    :param input_names: The input names to fix.
-    :param input_shapes: The shapes to set for the inputs.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_input_shape_fixed
-
-    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
-    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
-
-    for name, shape in zip(input_names, input_shapes):
-        make_input_shape_fixed(model_proto.graph, name, shape)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,

olive/passes/onnx/compose.py

@@ -112,10 +112,11 @@ def _get_composed_model(
         :return: Composed ONNX model.
         """
 
-        def shape_list(value: ir.Value):
-            if value.shape is None:
-                return None
-            return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
+        def merge_value_shapes(existing: ir.Value, consumer_input: ir.Value, name: str):
+            try:
+                existing.merge_shapes(consumer_input.shape)
+            except ValueError as e:
+                raise AssertionError(f"Input shape mismatch: {name}") from e
 
         ir_models = []
         for path in onnx_model_paths:
@@ -170,7 +171,7 @@ def get_value(name: str) -> ir.Value:
                 if name in composed_input_names or name in produced_names:
                     # already a graph input or an internal connection from a previous model
                     existing = composed_values[name]
-                    assert shape_list(inp) == shape_list(existing), f"Input shape mismatch: {name}"
+                    merge_value_shapes(existing, inp, name)
                     assert inp.dtype == existing.dtype, f"Input dtype mismatch: {name}"
                     continue