diff --git a/src/coreai_opt/inspection/model_inspector.py b/src/coreai_opt/inspection/model_inspector.py index a4c703c..22e2176 100644 --- a/src/coreai_opt/inspection/model_inspector.py +++ b/src/coreai_opt/inspection/model_inspector.py @@ -18,7 +18,10 @@ from coreai_opt.base_model_compressor import _BaseModelCompressor from coreai_opt.palettization import KMeansPalettizer from coreai_opt.quantization import Quantizer -from coreai_opt.quantization.config.quantization_config import ExecutionMode +from coreai_opt.quantization.config.quantization_config import ( + ExecutionMode, + InvalidExecutionModeError, +) from ._eager_mode import parse_ops_for_eager as _parse_ops_for_eager from ._formatting import format_model_summary as _format_model_summary @@ -138,8 +141,7 @@ def _validate_args( raise TypeError(msg) if execution_mode not in (ExecutionMode.GRAPH, ExecutionMode.EAGER): - msg = f"Unknown execution_mode {execution_mode}. Expected 'graph' or 'eager'." - raise ValueError(msg) + raise InvalidExecutionModeError(execution_mode) if example_inputs is None and not ( isinstance(model, torch.fx.GraphModule) and execution_mode == ExecutionMode.GRAPH diff --git a/src/coreai_opt/quantization/__init__.py b/src/coreai_opt/quantization/__init__.py index 6fab1d1..535983c 100644 --- a/src/coreai_opt/quantization/__init__.py +++ b/src/coreai_opt/quantization/__init__.py @@ -15,6 +15,7 @@ # Import config after spec to avoid circular imports from .config import ( ExecutionMode, + InvalidExecutionModeError, ModuleQuantizerConfig, QuantizerConfig, ) @@ -22,6 +23,7 @@ __all__ = [ "ExecutionMode", + "InvalidExecutionModeError", "ModuleQuantizerConfig", "QuantizationSpec", "Quantizer", diff --git a/src/coreai_opt/quantization/config/__init__.py b/src/coreai_opt/quantization/config/__init__.py index 41b69cd..9d04613 100644 --- a/src/coreai_opt/quantization/config/__init__.py +++ b/src/coreai_opt/quantization/config/__init__.py @@ -7,6 +7,7 @@ from .quantization_config import ( ExecutionMode, + InvalidExecutionModeError, KVCacheQuantConfig, ModuleQuantizerConfig, OpQuantizerConfig, @@ -16,6 +17,7 @@ __all__ = [ "ExecutionMode", + "InvalidExecutionModeError", # Configuration classes "KVCacheQuantConfig", "ModuleQuantizerConfig", diff --git a/src/coreai_opt/quantization/config/quantization_config.py b/src/coreai_opt/quantization/config/quantization_config.py index c0c30f7..c2c83b9 100644 --- a/src/coreai_opt/quantization/config/quantization_config.py +++ b/src/coreai_opt/quantization/config/quantization_config.py @@ -158,6 +158,13 @@ class ExecutionMode(_StrEnum, metaclass=_DeprecatedMemberEnumMeta): """Deprecated. Use ``ExecutionMode.GRAPH`` instead.""" +class InvalidExecutionModeError(ValueError): + """Raised when an execution mode is not a recognized ExecutionMode value.""" + + def __init__(self, execution_mode: object) -> None: + super().__init__(f"Unknown execution_mode {execution_mode}. Expected 'graph' or 'eager'.") + + class OpQuantizerConfig(OpCompressionConfig[QuantizationSpec]): """ Configuration class for quantization at the operation level. diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py index a6de40f..a8dcaa2 100644 --- a/src/coreai_opt/quantization/quantizer.py +++ b/src/coreai_opt/quantization/quantizer.py @@ -30,6 +30,7 @@ from coreai_opt.quantization.base_quantizer import _BaseQuantizer from coreai_opt.quantization.config.quantization_config import ( ExecutionMode, + InvalidExecutionModeError, QATSchedule, QuantizerConfig, ) @@ -154,7 +155,7 @@ def __init__( elif execution_mode == ExecutionMode.EAGER: self._quantizer = _EagerQuantizer(model, config) else: - raise ValueError(f"Unsupported execution mode: {execution_mode}") + raise InvalidExecutionModeError(execution_mode) super().__init__(model, config) diff --git a/tests/test_inspection.py b/tests/test_inspection.py index 483f3e4..99dd80c 100644 --- a/tests/test_inspection.py +++ b/tests/test_inspection.py @@ -24,7 +24,7 @@ from coreai_opt.inspection._eager_mode import _EagerOpDiscoveryMode from coreai_opt.inspection.types import BoundaryEdge, InputEdge, OpInfo from coreai_opt.palettization import KMeansPalettizer -from coreai_opt.quantization import Quantizer +from coreai_opt.quantization import InvalidExecutionModeError, Quantizer from coreai_opt.quantization.config.quantization_config import ExecutionMode execution_modes = pytest.mark.parametrize( @@ -1138,9 +1138,9 @@ def test_eager_with_graph_module_raises_type_error(self) -> None: ModelInspector(gm, (torch.randn(1, 10),), execution_mode="eager") def test_invalid_execution_mode_raises(self) -> None: - """Verify ValueError for unrecognized execution mode.""" + """Verify InvalidExecutionModeError for unrecognized execution mode.""" model = nn.Linear(10, 5) - with pytest.raises(ValueError, match="Unknown execution_mode"): + with pytest.raises(InvalidExecutionModeError): ModelInspector(model, (torch.randn(1, 10),), execution_mode="invalid") def test_unsupported_compressor_raises(self) -> None: