Fix inc quantization bugs (microsoft#2174)

## Describe your changes
- Unify the logic when bits input is either int or enum.
- Use IntEnum by default since IntEnum is available since python 3.4
## Checklist before requesting a review
- [ ] 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.

## (Optional) Issue link

Author: xiaoyu-work

olive/cli/capture_onnx.py

@@ -5,6 +5,7 @@
 import argparse
 from argparse import ArgumentParser
 from copy import deepcopy
+from enum import IntEnum
 
 from olive.cli.base import (
     BaseOliveCLICommand,
@@ -15,10 +16,10 @@
     get_input_model_config,
     update_shared_cache_options,
 )
-from olive.common.utils import IntEnumBase, set_nested_dict_value
+from olive.common.utils import set_nested_dict_value
 
 
-class ModelBuilderAccuracyLevel(IntEnumBase):
+class ModelBuilderAccuracyLevel(IntEnum):
     fp32 = 1
     fp16 = 2
     bf16 = 3

olive/common/utils.py

@@ -25,24 +25,18 @@
 
 
 if sys.version_info >= (3, 11):
-    from enum import IntEnum, StrEnum
+    from enum import StrEnum
 
     class StrEnumBase(StrEnum):
         pass
 
-    class IntEnumBase(IntEnum):
-        pass
-
 else:
     from enum import Enum
 
     class StrEnumBase(str, Enum):
         def __str__(self) -> str:
             return self.value
 
-    class IntEnumBase(int, Enum):
-        pass
-
 
 def run_subprocess(cmd, env=None, cwd=None, check=False):
     logger.debug("Running command: %s", cmd)

olive/constants.py

@@ -2,8 +2,10 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
+from enum import IntEnum
+
 from olive.common.config_utils import CaseInsensitiveEnum
-from olive.common.utils import IntEnumBase, StrEnumBase
+from olive.common.utils import StrEnumBase
 
 MSFT_DOMAIN = "com.microsoft"
 
@@ -52,7 +54,7 @@ class Precision(StrEnumBase):
     BF16 = "bf16"
 
 
-class PrecisionBits(IntEnumBase):
+class PrecisionBits(IntEnum):
     BITS4 = 4
     BITS8 = 8
     BITS16 = 16
@@ -103,7 +105,7 @@ class OpType(StrEnumBase):
     Constant = "Constant"
 
 
-class AccuracyLevel(IntEnumBase):
+class AccuracyLevel(IntEnum):
     unset = 0
     fp32 = 1
     fp16 = 2

olive/passes/onnx/inc_quantization.py

@@ -440,10 +440,10 @@ def _set_tuning_config(self, run_config):
     def _set_woq_config(self, run_config):
         # set weight only quantization config for INC API
         weight_only_config = run_config["weight_only_config"]
-        bits = weight_only_config.get("bits", PrecisionBits.BITS4).value
+        bits = int(weight_only_config.get("bits") or PrecisionBits.BITS4)
         group_size = weight_only_config.get("group_size", 32)
         scheme = weight_only_config.get("scheme", "asym")
-        algo = (weight_only_config.get("algorithm") or QuantAlgorithm.RTN).value.upper()
+        algo = (weight_only_config.get("algorithm") or QuantAlgorithm.RTN.value).upper()
         return {"bits": bits, "group_size": group_size, "scheme": scheme, "algorithm": algo}
 
     def _run_for_config(

olive/passes/onnx/model_builder.py

@@ -7,14 +7,14 @@
 import copy
 import json
 import logging
+from enum import IntEnum
 from pathlib import Path
 from typing import Any, ClassVar, Union
 
 import onnx
 import transformers
 from packaging import version
 
-from olive.common.utils import IntEnumBase
 from olive.constants import Precision
 from olive.hardware.accelerator import AcceleratorSpec, Device
 from olive.hardware.constants import ExecutionProvider
@@ -33,14 +33,14 @@ class ModelBuilder(Pass):
     See https://github.com/microsoft/onnxruntime-genai
     """
 
-    class BlockSize(IntEnumBase):
+    class BlockSize(IntEnum):
         B16 = 16
         B32 = 32
         B64 = 64
         B128 = 128
         B256 = 256
 
-    class AccuracyLevel(IntEnumBase):
+    class AccuracyLevel(IntEnum):
         fp32 = 1
         fp16 = 2
         bf16 = 3
@@ -239,7 +239,7 @@ def _run_for_config(
 
         extra_args.update(
             {
-                key: value.value if isinstance(value, IntEnumBase) else value
+                key: value.value if isinstance(value, IntEnum) else value
                 for key, value in config.dict().items()
                 if value is not None and key not in {"precision", "metadata_only", "search", "extra_options"}
             }