diff --git a/changelog.d/42.fixed b/changelog.d/42.fixed new file mode 100644 index 0000000..626e097 --- /dev/null +++ b/changelog.d/42.fixed @@ -0,0 +1 @@ +Reject per-channel activation quantization on CoreML export diff --git a/src/coreai_opt/_utils/export_utils.py b/src/coreai_opt/_utils/export_utils.py index 2ffbd19..1b6cc62 100644 --- a/src/coreai_opt/_utils/export_utils.py +++ b/src/coreai_opt/_utils/export_utils.py @@ -9,7 +9,9 @@ import torch from coreai_opt._utils.torch_utils import is_tensor_on_cpu -from coreai_opt.common import ExportBackend +from coreai_opt.common import CoreMLExportError, ExportBackend +from coreai_opt.config.spec import CompressionTargetTensor +from coreai_opt.quantization.spec.granularity import PerTensorGranularity, QuantizationGranularity COREML_SUPPORTED_WEIGHT_DTYPES: frozenset[torch.dtype] = frozenset( { @@ -34,6 +36,47 @@ } ) +COREML_SUPPORTED_ACTIVATION_GRANULARITIES: frozenset[type[QuantizationGranularity]] = frozenset( + {PerTensorGranularity} +) + +_COREML_SUPPORTED_DTYPES_BY_TARGET: dict[CompressionTargetTensor, frozenset[torch.dtype]] = { + CompressionTargetTensor.WEIGHT: COREML_SUPPORTED_WEIGHT_DTYPES, + CompressionTargetTensor.ACTIVATION: COREML_SUPPORTED_ACTIVATION_DTYPES, + CompressionTargetTensor.LUT: COREML_SUPPORTED_LUT_DTYPES, +} + + +def validate_coreml_compatibility( + target: CompressionTargetTensor, + dtype: torch.dtype, + context: str, + granularity: QuantizationGranularity | None = None, +) -> None: + """Raise CoreMLExportError if this weight/activation/LUT config isn't CoreML-exportable. + + Centralizes every reason CoreML export can reject a quantization config, so + new restrictions are added here once rather than at each call site. + + Args: + target (CompressionTargetTensor): Which tensor category is being checked. + dtype (torch.dtype): The quantization dtype to validate. + context (str): Human-readable description of what's being checked, used + in the error message (e.g. "weight 'conv.weight' of module 'conv'"). + granularity (QuantizationGranularity | None): The quantization + granularity, if applicable. Only checked for ACTIVATION — CoreML + only supports per-tensor activation quantization. + + Raises: + CoreMLExportError: If the dtype or granularity isn't supported. + """ + if dtype not in _COREML_SUPPORTED_DTYPES_BY_TARGET[target]: + raise CoreMLExportError.from_dtype(dtype, context) + if target == CompressionTargetTensor.ACTIVATION and not isinstance( + granularity, tuple(COREML_SUPPORTED_ACTIVATION_GRANULARITIES) + ): + raise CoreMLExportError.from_config(granularity, context) + def validate_mmap_backend_and_device( model: torch.nn.Module, diff --git a/src/coreai_opt/common.py b/src/coreai_opt/common.py index 24c8983..5e9aa5e 100644 --- a/src/coreai_opt/common.py +++ b/src/coreai_opt/common.py @@ -163,8 +163,15 @@ class ExportBackend(_StrEnum, metaclass=_DeprecatedMemberEnumMeta): class CoreMLExportError(ValueError): """Raised when a model cannot be exported to the CoreML backend.""" - def __init__(self, dtype: Any, context: str) -> None: - super().__init__( - f"CoreML export does not support dtype {dtype} on {context}. " - f"Use backend=ExportBackend.CoreAI instead." - ) + def __init__(self, message: str) -> None: + super().__init__(f"{message} Use backend=ExportBackend.CoreAI instead.") + + @classmethod + def from_dtype(cls, dtype: Any, context: str) -> CoreMLExportError: + """Build the error for an unsupported weight/activation/LUT dtype.""" + return cls(f"CoreML export does not support dtype {dtype} on {context}.") + + @classmethod + def from_config(cls, config: object, context: str) -> CoreMLExportError: + """Build the error for an unsupported quantization config attribute (e.g. granularity).""" + return cls(f"CoreML export does not support {type(config).__name__} on {context}.") diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 678188d..49512d0 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -12,16 +12,17 @@ import torch.nn.utils.parametrize as P from coreai_opt._utils.export_utils import ( - COREML_SUPPORTED_LUT_DTYPES, clear_parametrization_original as _clear_parametrization_original, prepare_mmap_dir as _prepare_mmap_dir, + validate_coreml_compatibility, ) from coreai_opt._utils.import_utils import lazy_import_coreai_torch from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata from coreai_opt._utils.torch_utils import ( mmap_module_state_dict as _mmap_module_state_dict, ) -from coreai_opt.common import CoreMLExportError, ExportBackend +from coreai_opt.common import ExportBackend +from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.palettization.spec.fake_palettize import ( _FakePalettizeImplBase, ) @@ -430,11 +431,11 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module: for param_name, parametrizations in module.parametrizations.items(): _, fake_palett_mod = _find_fake_palett_parametrization(parametrizations) if fake_palett_mod is not None and fake_palett_mod.lut_qspec is not None: - if fake_palett_mod.lut_qspec.dtype not in COREML_SUPPORTED_LUT_DTYPES: - raise CoreMLExportError( - fake_palett_mod.lut_qspec.dtype, - f"LUT of parameter '{param_name}' of module '{module_name}'", - ) + validate_coreml_compatibility( + CompressionTargetTensor.LUT, + fake_palett_mod.lut_qspec.dtype, + f"LUT of parameter '{param_name}' of module '{module_name}'", + ) _process_weight_palettization(model, backend=ExportBackend.CoreML) diff --git a/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py b/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py index 9670516..f0679a6 100644 --- a/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py +++ b/src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py @@ -17,15 +17,11 @@ import torch.nn.utils.parametrize as P from torch import nn -from coreai_opt._utils.export_utils import ( - COREML_SUPPORTED_ACTIVATION_DTYPES, - COREML_SUPPORTED_WEIGHT_DTYPES, -) +from coreai_opt._utils.export_utils import validate_coreml_compatibility from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata from coreai_opt._utils.torch_utils import ( get_parent_module_and_attr_name as _get_parent_module_and_attr_name, ) -from coreai_opt.common import CoreMLExportError from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization._export_utils import ( convert_dtype_for_torch_quantize as _convert_dtype_for_torch_quantize, @@ -35,7 +31,6 @@ validate_qformulation_for_mil_export as _validate_qformulation_for_mil_export, ) from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase -from coreai_opt.quantization.spec.granularity import PerBlockGranularity def _process_weight_quantization( @@ -96,22 +91,16 @@ def _process_activation_quantization( ) -> None: """Process activation quantization by replacing with Sequential quantize/dequantize. - Supports both per-tensor and per-channel quantization based on the - granularity axis of the fake quantization module. + By the time this runs, CoreML export compatibility has already been + validated, so only per-tensor activation granularity ever reaches here. Args: parent_module: The parent module containing the fake quantizer attr_name: The attribute name of the fake quantizer in the parent fake_quant_mod: The fake quantization module to replace - Raises: - ValueError: If the granularity is per-block (not supported for MIL - activation export). - """ _validate_qformulation_for_mil_export(fake_quant_mod) - if isinstance(fake_quant_mod.granularity, PerBlockGranularity): - raise ValueError("MIL export does not support per-block granularity for activations.") scale, zero_point, _ = _extract_quantization_params(fake_quant_mod) converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize( @@ -151,24 +140,24 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module: """ processed_fq_ids: set[int] = set() - # CoreML does not support certain dtypes (FP4, FP8, INT2, UINT2). Fail fast - # before mutating the model below. + # Fail fast if model is not coreml-exportable for module_name, module in model.named_modules(): if P.is_parametrized(module): for param_name, parametrizations in module.parametrizations.items(): for p in parametrizations: - if ( - _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT) - and p.dtype not in COREML_SUPPORTED_WEIGHT_DTYPES - ): - raise CoreMLExportError( - p.dtype, f"weight '{param_name}' of module '{module_name}'" + if _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT): + validate_coreml_compatibility( + CompressionTargetTensor.WEIGHT, + p.dtype, + f"weight '{param_name}' of module '{module_name}'", ) - if ( - _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION) - and module.dtype not in COREML_SUPPORTED_ACTIVATION_DTYPES - ): - raise CoreMLExportError(module.dtype, f"activation quantizer of module '{module_name}'") + if _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION): + validate_coreml_compatibility( + CompressionTargetTensor.ACTIVATION, + module.dtype, + f"activation quantizer of module '{module_name}'", + module.granularity, + ) for name, module in list(model.named_modules()): # Handle weight quantization parametrizations diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index 60f1545..d7ef4ba 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -17,10 +17,7 @@ import torch from torch.fx import Node -from coreai_opt._utils.export_utils import ( - COREML_SUPPORTED_ACTIVATION_DTYPES, - COREML_SUPPORTED_WEIGHT_DTYPES, -) +from coreai_opt._utils.export_utils import validate_coreml_compatibility from coreai_opt._utils.fx_utils import get_node_type as _get_node_type from coreai_opt._utils.import_utils import lazy_import_coreai_torch from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata @@ -28,7 +25,6 @@ is_float4_dtype as _is_float4_dtype, sanitize_module_name as _sanitize_module_name, ) -from coreai_opt.common import CoreMLExportError from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization._export_utils import ( canonicalize_qparam_shape as _canonicalize_qparam_shape, @@ -45,7 +41,6 @@ resolve_attr, ) from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase -from coreai_opt.quantization.spec.granularity import PerBlockGranularity logger = logging.getLogger(__name__) @@ -503,23 +498,17 @@ def _process_mil_activation_quantization( Converts activation FakeQuantize modules to Sequential modules containing quantize and dequantize operations that CoreMLTools can understand. - Supports both per-tensor (quantize_per_tensor) and per-channel - (quantize_per_channel) quantization based on the granularity axis. - Buffers are stored inside the modules. + Buffers are stored inside the modules. By the time this runs, CoreML + export compatibility has already been validated, so only per-tensor + activation granularity ever reaches here. Args: model: The graph module being modified fake_quant_node: The fake quantization node to process fake_quant_mod: The fake quantization module - Raises: - ValueError: If the granularity is per-block (not supported for MIL - activation export). - """ _validate_qformulation_for_mil_export(fake_quant_mod) - if isinstance(fake_quant_mod.granularity, PerBlockGranularity): - raise ValueError("MIL export does not support per-block granularity for activations.") scale, zero_point, _ = _extract_quantization_params(fake_quant_mod) converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize( @@ -599,22 +588,22 @@ def prepare_for_mil_export(model: torch.fx.GraphModule) -> torch.fx.GraphModule: msg = "Model contains no fake quantization nodes" raise ValueError(msg) - # CoreML does not support certain dtypes. Fail fast before mutating the - # graph below. + # Fail fast if model is not coreml-exportable for fake_quant_node, fake_quant_mod in fake_quant_nodes: node_id = str(fake_quant_node.target) if _is_weight_fake_quant(fake_quant_node, fake_quant_mod): - if fake_quant_mod.dtype not in COREML_SUPPORTED_WEIGHT_DTYPES: - raise CoreMLExportError( - fake_quant_mod.dtype, - f"weight quantizer '{node_id}'", - ) + validate_coreml_compatibility( + CompressionTargetTensor.WEIGHT, + fake_quant_mod.dtype, + f"weight quantizer '{node_id}'", + ) else: - if fake_quant_mod.dtype not in COREML_SUPPORTED_ACTIVATION_DTYPES: - raise CoreMLExportError( - fake_quant_mod.dtype, - f"activation quantizer '{node_id}'", - ) + validate_coreml_compatibility( + CompressionTargetTensor.ACTIVATION, + fake_quant_mod.dtype, + f"activation quantizer '{node_id}'", + fake_quant_mod.granularity, + ) # Process all fake quantization nodes for fake_quant_node, fake_quant_mod in fake_quant_nodes: diff --git a/tests/export/export_utils.py b/tests/export/export_utils.py index 9553496..d622324 100644 --- a/tests/export/export_utils.py +++ b/tests/export/export_utils.py @@ -29,9 +29,10 @@ if platform.system() == "Darwin": from coreai.runtime import ComputeUnitKind, SpecializationOptions -# Substring of the dtype guard message raised by the CoreML export validation. Shared so -# test files asserting the rejection don't drift from one another. -COREML_DTYPE_REJECTION_MATCH = "CoreML export does not support" +# Substring of the guard message raised by CoreML export validation (dtype or +# granularity/config rejection). Shared so test files asserting the rejection +# don't drift from one another. +COREML_REJECTION_MATCH = "CoreML export does not support" # Compute unit selection driven by the --compute-unit-kind pytest option (see # tests/conftest.py). Default is "interpreter" so a plain `pytest` run uses the @@ -87,17 +88,18 @@ def _get_test_specialization_options() -> "SpecializationOptions | None": raise ValueError(msg) -def assert_coreml_finalize_rejects_unsupported_dtype(finalizer: Any) -> None: - """Assert ``finalizer.finalize(backend=CoreML)`` rejects an unsupported dtype. +def assert_coreml_finalize_rejects(finalizer: Any) -> None: + """Assert ``finalizer.finalize(backend=CoreML)`` rejects an unsupported config. CoreML does not support FP4, FP8, INT2, or UINT2 quantization or - palettization dtypes, so finalize must raise a ``CoreMLExportError`` rather - than emit an invalid model. + palettization dtypes, nor per-channel/per-block activation quantization + granularity, so finalize must raise a ``CoreMLExportError`` rather than + emit an invalid model. Args: finalizer (Any): A prepared ``Quantizer`` or ``KMeansPalettizer``. """ - with pytest.raises(CoreMLExportError, match=COREML_DTYPE_REJECTION_MATCH): + with pytest.raises(CoreMLExportError, match=COREML_REJECTION_MATCH): finalizer.finalize(backend=ExportBackend.CoreML) diff --git a/tests/export/test_eager_mil_export.py b/tests/export/test_eager_mil_export.py index 6ab3eb9..ffe9835 100644 --- a/tests/export/test_eager_mil_export.py +++ b/tests/export/test_eager_mil_export.py @@ -94,6 +94,10 @@ def test_simple_model_export( parametrized_quant_config_general: ParametrizedQuantConfigs, ) -> None: """Test eager CoreML export with various quantization configurations.""" + # CoreML rejects per-channel activation quantization outright. + if parametrized_quant_config_general.has_per_channel_activation_granularity: + pytest.skip("CoreML export rejects per-channel activation quantization") + has_act_quant = parametrized_quant_config_general.has_activation_quantization _run_eager_mil_export_test( @@ -114,6 +118,10 @@ def test_mnist_export( parametrized_quant_config_general: ParametrizedQuantConfigs, ) -> None: """Test eager CoreML export on MNIST model with various quantization configurations.""" + # CoreML rejects per-channel activation quantization outright. + if parametrized_quant_config_general.has_per_channel_activation_granularity: + pytest.skip("CoreML export rejects per-channel activation quantization") + has_act_quant = parametrized_quant_config_general.has_activation_quantization _run_eager_mil_export_test( @@ -167,22 +175,31 @@ def test_gated_mlp_perchannel_act_export( parametrized_quant_config_perchannel_act_axis_coverage: ParametrizedQuantConfigs, ) -> None: """Test eager CoreML export with per-channel activation quantization axes. - Uses GatedMLPModel (uniform rank-3 activations throughout the model) to - test per-channel activation quantization across all valid axis values without - out-of-bounds errors. + + Uses GatedMLPModel (uniform rank-3 activations throughout the model). CoreML + export rejects per-channel activation quantization outright; per-tensor cases + from this fixture should still export and verify normally. """ - has_act_quant = ( - parametrized_quant_config_perchannel_act_axis_coverage.has_activation_quantization - ) + config = parametrized_quant_config_perchannel_act_axis_coverage + + if config.has_per_channel_activation_granularity: + with pytest.raises(CoreMLExportError): + _run_eager_mil_export_test( + model=gated_mlp_model, + input_data=gated_mlp_model_input, + parametrized_quant_config=config, + expected_ops={}, + ) + return _run_eager_mil_export_test( model=gated_mlp_model, input_data=gated_mlp_model_input, - parametrized_quant_config=parametrized_quant_config_perchannel_act_axis_coverage, + parametrized_quant_config=config, expected_ops={ "constexpr_blockwise_shift_scale": 3, - "quantize": 6 if has_act_quant else 0, - "dequantize": 6 if has_act_quant else 0, + "quantize": 6, + "dequantize": 6, }, ) diff --git a/tests/export/test_kmeans_export.py b/tests/export/test_kmeans_export.py index 7af7651..1858ca0 100644 --- a/tests/export/test_kmeans_export.py +++ b/tests/export/test_kmeans_export.py @@ -76,7 +76,7 @@ def _assert_coreml_rejects_unsupported_lut( model.eval() palettizer = KMeansPalettizer(model, config) palettizer.prepare((input_data,)) - export_utils.assert_coreml_finalize_rejects_unsupported_dtype(palettizer) + export_utils.assert_coreml_finalize_rejects(palettizer) def _skip_heavy_mnist_configs(config: ParametrizedPalettConfigs) -> None: diff --git a/tests/export/test_pt2e_mil_export.py b/tests/export/test_pt2e_mil_export.py index 4fa4b2f..b59c726 100644 --- a/tests/export/test_pt2e_mil_export.py +++ b/tests/export/test_pt2e_mil_export.py @@ -11,6 +11,7 @@ from coreai_opt import ExportBackend from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig from coreai_opt.quantization.spec import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationGranularity, @@ -28,29 +29,50 @@ COREML_WEIGHT_REJECT_DTYPES, make_quant_config, ) -from tests.test_utils.general import SNRBelowThresholdError from . import export_utils # TODO: migrate to using conftest.py for fixtures. + +def _build_simple_model_quant_config( + weight_dtype: torch.dtype, + weight_granularity: QuantizationGranularity, + act_granularity: QuantizationGranularity, + qscheme: QuantizationScheme, +) -> QuantizerConfig: + """Build a QuantizerConfig for the simple conv/linear model's weight and activations.""" + weight_qspec = QuantizationSpec( + dtype=weight_dtype, + qscheme=qscheme, + granularity=weight_granularity, + fake_quantize_cls=_DefaultFakeQuantizeImpl, + qparam_calculator_cls=StaticQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + activation_qspec = QuantizationSpec( + dtype=torch.uint8, + qscheme=qscheme, + granularity=act_granularity, + fake_quantize_cls=_DefaultFakeQuantizeImpl, + qparam_calculator_cls=MovingAverageQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_qspec}, + op_input_spec={"*": activation_qspec}, + op_output_spec={"*": activation_qspec}, + ), + ) + + _test_params = ( [ # int8 weights with per-tensor granularity (torch.int8, PerTensorGranularity(), PerTensorGranularity(), qscheme) for qscheme in QuantizationScheme ] - + [ - # uint8 weights with per-channel (axis=0) granularity - (torch.uint8, PerChannelGranularity(axis=0), PerChannelGranularity(axis=0), qscheme) - for qscheme in QuantizationScheme - ] - + [ - # uint8 weights with per-channel (axis=0) granularity and per-channel activations - # with negative axis - (torch.uint8, PerChannelGranularity(axis=0), PerChannelGranularity(axis=-1), qscheme) - for qscheme in QuantizationScheme - ] + [ # int4 weights with per-tensor granularity (low-bit quantization) (torch.int4, PerTensorGranularity(), PerTensorGranularity(), qscheme) @@ -61,23 +83,11 @@ (torch.int8, PerChannelGranularity(axis=1), PerTensorGranularity(), qscheme) for qscheme in QuantizationScheme ] - + [ - # int8 weights with per-channel (axis=1) and per-channel activations with - # negative axis - (torch.int8, PerChannelGranularity(axis=1), PerChannelGranularity(axis=-1), qscheme) - for qscheme in QuantizationScheme - ] + [ # int4 weights with per-channel (axis=1) - low-bit with per-channel (torch.int4, PerChannelGranularity(axis=1), PerTensorGranularity(), qscheme) for qscheme in QuantizationScheme ] - + [ - # int4 weights with per-channel (axis=1) - low-bit with per-channel and per-channel - # activations with negative axis - (torch.int4, PerChannelGranularity(axis=1), PerChannelGranularity(axis=-1), qscheme) - for qscheme in QuantizationScheme - ] + [ # uint8 weights with per-tensor granularity, asymmetric qscheme ( @@ -98,31 +108,10 @@ for dtype, wg, ag, qscheme in _test_params ] -# Mark known failing cases as xfail -_test_params_with_xfail: list = [] -for params in _test_params: - dtype, wg, ag, qscheme = params - # Per-channel activation with negative axis has SNR below threshold. - # TODO: SNR below threshold for per-channel activation with negative axis on CoreML export. - if isinstance(ag, PerChannelGranularity) and ag.axis is not None and ag.axis < 0: - _test_params_with_xfail.append( - pytest.param( - *params, - # TODO: SNR below threshold for per-channel activation with negative axis - # on CoreML export. - marks=pytest.mark.xfail( - raises=SNRBelowThresholdError, - reason="SNR below threshold for per-channel activation with negative axis.", - ), - ), - ) - else: - _test_params_with_xfail.append(params) - @pytest.mark.parametrize( ("weight_dtype", "weight_granularity", "act_granularity", "qscheme"), - _test_params_with_xfail, + _test_params, ids=_test_ids, ) def test_simple_model_export( @@ -137,28 +126,8 @@ def test_simple_model_export( model = simple_conv_linear_model model.eval() - weight_qspec = QuantizationSpec( - dtype=weight_dtype, - qscheme=qscheme, - granularity=weight_granularity, - fake_quantize_cls=_DefaultFakeQuantizeImpl, - qparam_calculator_cls=StaticQParamsCalculator, - range_calculator_cls=MinMaxRangeCalculator, - ) - activation_qspec = QuantizationSpec( - dtype=torch.uint8, - qscheme=qscheme, - granularity=act_granularity, - fake_quantize_cls=_DefaultFakeQuantizeImpl, - qparam_calculator_cls=MovingAverageQParamsCalculator, - range_calculator_cls=MinMaxRangeCalculator, - ) - config = QuantizerConfig( - global_config=ModuleQuantizerConfig( - op_state_spec={"weight": weight_qspec}, - op_input_spec={"*": activation_qspec}, - op_output_spec={"*": activation_qspec}, - ), + config = _build_simple_model_quant_config( + weight_dtype, weight_granularity, act_granularity, qscheme ) quantizer = Quantizer(model, config) prepared_model = quantizer.prepare((simple_model_input,)) @@ -181,6 +150,57 @@ def test_simple_model_export( ) +# CoreML export rejects per-channel and per-block activation quantization outright, +# regardless of axis, dtype, or qscheme: an upstream coremltools MIL pass can +# silently corrupt it. +_perchannel_act_reject_params = [ + ( + torch.uint8, + PerChannelGranularity(axis=0), + PerChannelGranularity(axis=0), + QuantizationScheme.SYMMETRIC, + ), + ( + torch.int8, + PerChannelGranularity(axis=1), + PerChannelGranularity(axis=-1), + QuantizationScheme.ASYMMETRIC, + ), + ( + torch.int8, + PerTensorGranularity(), + PerBlockGranularity(axis=1, block_size=2), + QuantizationScheme.SYMMETRIC, + ), +] +_perchannel_act_reject_ids = ["axis0--symmetric", "axis-1--asymmetric", "per-block--symmetric"] + + +@pytest.mark.parametrize( + ("weight_dtype", "weight_granularity", "act_granularity", "qscheme"), + _perchannel_act_reject_params, + ids=_perchannel_act_reject_ids, +) +def test_simple_model_export_rejects_perchannel_activation( + simple_conv_linear_model: torch.nn.Module, + simple_model_input: torch.Tensor, + weight_dtype: torch.dtype, + weight_granularity: QuantizationGranularity, + act_granularity: QuantizationGranularity, + qscheme: QuantizationScheme, +) -> None: + """Per-channel activation quantization must be rejected on CoreML export.""" + model = simple_conv_linear_model + model.eval() + + config = _build_simple_model_quant_config( + weight_dtype, weight_granularity, act_granularity, qscheme + ) + quantizer = Quantizer(model, config) + quantizer.prepare((simple_model_input,)) + export_utils.assert_coreml_finalize_rejects(quantizer) + + _mnist_test_params = [ # Per-tensor weight + per-tensor activation (torch.int8, PerTensorGranularity(), PerTensorGranularity(), qscheme) @@ -322,20 +342,10 @@ def test_resnet_export( ) -# Per-channel activation axis test params for GatedMLPModel. -_gated_mlp_test_params = [ - (act_gran, qscheme) - for act_gran in [ - PerTensorGranularity(), - PerChannelGranularity(axis=0), - PerChannelGranularity(axis=1), - PerChannelGranularity(axis=2), - PerChannelGranularity(axis=-1), - PerChannelGranularity(axis=-2), - PerChannelGranularity(axis=-3), - ] - for qscheme in QuantizationScheme -] +# Activation granularity test params for GatedMLPModel. CoreML export rejects +# PerChannelGranularity activations outright (see test_gated_mlp_export_rejects_ +# perchannel_activation below), so only PerTensorGranularity actually exports here. +_gated_mlp_test_params = [(PerTensorGranularity(), qscheme) for qscheme in QuantizationScheme] _gated_mlp_test_ids = [ f"ag:{ag.__class__.__name__.replace('Granularity', '')}--" @@ -345,24 +355,11 @@ def test_resnet_export( ] -@pytest.mark.parametrize( - ("act_granularity", "qscheme"), - _gated_mlp_test_params, - ids=_gated_mlp_test_ids, -) -def test_gated_mlp_export( - gated_mlp_model: torch.nn.Module, - gated_mlp_model_input: torch.Tensor, +def _build_gated_mlp_quant_config( act_granularity: QuantizationGranularity, qscheme: QuantizationScheme, -) -> None: - """Test PT2E CoreML export with per-channel activation quantization axes. - Uses GatedMLPModel (uniform rank-3 activations throughout the model) to - test per-channel activation quantization across all valid axis values. - """ - model = gated_mlp_model - model.eval() - +) -> QuantizerConfig: + """Build a QuantizerConfig for GatedMLPModel's weight and activations.""" weight_qspec = QuantizationSpec( dtype=torch.int8, qscheme=qscheme, @@ -379,13 +376,31 @@ def test_gated_mlp_export( qparam_calculator_cls=MovingAverageQParamsCalculator, range_calculator_cls=MinMaxRangeCalculator, ) - config = QuantizerConfig( + return QuantizerConfig( global_config=ModuleQuantizerConfig( op_state_spec={"weight": weight_qspec}, op_input_spec={"*": activation_qspec}, op_output_spec={"*": activation_qspec}, ), ) + + +@pytest.mark.parametrize( + ("act_granularity", "qscheme"), + _gated_mlp_test_params, + ids=_gated_mlp_test_ids, +) +def test_gated_mlp_export( + gated_mlp_model: torch.nn.Module, + gated_mlp_model_input: torch.Tensor, + act_granularity: QuantizationGranularity, + qscheme: QuantizationScheme, +) -> None: + """Test PT2E CoreML export with per-tensor activation quantization.""" + model = gated_mlp_model + model.eval() + + config = _build_gated_mlp_quant_config(act_granularity, qscheme) quantizer = Quantizer(model, config) prepared_model = quantizer.prepare((gated_mlp_model_input,)) @@ -407,6 +422,37 @@ def test_gated_mlp_export( ) +# CoreML export rejects per-channel and per-block activation quantization outright, +# regardless of axis: an upstream coremltools MIL pass can silently corrupt it. +_gated_mlp_reject_params = [ + (PerChannelGranularity(axis=0), QuantizationScheme.SYMMETRIC), + (PerChannelGranularity(axis=-1), QuantizationScheme.ASYMMETRIC), + (PerBlockGranularity(axis=1, block_size=2), QuantizationScheme.SYMMETRIC), +] +_gated_mlp_reject_ids = ["axis0--symmetric", "axis-1--asymmetric", "per-block--symmetric"] + + +@pytest.mark.parametrize( + ("act_granularity", "qscheme"), + _gated_mlp_reject_params, + ids=_gated_mlp_reject_ids, +) +def test_gated_mlp_export_rejects_perchannel_activation( + gated_mlp_model: torch.nn.Module, + gated_mlp_model_input: torch.Tensor, + act_granularity: QuantizationGranularity, + qscheme: QuantizationScheme, +) -> None: + """Per-channel activation quantization must be rejected regardless of axis.""" + model = gated_mlp_model + model.eval() + + config = _build_gated_mlp_quant_config(act_granularity, qscheme) + quantizer = Quantizer(model, config) + quantizer.prepare((gated_mlp_model_input,)) + export_utils.assert_coreml_finalize_rejects(quantizer) + + # Unsupported dtypes (FP4, FP8, INT2, UINT2) must be rejected on CoreML export; # finalize must reject them. Dtype lists and the config builder live in conftest # (shared with the eager tests). @@ -422,7 +468,7 @@ def test_unsupported_weight_quant_coreml_export_rejected( model.eval() quantizer = Quantizer(model, config) quantizer.prepare((simple_model_input,)) - export_utils.assert_coreml_finalize_rejects_unsupported_dtype(quantizer) + export_utils.assert_coreml_finalize_rejects(quantizer) @pytest.mark.parametrize("act_dtype", COREML_ACT_REJECT_DTYPES) @@ -437,4 +483,4 @@ def test_unsupported_activation_quant_coreml_export_rejected( model.eval() quantizer = Quantizer(model, config) quantizer.prepare((simple_model_input,)) - export_utils.assert_coreml_finalize_rejects_unsupported_dtype(quantizer) + export_utils.assert_coreml_finalize_rejects(quantizer) diff --git a/tests/fixtures/quantization.py b/tests/fixtures/quantization.py index 82d17ae..4ee165b 100644 --- a/tests/fixtures/quantization.py +++ b/tests/fixtures/quantization.py @@ -183,6 +183,19 @@ def has_activation_quantization(self) -> bool: else False ) + @property + def has_per_channel_activation_granularity(self) -> bool: + """Check if activation quantization in this config uses PerChannelGranularity. + + Returns: + True if activation quantization is enabled and per-channel + + """ + if not self.eager.global_config: + return False + act_qspec = self.eager.global_config.op_input_spec.get("*") + return act_qspec is not None and isinstance(act_qspec.granularity, PerChannelGranularity) + def skip_if_unsupported( self, mode: Literal["eager", "graph"],