Update matmul 4bit for ort>=1.17.0

@@ -594,6 +594,29 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
                 default_value=None,
                 description="List of node names to exclude from quantization.",
             ),
+            "accuracy_level": PassConfigParam(
+                type_=Any,
+                default_value=None,
+                description="Available from onnxruntime>=1.17.0 "
+                "The minimum accuracy level of input A, can be: 0(unset), 1(fp32), 2(fp16), 3(bf16), "
+                "or 4(int8) (default unset when 0 or None). It is used to control how input A is quantized or downcast "
+                "internally while doing computation, for example: 0 means input A will not be quantized "
+                "or downcast while doing computation. 4 means input A can be quantized with the same "
+                "block_size to int8 internally from type T1. "
+                "Refer to the MatMulNBits contrib op's 'accuracy_level' attribute for details "
+                "(https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits).",
+            ),
+            "algorithm": PassConfigParam(
+                type_=str,
+                default_value=None,
+                description="Available from onnxruntime>=1.17.0 "
+                "Algorithm config for quantization, can be 'DEFAULT', 'HQQ' or 'RTN', 'GPTQ' "
+                "(default DEFAULT when not set). For 4b quantize a model with RTN or GPTQ algorithm. "
+                "Please refer to https://github.com/intel/neural-compressor/blob/master/docs/source/quantization_weight_only.md "  # noqa: E501
+                "for more details on weight only quantization using Intel® Neural Compressor. "
+                "For HQQ, please refer to onnxruntime for more details: "
+                "https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/quantization/matmul_4bits_quantizer.py#L102C1-L126C25",
+            ),
         }
         config.update(get_external_data_config())
         return config
@@ -610,9 +633,25 @@ def _run_for_config(

         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)

-        quant = MatMul4BitsQuantizer(
-            model.load_model(), config["block_size"], config["is_symmetric"], config["nodes_to_exclude"]
-        )
+        if version.parse(OrtVersion) >= version.parse("1.17.0"):
+            from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig
+
+            weight_only_quant_config = WeightOnlyQuantConfig(algorithm=config["algorithm"])
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+                accuracy_level=config["accuracy_level"],
+                algo_config=weight_only_quant_config,
+            )
+        else:
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+            )
         quant.process()
         # topologically sort the graph at the end since previous optimizations may have broken it
         quant.model.topological_sort()

@@ -7,7 +7,7 @@

 from olive.hardware.accelerator import AcceleratorSpec
 from olive.passes.olive_pass import create_pass_from_dict
-from olive.passes.onnx import OnnxQuantization, OnnxStaticQuantization
+from olive.passes.onnx import OnnxMatMul4Quantizer, OnnxQuantization, OnnxStaticQuantization

 class DummyCalibrationDataReader(CalibrationDataReader):
@@ -86,3 +86,25 @@ def test_qnn_quantization(tmp_path):
     p = create_pass_from_dict(OnnxStaticQuantization, config, disable_search=True, accelerator_spec=accelerator_spec)
     out = p.run(input_model, None, tmp_path)
     assert out is not None
+
+
+@pytest.mark.skipif(
+    version.parse(OrtVersion) < version.parse("1.16.2"),
+    reason="matmul 4bit quantization is only supported in onnxruntime>=1.16.2",
+)
+def test_matmul_4bit_quantization(tmp_path):
+    input_model = get_onnx_model()
+    config = {
+        "block_size": 32,
+        "is_symmetric": True,
+        "nodes_to_exclude": [],
+        "accuracy_level": 4,
+        "algorithm": "DEFAULT",
+    }
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type="CPU",
+        execution_provider="CPUExecutionProvider",
+    )
+    p = create_pass_from_dict(OnnxMatMul4Quantizer, config, disable_search=True, accelerator_spec=accelerator_spec)
+    out = p.run(input_model, None, tmp_path)
+    assert out is not None


Update matmul 4bit for ort>=1.17.0


🆙 Update matmul 4bit for ort>=1.17.0

## Describe your changes
Update arguments of matmul 4bit for ort>=1.17.0
1. accuracy_level
2. algorithm

## Checklist before requesting a review
- [x] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [ ] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this change to be included in the release notes.
- [ ] Is this PR including examples changes? If yes, please remember to update [example documentation](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md) in a follow-up PR.

## (Optional) Issue link

Olive

Olive: Simplify ML Model Finetuning, Conversion, Quantization, and Optimization for CPUs, GPUs and NPUs.


@@ -594,6 +594,29 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
                 default_value=None,
                 description="List of node names to exclude from quantization.",
             ),
+            "accuracy_level": PassConfigParam(
+                type_=Any,

Should this be `int`?


@@ -594,6 +594,29 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
                 default_value=None,
                 description="List of node names to exclude from quantization.",
             ),
