diff --git a/changelog.d/133.added b/changelog.d/133.added new file mode 100644 index 00000000..ac9be5a2 --- /dev/null +++ b/changelog.d/133.added @@ -0,0 +1 @@ +Add MXINT8 support. diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 4ddf337b..0fb4f12f 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -24,6 +24,7 @@ # Mapping from dtype to the largest power-of-2 component of its max value. # Used by e8m0 scale computation (OCP Microscaling spec, FLOOR mode). E8M0_TARGET_MAX_POW2: dict[torch.dtype, int] = { + torch.int8: 6, # max = 127 = 1.98 * 2^6 torch.float4_e2m1fn_x2: 2, torch.float8_e4m3fn: 8, # max = 448.0 = 1.75 * 2^8 torch.float8_e5m2: 15, # max = 57344.0 = 1.75 * 2^15 diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index 24baeb99..97836851 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -132,6 +132,13 @@ def validate_lut_qspec(self) -> "PalettizationSpec": "Use lut_qspec.qformulation=ZP instead." ) + if not self.lut_qspec.dtype.is_floating_point and self.lut_qspec.scale_dtype is not None: + msg = ( + "lut_qspec.scale_dtype must be None for integer lut_qspec.dtype, " + f"got scale_dtype={self.lut_qspec.scale_dtype}" + ) + raise ValueError(msg) + return self def model_dump_preserve_objects(self) -> dict[str, Any]: diff --git a/src/coreai_opt/quantization/_export_utils.py b/src/coreai_opt/quantization/_export_utils.py index 0d97a1e1..a3c82b52 100644 --- a/src/coreai_opt/quantization/_export_utils.py +++ b/src/coreai_opt/quantization/_export_utils.py @@ -371,20 +371,18 @@ def get_activation_export_handler( return _GRAPH_ACTIVATION_EXPORT_HANDLERS.get(type(granularity)) -def can_export_stateless_fake_quant( +def has_activation_export_handler( fake_quant_mod: FakeQuantizeImplBase, backend: ExportBackend, execution_mode: ExecutionMode, ) -> bool: - """Return whether a recompute-every-forward calculator can still be exported. + """Return whether a registered activation handler will export this module. - The built-in export paths bake qparams into buffers, so a calculator that - recomputes them every forward has nothing to bake. A registered activation handler - emits its own ops and can express the recomputation, so it lifts the - restriction for the activation it covers. + A handler emits its own ops, so callers use this to lift restrictions of the + built-in export paths. Args: - fake_quant_mod (FakeQuantizeImplBase): The module whose calculator is stateless. + fake_quant_mod (FakeQuantizeImplBase): The module to export. backend (ExportBackend): The export target. execution_mode (ExecutionMode): The quantizer's execution mode. diff --git a/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py index d8fb92ca..992f4540 100644 --- a/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py +++ b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py @@ -280,10 +280,11 @@ def _apply_op_intrinsic_override( Nothing else: the op has no opinion on the other fields. - The ``QSCHEME`` half is skipped for a floating-point dtype, which admits only - the symmetric scheme — see ``QuantizationSpec.validate_qscheme_for_fp_quant``. + The ``QSCHEME`` half is skipped for a floating-point dtype or an e8m0 scale, + which admit only the symmetric scheme — see + ``QuantizationSpec.validate_qscheme_for_fp_quant``. An op like ``relu`` proposes an asymmetric scheme to exploit its one-sided - range, but that is a representational choice the dtype forbids, so proposing it + range, but that is a representational choice the spec forbids, so proposing it would make the group unbuildable rather than merely lower quality. The ``FLOAT_RANGE`` half still applies: the bound is a fact about the data and holds whatever the dtype. @@ -292,7 +293,7 @@ def _apply_op_intrinsic_override( assert intrinsic is not None, "caller must gate on _get_op_intrinsic" scheme, float_range = intrinsic - if not user_spec.dtype.is_floating_point: + if not user_spec.dtype.is_floating_point and user_spec.scale_dtype is None: _override_field( qspecs, output_slot, diff --git a/src/coreai_opt/quantization/_graph/_qspec_resolution.py b/src/coreai_opt/quantization/_graph/_qspec_resolution.py index e8a7d541..626700f0 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_resolution.py +++ b/src/coreai_opt/quantization/_graph/_qspec_resolution.py @@ -26,10 +26,11 @@ ) from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY -from coreai_opt.quantization.spec import QuantizationSpec +from coreai_opt.quantization.spec import QuantizationScheme, QuantizationSpec from ._annotation_config import AnnotationConfig from ._qspec_types import ( + OP_INTRINSIC_PRIORITY, FieldName, NodeSlot, ProvisionalQSpec, @@ -207,12 +208,15 @@ def _build_concrete_spec(qspec: ProvisionalQSpec) -> TorchAOQuantizationSpec | N f"{sorted(f.name for f in qspec.fields)}." ) - spec = QuantizationSpec( - **{ - kwarg: qspec.fields[field_name].value - for field_name, kwarg in _SPEC_KWARG_FROM_FIELD.items() - } - ) + kwargs = { + kwarg: qspec.fields[field_name].value + for field_name, kwarg in _SPEC_KWARG_FROM_FIELD.items() + } + if qspec.fields[FieldName.QSCHEME].priority == OP_INTRINSIC_PRIORITY and ( + kwargs["dtype"].is_floating_point or kwargs["scale_dtype"] is not None + ): + kwargs["qscheme"] = QuantizationScheme.SYMMETRIC + spec = QuantizationSpec(**kwargs) return _convert_to_pt2e_spec(spec, qspec.fields[FieldName.QUANTIZATION_TARGET].value) diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py index 501396ae..6474f620 100644 --- a/src/coreai_opt/quantization/_graph/quantizer.py +++ b/src/coreai_opt/quantization/_graph/quantizer.py @@ -1055,7 +1055,10 @@ def prepare( try: prepared_model = prepare_qat_pt2e(exported_model, quantizer) except Exception as e: - raise type(e)(f"prepare_qat_pt2e call failed, with error: {e}") from e + # Not type(e)(msg): exceptions such as pydantic's ValidationError take + # other constructor arguments, and rebuilding one raises a TypeError. + e.add_note("prepare_qat_pt2e call failed") + raise if torchao_requires_empty_kwargs: restore_kwargs(prepared_model.graph, saved_kwargs) diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py index 581edc0a..02c6117d 100644 --- a/src/coreai_opt/quantization/quantizer.py +++ b/src/coreai_opt/quantization/quantizer.py @@ -6,10 +6,12 @@ from __future__ import annotations import warnings +from collections.abc import Callable from contextlib import contextmanager from os import PathLike from typing import Any +import torch import torch.nn as nn from torch import fx from torchao.quantization.pt2e import ( @@ -27,7 +29,7 @@ from coreai_opt.common import ExportBackend from coreai_opt.quantization._eager import EagerQuantizer as _EagerQuantizer from coreai_opt.quantization._export_utils import ( - can_export_stateless_fake_quant as _can_export_stateless_fake_quant, + has_activation_export_handler as _has_activation_export_handler, ) from coreai_opt.quantization._graph import GraphQuantizer as _GraphQuantizer from coreai_opt.quantization.base_quantizer import _BaseQuantizer @@ -417,16 +419,9 @@ def _validate_no_persistent_observer_calculators( ``EagerActivationExportHandler`` in eager mode. """ - if backend == ExportBackend._TORCH: - return - model_to_check = model if model is not None else self._model - stateless_fq_names = [ - name - for name, mod in model_to_check.named_modules() - if isinstance(mod, FakeQuantizeImplBase) - and mod.is_stateless - and not _can_export_stateless_fake_quant(mod, backend, self._execution_mode) - ] + stateless_fq_names = self._fake_quant_names_without_export_handler( + model, backend, lambda mod: mod.is_stateless + ) if stateless_fq_names: raise NotImplementedError( f"backend={backend} does not yet support qparams calculators that " @@ -435,6 +430,54 @@ def _validate_no_persistent_observer_calculators( f"backend=ExportBackend._TORCH for torch-only inference." ) + def _validate_no_integer_e8m0_scales( + self, + model: nn.Module | fx.GraphModule | None, + backend: ExportBackend, + ) -> None: + """Reject CoreAI/CoreML export of integer quantization with e8m0 scales. + + The built-in integer export ops cannot take an e8m0 scale; a registered + activation export handler lifts the restriction. + """ + e8m0_int_fq_names = self._fake_quant_names_without_export_handler( + model, + backend, + lambda mod: ( + not mod.dtype.is_floating_point + and mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu + ), + ) + if e8m0_int_fq_names: + msg = ( + f"backend={backend} does not support integer quantization with " + f"scale_dtype={torch.float8_e8m0fnu}. Affected FakeQuantize modules: " + f"{e8m0_int_fq_names}. Use backend=ExportBackend._TORCH for torch-only " + f"inference." + ) + raise NotImplementedError(msg) + + def _fake_quant_names_without_export_handler( + self, + model: nn.Module | fx.GraphModule | None, + backend: ExportBackend, + predicate: Callable[[FakeQuantizeImplBase], bool], + ) -> list[str]: + """Return FakeQuantize modules matching ``predicate`` with no activation export handler. + + Always empty for ``ExportBackend._TORCH``. + """ + if backend == ExportBackend._TORCH: + return [] + model_to_check = model if model is not None else self._model + return [ + name + for name, mod in model_to_check.named_modules() + if isinstance(mod, FakeQuantizeImplBase) + and predicate(mod) + and not _has_activation_export_handler(mod, backend, self._execution_mode) + ] + def finalize( self, model: nn.Module | fx.GraphModule | None = None, @@ -484,6 +527,7 @@ def finalize( """ self._validate_mmap_dir_constraints(model, backend, mmap_dir) self._validate_no_persistent_observer_calculators(model, backend) + self._validate_no_integer_e8m0_scales(model, backend) return self._quantizer.finalize(model, backend, mmap_dir=mmap_dir) @contextmanager diff --git a/src/coreai_opt/quantization/spec/qparams_calculator.py b/src/coreai_opt/quantization/spec/qparams_calculator.py index 28aa49a7..3e2ee6cd 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -221,12 +221,16 @@ def _compute_scale_zero_point_minval( scale = 2^(floor(log2(max_abs)) - target_max_pow2) where ``target_max_pow2`` comes from ``E8M0_TARGET_MAX_POW2``, which is the single source of truth for which dtypes support an e8m0 scale. + The spec pins this mode to symmetric ZP, so the zero point is zero. """ # e8m0 path: power-of-2 scales if self.scale_dtype == torch.float8_e8m0fnu: max_abs = torch.maximum(torch.abs(min_val), torch.abs(max_val)) - return self._compute_e8m0_scale(max_abs), None, None + scale = self._compute_e8m0_scale(max_abs) + if self.dtype.is_floating_point: + return scale, None, None + return scale, torch.zeros_like(scale, dtype=torch.int32), -max_abs # Default path: torchao handles integer and FP8 dtypes scale, zero_point = choose_qparams_affine_with_min_max( diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index c721cf3f..5ee96884 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -85,9 +85,9 @@ class QuantizationSpec(CompressionSpec): - ``x``: unquantized data. - ``q``: quantized data (dtype as specified by ``QuantizationSpec.dtype``). - ``x'``: dequantized data (same dtype as ``x``). - - ``scale``: for INT quantization, defaults to the same dtype as - ``x``. For FP quantization, see the ``scale_dtype`` description - below. + - ``scale``: for INT quantization without an e8m0 ``scale_dtype``, + defaults to the same dtype as ``x``. Otherwise, see the + ``scale_dtype`` description below. Formulae: @@ -145,16 +145,19 @@ class QuantizationSpec(CompressionSpec): | | | | | min_val_neg/scale), 0, 255) | +-------+------------+-------------+-----------------+--------------------------------+ - And for FP4/FP8 dtypes, zero-point is always set to 0 (FP supports only the - symmetric qscheme). The scale formula depends on ``scale_dtype``: + And for FP4/FP8 dtypes, and INT8 with an e8m0 ``scale_dtype``, zero-point is + always set to 0 (both support only the symmetric qscheme). The scale formula + depends on ``scale_dtype``: - ``scale_dtype=None`` (FP8 only): ``scale = max_abs / fp_max``, where ``fp_max`` is the largest representable value for the target FP dtype (448.0 for FP8 E4M3, 57344.0 for FP8 E5M2). - - ``scale_dtype=torch.float8_e8m0fnu`` (FP4 and FP8): + - ``scale_dtype=torch.float8_e8m0fnu`` (FP4, FP8 and INT8): power-of-2 scale per OCP MX spec — ``scale = 2^(floor(log2(max_abs)) - target_max_pow2)``, with - ``target_max_pow2`` of 2 for FP4 E2M1, 8 for FP8 E4M3, 15 for FP8 E5M2. + ``target_max_pow2`` of 2 for FP4 E2M1, 8 for FP8 E4M3, 15 for FP8 E5M2, + and 6 for INT8. OCP MXINT8 keeps a ``2^-6`` factor in each element and a + shared scale 64x larger; here that factor is folded into ``scale``. **MINVAL formulation**, e.g. with 8 bit signed and unsigned fixed point types: @@ -310,15 +313,16 @@ class type: MinMaxRangeCalculator or custom registered class type choose_qparams_affine_with_min_max (integer and FP8 dtypes). For FP4, None is resolved to torch.float8_e8m0fnu automatically. - torch.float8_e8m0fnu: Power-of-2 scales following OCP Microscaling (MX) - spec. Required for FP4 quantization, optional for FP8. + spec. Required for FP4 quantization, optional for FP8 and int8. Constraints: - FP4 (float4_e2m1fn_x2): scale_dtype must be torch.float8_e8m0fnu or None (defaults to e8m0) - - FP8 (float8_e4m3fn, float8_e5m2): scale_dtype must be + - FP8 (float8_e4m3fn, float8_e5m2) and int8: scale_dtype must be torch.float8_e8m0fnu or None (defaults to None) - - Integer dtypes: scale_dtype must be None (defaults to None) + - Other integer dtypes: scale_dtype must be None (defaults to None) + - torch.float8_e8m0fnu requires qscheme="symmetric" and qformulation="zp" Default: None @@ -502,46 +506,18 @@ def resolve_scale_dtype(cls, data: Any) -> Any: data["scale_dtype"] = torch.float8_e8m0fnu return data - @model_validator(mode="after") - def validate_qscheme_for_fp_quant(self) -> QuantizationSpec: - """ - Validate that FP quantization uses symmetric quantization scheme. - """ - if self.dtype.is_floating_point: - if self.qscheme != QuantizationScheme.SYMMETRIC: - error_msg = ( - f"FP quantization (dtype={self.dtype}) requires " - f"symmetric quantization scheme, got " - f"qscheme={self.qscheme}. Valid option: 'symmetric'" - ) - raise ValueError(error_msg) - return self - - @model_validator(mode="after") - def validate_qformulation_for_fp_quant(self) -> QuantizationSpec: - """ - Validate that FP quantization uses zero-point quantization formulation. - """ - if self.dtype.is_floating_point: - if self.qformulation != QuantizationFormulation.ZP: - error_msg = ( - f"FP quantization (dtype={self.dtype}) requires " - f"zero-point quantization formulation, got " - f"qformulation={self.qformulation}. Valid option: 'zp'" - ) - raise ValueError(error_msg) - return self - + # Defined before the qscheme and qformulation checks so it runs first: pydantic + # runs after-validators in definition order, and those checks would otherwise + # ask a user with an unsupported scale_dtype to fix a field that isn't the problem. @model_validator(mode="after") def validate_scale_dtype(self) -> QuantizationSpec: - """ - Validate scale_dtype based on element dtype. + """Validate scale_dtype based on element dtype. Rules: - Only None or torch.float8_e8m0fnu are supported. - When scale_dtype is set, dtype must be a key of ``E8M0_TARGET_MAX_POW2`` - (the dtypes with an e8m0 target exponent); for any other dtype, including - all integer dtypes, scale_dtype must be None. + (the dtypes with an e8m0 target exponent); for any other dtype, + scale_dtype must be None. - FP4 dtypes: scale_dtype is resolved to torch.float8_e8m0fnu by resolve_scale_dtype (before validator). """ @@ -549,20 +525,54 @@ def validate_scale_dtype(self) -> QuantizationSpec: return self if self.scale_dtype != torch.float8_e8m0fnu: - raise ValueError( + msg = ( f"Unsupported scale_dtype: {self.scale_dtype}. " f"Only None or torch.float8_e8m0fnu are supported." ) + raise ValueError(msg) if self.dtype not in _E8M0_TARGET_MAX_POW2: - raise ValueError( + msg = ( f"scale_dtype must be None for dtype={self.dtype}, got " f"scale_dtype={self.scale_dtype}. No e8m0 target exponent is defined " f"for it; supported dtypes are {list(_E8M0_TARGET_MAX_POW2)}." ) + raise ValueError(msg) return self + @model_validator(mode="after") + def validate_qscheme_for_fp_quant(self) -> QuantizationSpec: + """ + Validate that FP quantization and e8m0 scales use symmetric quantization scheme. + """ + if self.dtype.is_floating_point or self.scale_dtype is not None: + if self.qscheme != QuantizationScheme.SYMMETRIC: + error_msg = ( + f"FP quantization or an e8m0 scale (dtype={self.dtype}, " + f"scale_dtype={self.scale_dtype}) requires " + f"symmetric quantization scheme, got " + f"qscheme={self.qscheme}. Valid option: 'symmetric'" + ) + raise ValueError(error_msg) + return self + + @model_validator(mode="after") + def validate_qformulation_for_fp_quant(self) -> QuantizationSpec: + """ + Validate that FP quantization and e8m0 scales use zero-point quantization formulation. + """ + if self.dtype.is_floating_point or self.scale_dtype is not None: + if self.qformulation != QuantizationFormulation.ZP: + error_msg = ( + f"FP quantization or an e8m0 scale (dtype={self.dtype}, " + f"scale_dtype={self.scale_dtype}) requires " + f"zero-point quantization formulation, got " + f"qformulation={self.qformulation}. Valid option: 'zp'" + ) + raise ValueError(error_msg) + return self + @model_validator(mode="after") def validate_qscheme_for_dynamic_quant(self) -> QuantizationSpec: """Reject ``SYMMETRIC_WITH_CLIPPING`` with a dynamic (stateless) calculator.""" diff --git a/tests/palettization/test_palettization_spec.py b/tests/palettization/test_palettization_spec.py index 2d6b98a7..910c9956 100644 --- a/tests/palettization/test_palettization_spec.py +++ b/tests/palettization/test_palettization_spec.py @@ -267,6 +267,16 @@ def test_lut_qspec_validation(): ) ) + # Invalid: e8m0 scale on an integer LUT + with pytest.raises(ValidationError, match="lut_qspec.scale_dtype must be None"): + PalettizationSpec( + lut_qspec=QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + scale_dtype=torch.float8_e8m0fnu, + ) + ) + @pytest.mark.parametrize( "granularity", [PerTensorGranularity(), PerGroupedChannelGranularity(axis=1, group_size=8)] diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index 658532fe..565f063f 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -14,6 +14,7 @@ import pytest import torch import torch.nn as nn +from pydantic import ValidationError from torch.export.dynamic_shapes import Dim from torch.ops import coreai from torchao.quantization.pt2e.quantizer import ( @@ -146,6 +147,25 @@ def test_prepare_basic(self, simple_conv_linear_model, basic_config, simple_mode output = prepared_model(simple_model_input) assert output.shape == (1, 10) + def test_prepare_failure_reaches_the_caller_unchanged( + self, simple_conv_linear_model, basic_config, simple_model_input, monkeypatch + ): + """An exception that takes more than a message to construct keeps its type.""" + + def _failing_prepare(*_args, **_kwargs): + QuantizationSpec( + dtype=torch.int8, qscheme="asymmetric", scale_dtype=torch.float8_e8m0fnu + ) + + monkeypatch.setattr( + "coreai_opt.quantization._graph.quantizer.prepare_qat_pt2e", _failing_prepare + ) + quantizer = Quantizer(simple_conv_linear_model, basic_config) + + with pytest.raises(ValidationError, match="requires symmetric") as exc_info: + quantizer.prepare((simple_model_input,)) + assert exc_info.value.__notes__ == ["prepare_qat_pt2e call failed"] + def test_prepare_with_dynamic_shapes( self, simple_conv_linear_model, basic_config, simple_model_input ): diff --git a/tests/quantization/test_qparams_calculator.py b/tests/quantization/test_qparams_calculator.py index 2b5a9628..435c17c5 100644 --- a/tests/quantization/test_qparams_calculator.py +++ b/tests/quantization/test_qparams_calculator.py @@ -348,6 +348,36 @@ def test_per_block_scale_zero_point( ) +@pytest.mark.parametrize("qparam_calculator_cls", ["default", "dynamic"]) +def test_int8_e8m0_scale_zero_point_minval(qparam_calculator_cls): + """int8 with an e8m0 scale gets a power-of-2 scale and zero offsets for integer fake-quant.""" + x = torch.tensor([[-3.0, 0.01, 1.01, 3.0]]) + spec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerTensorGranularity(), + qparam_calculator_cls=qparam_calculator_cls, + scale_dtype=torch.float8_e8m0fnu, + ) + fake_quant = QuantizationComponentFactory.create_fake_quantizer( + spec, CompressionTargetTensor.ACTIVATION + ) + + scale, zero_point, minval = fake_quant.qparams_calculator(x) + + # 2^(floor(log2(3.0)) - E8M0_TARGET_MAX_POW2[int8]) = 2^(1 - 6) + torch.testing.assert_close(scale, torch.full((1, 1), 2.0**-5)) + assert torch.equal(zero_point, torch.zeros((1, 1), dtype=torch.int32)) + torch.testing.assert_close(minval, torch.full((1, 1), -3.0)) + + # 0.01 and 1.01 are not multiples of the scale, so they round to codes 0 and 32. + # disable_observer() switches the stateful calculator to its cached qparams. + expected = torch.tensor([[-3.0, 0.0, 1.0, 3.0]]) + torch.testing.assert_close(fake_quant(x), expected) + fake_quant.disable_observer() + torch.testing.assert_close(fake_quant(x), expected) + + def test_extra_repr(): qparams_calc = StaticQParamsCalculator( dtype=torch.int8, diff --git a/tests/quantization/test_qspec_reconcile.py b/tests/quantization/test_qspec_reconcile.py index 20b7cd1c..0d21d094 100644 --- a/tests/quantization/test_qspec_reconcile.py +++ b/tests/quantization/test_qspec_reconcile.py @@ -668,6 +668,34 @@ def test_reconciled_qscheme_reaches_the_calculator(self) -> None: spec = _build_concrete_spec(self._fields(QSCHEME=_fv(QuantizationScheme.ASYMMETRIC))) assert self._calculator(spec).qscheme is QuantizationScheme.ASYMMETRIC + @pytest.mark.parametrize( + "dtype,scale_dtype", + [(torch.int8, torch.float8_e8m0fnu), (torch.float8_e4m3fn, None)], + ids=["int8-e8m0", "fp8"], + ) + def test_op_intrinsic_asymmetric_yields_to_a_symmetric_only_spec( + self, dtype: torch.dtype, scale_dtype: torch.dtype | None + ) -> None: + """Another slot in the group can bring in a dtype or scale the op's own config lacked.""" + spec = _build_concrete_spec( + self._fields( + DTYPE=_fv(dtype), + SCALE_DTYPE=_fv(scale_dtype), + QSCHEME=_fv(QuantizationScheme.ASYMMETRIC, OP_INTRINSIC_PRIORITY), + ) + ) + assert self._calculator(spec).qscheme is QuantizationScheme.SYMMETRIC + + def test_user_asymmetric_with_an_e8m0_scale_still_raises(self) -> None: + """Only the op's preference yields; a conflicting user config still surfaces.""" + with pytest.raises(ValueError, match="requires symmetric"): + _build_concrete_spec( + self._fields( + SCALE_DTYPE=_fv(torch.float8_e8m0fnu), + QSCHEME=_fv(QuantizationScheme.ASYMMETRIC), + ) + ) + def test_reconciled_float_range_reaches_the_calculator(self) -> None: spec = _build_concrete_spec(self._fields(FLOAT_RANGE=_fv([0.0, 1.0]))) assert tuple(self._calculator(spec).float_range) == (0.0, 1.0) diff --git a/tests/quantization/test_quantization.py b/tests/quantization/test_quantization.py index f369401f..7ee667a2 100644 --- a/tests/quantization/test_quantization.py +++ b/tests/quantization/test_quantization.py @@ -28,6 +28,7 @@ from coreai_opt.quantization.spec import ( PerBlockGranularity, PerTensorGranularity, + QuantizationGranularity, QuantizationScheme, default_activation_quantization_spec, default_weight_quantization_spec, @@ -634,6 +635,126 @@ def test_finalize_rejects_per_block_activation(self, execution_mode, backend): quantizer.finalize(prepared_model, backend=backend) +class TestIntegerE8M0Quantization: + """int8 with e8m0 scales runs in torch at any granularity; built-in export rejects it.""" + + @staticmethod + def _int8_e8m0_spec(granularity: QuantizationGranularity) -> QuantizationSpec: + return QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=granularity, + scale_dtype=torch.float8_e8m0fnu, + ) + + @staticmethod + def _e8m0_fake_quants(model: nn.Module) -> list[FakeQuantizeImplBase]: + return [ + m + for m in model.modules() + if isinstance(m, FakeQuantizeImplBase) + and m.qparams_calculator.scale_dtype == torch.float8_e8m0fnu + ] + + @classmethod + def _make_config(cls, execution_mode: str, target: str) -> QuantizerConfig: + """int8 with e8m0 scales on per-block weights or on per-tensor activations.""" + is_weight = target == "weight" + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": cls._int8_e8m0_spec(PerBlockGranularity(axis=1, block_size=32)) + if is_weight + else default_weight_quantization_spec() + }, + op_input_spec=None + if is_weight + else {"*": cls._int8_e8m0_spec(PerTensorGranularity())}, + op_output_spec=None, + ) + ).set_execution_mode(execution_mode) + + @pytest.mark.parametrize("target", ["weight", "activation"]) + def test_scales_are_powers_of_two(self, execution_mode, target): + quantizer = Quantizer(SimpleLinearModel(), self._make_config(execution_mode, target)) + prepared_model = quantizer.prepare((torch.randn(4, 64),)) + with quantizer.calibration_mode(): + prepared_model(torch.randn(4, 64)) + prepared_model(torch.randn(4, 64)) + + scales = [m.calculate_qparams()[0] for m in self._e8m0_fake_quants(prepared_model)] + assert scales, f"No e8m0 FakeQuantize modules for target={target}" + for scale in scales: + assert torch.equal(scale.log2(), scale.log2().round()) + + @pytest.mark.parametrize("target", ["weight", "activation"]) + @pytest.mark.parametrize( + "backend,is_supported", + [ + (ExportBackend.CoreAI, False), + (ExportBackend.CoreML, False), + (ExportBackend._TORCH, True), + ], + ) + def test_finalize_rejects_non_torch_backends( + self, execution_mode, target, backend, is_supported + ): + quantizer = Quantizer(SimpleLinearModel(), self._make_config(execution_mode, target)) + prepared_model = quantizer.prepare((torch.randn(4, 64),)) + prepared_model(torch.randn(4, 64)) + + if is_supported: + assert quantizer.finalize(prepared_model, backend=backend) is not None + else: + with pytest.raises(NotImplementedError, match="float8_e8m0fnu"): + quantizer.finalize(prepared_model, backend=backend) + + def test_graph_relu_output_keeps_symmetric_scheme(self): + """relu's asymmetric proposal yields to the symmetric scheme an e8m0 spec requires.""" + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": default_weight_quantization_spec()}, + op_input_spec=None, + op_output_spec={"*": self._int8_e8m0_spec(PerTensorGranularity())}, + ) + ).set_execution_mode("graph") + model = nn.Sequential(nn.Linear(64, 64), nn.ReLU()) + + prepared_model = Quantizer(model, config).prepare((torch.randn(4, 64),)) + + e8m0_qschemes = {m.qscheme for m in self._e8m0_fake_quants(prepared_model)} + assert e8m0_qschemes == {QuantizationScheme.SYMMETRIC} + + def test_graph_shared_observer_keeps_symmetric_scheme(self): + """sigmoid's asymmetric proposal yields to the e8m0 scale its consumer brings in.""" + int8_spec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerTensorGranularity(), + ) + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": default_weight_quantization_spec()}, + op_input_spec=None, + op_output_spec={"*": int8_spec}, + ), + module_name_configs={ + "2": ModuleQuantizerConfig( + op_state_spec={"weight": default_weight_quantization_spec()}, + op_input_spec={"*": self._int8_e8m0_spec(PerTensorGranularity())}, + op_output_spec=None, + ) + }, + ).set_execution_mode("graph") + # sigmoid's output and the second Linear's input share one observer. + model = nn.Sequential(nn.Linear(64, 64), nn.Sigmoid(), nn.Linear(64, 64)) + + prepared_model = Quantizer(model, config).prepare((torch.randn(4, 64),)) + + e8m0_qschemes = {m.qscheme for m in self._e8m0_fake_quants(prepared_model)} + assert e8m0_qschemes == {QuantizationScheme.SYMMETRIC} + + class TestSharedWeightQuantization: class _LeafA(nn.Module): def __init__(self): diff --git a/tests/quantization/test_quantization_spec.py b/tests/quantization/test_quantization_spec.py index a4d845e5..1a0ec6b1 100644 --- a/tests/quantization/test_quantization_spec.py +++ b/tests/quantization/test_quantization_spec.py @@ -876,6 +876,7 @@ def test_partial_defaults(): ), (torch.float4_e2m1fn_x2, None, torch.float8_e8m0fnu), (torch.float4_e2m1fn_x2, torch.float8_e8m0fnu, torch.float8_e8m0fnu), + (torch.int8, torch.float8_e8m0fnu, torch.float8_e8m0fnu), ], ) def test_valid_scale_dtype(dtype, scale_dtype, expected_scale_dtype): @@ -899,6 +900,7 @@ def test_valid_scale_dtype(dtype, scale_dtype, expected_scale_dtype): torch.int4, torch.float8_e8m0fnu, ), + (torch.uint8, torch.float8_e8m0fnu), ], ) def test_invalid_scale_dtype(dtype, scale_dtype): @@ -911,6 +913,65 @@ def test_invalid_scale_dtype(dtype, scale_dtype): ) +@pytest.mark.parametrize( + "qscheme,qformulation", + [ + (QuantizationScheme.ASYMMETRIC, QuantizationFormulation.ZP), + (QuantizationScheme.SYMMETRIC_WITH_CLIPPING, QuantizationFormulation.ZP), + (QuantizationScheme.SYMMETRIC, QuantizationFormulation.MINVAL), + ], +) +def test_e8m0_scale_requires_symmetric_zp(qscheme, qformulation): + with pytest.raises(ValueError, match="FP quantization or an e8m0 scale"): + QuantizationSpec( + dtype=torch.int8, + qscheme=qscheme, + qformulation=qformulation, + granularity=PerBlockGranularity(axis=1, block_size=32), + scale_dtype=torch.float8_e8m0fnu, + ) + + +@pytest.mark.parametrize( + "dtype,scale_dtype,qscheme,qformulation,match", + [ + ( + torch.uint8, + torch.float8_e8m0fnu, + QuantizationScheme.ASYMMETRIC, + QuantizationFormulation.ZP, + "scale_dtype must be None for dtype=torch.uint8", + ), + ( + torch.int4, + torch.float8_e8m0fnu, + QuantizationScheme.SYMMETRIC, + QuantizationFormulation.MINVAL, + "scale_dtype must be None for dtype=torch.int4", + ), + ( + torch.int8, + torch.float16, + QuantizationScheme.ASYMMETRIC, + QuantizationFormulation.ZP, + "Unsupported scale_dtype", + ), + ], +) +def test_scale_dtype_error_is_reported_before_qscheme_and_qformulation( + dtype, scale_dtype, qscheme, qformulation, match +): + """Fixing the qscheme or qformulation first would only lead to the scale_dtype error.""" + with pytest.raises(ValueError, match=match): + QuantizationSpec( + dtype=dtype, + qscheme=qscheme, + qformulation=qformulation, + granularity=PerBlockGranularity(axis=1, block_size=32), + scale_dtype=scale_dtype, + ) + + @pytest.mark.parametrize( "spec", [