+            "accuracy_level": PassConfigParam(
+                type_=Any,
+                default_value=None,
+                description="Available from onnxruntime>=1.17.0 "
+                "The minimum accuracy level of input A, can be: 0(unset), 1(fp32), 2(fp16), 3(bf16), "
+                "or 4(int8) (default unset when 0 or None). It is used to control how input A is quantized or downcast "
+                "internally while doing computation, for example: 0 means input A will not be quantized "
+                "or downcast while doing computation. 4 means input A can be quantized with the same "
+                "block_size to int8 internally from type T1. "
+                "Refer to the MatMulNBits contrib op's 'accuracy_level' attribute for details "
+                "(https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits).",
+            ),
+            "algorithm": PassConfigParam(

What if the input str value is not desired? Will the api raise an error? Should we either add a value validation here, or having an enum class defined for validation for this value?


Any user requires this feature?


add validators

@@ -13,6 +13,7 @@

 from olive.cache import get_local_path_from_root
 from olive.common.config_utils import validate_config
+from olive.common.pydantic_v1 import validator
 from olive.common.utils import hash_string
 from olive.data.config import DataConfig
 from olive.exception import OlivePassError
@@ -595,7 +596,8 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
                 description="List of node names to exclude from quantization.",
             ),
             "accuracy_level": PassConfigParam(
-                type_=Any,
+                # TODO(trajep): to make it searchable
+                type_=int,
                 default_value=None,
                 description="Available from onnxruntime>=1.17.0 "
                 "The minimum accuracy level of input A, can be: 0(unset), 1(fp32), 2(fp16), 3(bf16), "
@@ -621,6 +623,13 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassCon
         config.update(get_external_data_config())
         return config

+    @classmethod
+    def _validators(cls) -> Dict[str, Callable]:
+        return {
+            "validate_accuracy_level": validator("accuracy_level", allow_reuse=True)(_validate_accuracy_level),
+            "validate_algorithm": validator("algorithm", allow_reuse=True)(_validate_algorithm),
+        }
+
     def _run_for_config(
         self, model: ONNXModelHandler, data_root: str, config: Dict[str, Any], output_model_path: str
     ) -> ONNXModelHandler:
@@ -659,3 +668,31 @@ def _run_for_config(

         # save the model to the output path and return the model
         return model_proto_to_olive_model(quant.model.model, output_model_path, config)
+
+
+def _validate_accuracy_level(v, values, field):
+    if not v:
+        return v
+
+    if v not in (0, 1, 2, 3, 4):
+        raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 0(unset), 1(fp32), 2(fp16), 3(bf16) or 4(int8)")
+
+    return v
+
+
+def _validate_algorithm(v, values, field):
+    if not v:
+        return v
+
+    if v not in ("DEFAULT", "HQQ", "RTN", "GPTQ"):
+        raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 'DEFAULT', 'HQQ', 'RTN', 'GPTQ'")
+
+    if v in ("RTN", "GPTQ"):
+        try:
+            import neural_compressor  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                f"OnnxMatMul4Quantizer {field.name}={v} requires neural_compressor package to be installed"
+            ) from None
+
+    return v

@@ -5,6 +5,7 @@
 from onnxruntime.quantization.calibrate import CalibrationDataReader
 from packaging import version

+from olive.common.pydantic_v1 import ValidationError
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx import OnnxMatMul4Quantizer, OnnxQuantization, OnnxStaticQuantization
@@ -108,3 +109,23 @@ def test_matmul_4bit_quantization(tmp_path):
     p = create_pass_from_dict(OnnxMatMul4Quantizer, config, disable_search=True, accelerator_spec=accelerator_spec)
     out = p.run(input_model, None, tmp_path)
     assert out is not None
+
+
+@pytest.mark.skipif(
+    version.parse(OrtVersion) < version.parse("1.16.2"),
+    reason="matmul 4bit quantization is only supported in onnxruntime>=1.16.2",
+)
+def test_invalid_config_for_matmul_4bits():
+    config = {
+        "block_size": 32,
+        "is_symmetric": True,
+        "nodes_to_exclude": [],
+        "accuracy_level": 5,
+        "algorithm": "TE",
+    }
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type="CPU",
+        execution_provider="CPUExecutionProvider",
+    )
+    with pytest.raises(ValidationError):
+        create_pass_from_dict(OnnxMatMul4Quantizer, config, disable_search=True, accelerator_spec=accelerator_spec)


add validators


Updated.


@@ -610,13 +642,57 @@

         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)

-        quant = MatMul4BitsQuantizer(
-            model.load_model(), config["block_size"], config["is_symmetric"], config["nodes_to_exclude"]
-        )
+        if version.parse(OrtVersion) >= version.parse("1.17.0"):
+            from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig
+
+            weight_only_quant_config = WeightOnlyQuantConfig(algorithm=config["algorithm"])
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+                accuracy_level=config["accuracy_level"],
+                algo_config=weight_only_quant_config,
+            )
+        else:
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+            )
         quant.process()
         # topologically sort the graph at the end since previous optimizations may have broken it
         quant.model.topological_sort()
         # quant.model._check_init is not needed since it's only meant for float8 quantization

         # save the model to the output path and return the model
         return model_proto_to_olive_model(quant.model.model, output_model_path, config)
+
+
+def _validate_accuracy_level(v, values, field):
+    if not v:
+        return v
+
+    if v not in (0, 1, 2, 3, 4):
+        raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 0(unset), 1(fp32), 2(fp16), 3(bf16) or 4(int8)")
+
+    return v
+
+
+def _validate_algorithm(v, values, field):
+    if not v:
+        return v
+
+    if v not in ("DEFAULT", "HQQ", "RTN", "GPTQ"):
+        raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 'DEFAULT', 'HQQ', 'RTN', 'GPTQ'")
+
+    if v in ("RTN", "GPTQ"):
+        try:
+            import neural_compressor  # noqa: F401

## PYLINT/W0611

Unused import neural_compressor (unused-import)
See [unused-import](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/unused-import.html).

[Show more details](https://github.com/microsoft/Olive/security/code-scanning/10537)


@@ -610,13 +642,57 @@ def _run_for_config(

         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)

-        quant = MatMul4BitsQuantizer(
-            model.load_model(), config["block_size"], config["is_symmetric"], config["nodes_to_exclude"]
-        )
+        if version.parse(OrtVersion) >= version.parse("1.17.0"):
+            from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig
+
+            weight_only_quant_config = WeightOnlyQuantConfig(algorithm=config["algorithm"])
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+                accuracy_level=config["accuracy_level"],
+                algo_config=weight_only_quant_config,
+            )
+        else:
+            quant = MatMul4BitsQuantizer(
+                model.load_model(),
+                block_size=config["block_size"],
+                is_symmetric=config["is_symmetric"],
+                nodes_to_exclude=config["nodes_to_exclude"],
+            )
         quant.process()
         # topologically sort the graph at the end since previous optimizations may have broken it
         quant.model.topological_sort()
         # quant.model._check_init is not needed since it's only meant for float8 quantization

         # save the model to the output path and return the model
         return model_proto_to_olive_model(quant.model.model, output_model_path, config)
+
+
+def _validate_accuracy_level(v, values, field):
+    if not v:
+        return v
+
+    if v not in (0, 1, 2, 3, 4):
+        raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 0(unset), 1(fp32), 2(fp16), 3(bf16) or 4(int8)")
+
+    return v
+
+
+def _validate_algorithm(v, values, field):

the validate_algorithm seems has additional logic to verify neural compressor installation. Better to make sure only algo valiation logic in this function and move the validation inc logic somewhere else.


move import checker into run

@@ -645,6 +645,15 @@ def _run_for_config(
         if version.parse(OrtVersion) >= version.parse("1.17.0"):
             from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig

+            if config["algorithm"] in ("RTN", "GPTQ"):
+                try:
+                    import neural_compressor  # noqa: F401 # pylint: disable=unused-import
+                except ImportError:
+                    raise ImportError(
+                        f"OnnxMatMul4Quantizer algorithm={config['algorithm']} "
+                        "requires neural_compressor package to be installed"
+                    ) from None
+
             weight_only_quant_config = WeightOnlyQuantConfig(algorithm=config["algorithm"])
             quant = MatMul4BitsQuantizer(
                 model.load_model(),
@@ -687,12 +696,4 @@ def _validate_algorithm(v, values, field):
     if v not in ("DEFAULT", "HQQ", "RTN", "GPTQ"):
         raise ValueError(f"OnnxMatMul4Quantizer {field.name} must be 'DEFAULT', 'HQQ', 'RTN', 'GPTQ'")

-    if v in ("RTN", "GPTQ"):
-        try:
-            import neural_compressor  # noqa: F401
-        except ImportError:
-            raise ImportError(
-                f"OnnxMatMul4Quantizer {field.name}={v} requires neural_compressor package to be installed"
-            ) from None
-
     return v


move import checker into run


@@ -610,13 +642,58 @@ def _run_for_config(

         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)

-        quant = MatMul4BitsQuantizer(
-            model.load_model(), config["block_size"], config["is_symmetric"], config["nodes_to_exclude"]
-        )
+        if version.parse(OrtVersion) >= version.parse("1.17.0"):
+            from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig
+
+            if config["algorithm"] in ("RTN", "GPTQ"):

I suppose the underlying ORT will do raise the ImportError and we might need not this check?


saw the changes have already been made but just wanted to agree to Mikes point. Version checks and import checkers are runtime decisions to be run on the pass host. doing it during config validation or search point validation runs it on the engine's machine which is not necessarily the pass host.


Yes, the import error will be raised in ort as well. But no extra info to tell which config needs inc package.
Maybe we can keep it here to provide clear description once user choose the wrong config.


remove import error check

@@ -645,15 +645,6 @@ def _run_for_config(
         if version.parse(OrtVersion) >= version.parse("1.17.0"):
             from onnxruntime.quantization.matmul_4bits_quantizer import WeightOnlyQuantConfig

-            if config["algorithm"] in ("RTN", "GPTQ"):
-                try:
-                    import neural_compressor  # noqa: F401 # pylint: disable=unused-import
-                except ImportError:
-                    raise ImportError(
-                        f"OnnxMatMul4Quantizer algorithm={config['algorithm']} "
-                        "requires neural_compressor package to be installed"
-                    ) from None
-
             weight_only_quant_config = WeightOnlyQuantConfig(algorithm=config["algorithm"])
             quant = MatMul4BitsQuantizer(
                 model.load_model(),


remove import error check