diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 475c180..6ea9959 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -23,10 +23,11 @@ # Mapping from dtype to the largest power-of-2 component of its max value. # Used by e8m0 scale computation (OCP Microscaling spec, FLOOR mode). -FP_DTYPE_TO_MAX_POW2: dict[torch.dtype, int] = { +E8M0_TARGET_MAX_POW2: dict[torch.dtype, int] = { 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 + torch.int8: 0, # max = 127 / 64 = 1.984375 } # Constants for e8m0 scale computation. diff --git a/src/coreai_opt/quantization/_eager/_prepare_for_export.py b/src/coreai_opt/quantization/_eager/_prepare_for_export.py index bedf246..a6f26fd 100644 --- a/src/coreai_opt/quantization/_eager/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_eager/_prepare_for_export.py @@ -21,14 +21,17 @@ ) from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization._export_utils import ( - canonicalize_qparam_shape, + dequant_output_dtype, + extract_export_qparams, extract_quantization_params, + get_activation_export_handler, pack_fp4_to_float4tensor, select_export_qparams_by_formulation, + validate_activation_export_supported, validate_fp4_export, ) +from coreai_opt.quantization.config.quantization_config import ExecutionMode from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase -from coreai_opt.quantization.spec.granularity import PerBlockGranularity logger = logging.getLogger(__name__) @@ -169,11 +172,33 @@ def _import_coreai_torch_modules(): fq_id_to_dequant_mod[id(fake_quant_mod)] = weight_dequant_mod -def _process_activation_quantization(model: nn.Module): +def _build_activation_replacement_module( + fake_quant_module: FakeQuantizeImplBase, +) -> nn.Module: """ - Replace FakeQuantizeImplBase modules with ActivationQuantizeParametrization - followed by ActivationDequantizeParametrization. + Build the module that replaces one activation fake quantization module. + + Mirrors the graph path's ``_process_mlir_activation_quantization``: a registered + handler for the granularity takes precedence, and the built-in axis-based path is + the fallback. Only the shape differs -- graph mode rewrites the node in place, + while this returns a module for the caller to swap in. + + Args: + fake_quant_module (FakeQuantizeImplBase): The activation fake quantization + module being replaced. + + Returns: + nn.Module: The replacement, a ``Sequential`` of a quantize and a dequantize + module for the built-in path, or whatever a registered handler returns. """ + # A registered handler owns this granularity end to end, including its own + # dtype and shape validation, and returns the replacement module; the + # caller's swap is shared. + handler = get_activation_export_handler(fake_quant_module.granularity, ExecutionMode.EAGER) + if handler is not None: + return handler(fake_quant_module) # type: ignore[call-arg,return-value] + + validate_activation_export_supported(fake_quant_module) # Lazy import: coreai_torch is required for MLIR export def _import_coreai_torch_modules(): @@ -188,97 +213,69 @@ def _import_coreai_torch_modules(): _import_coreai_torch_modules ) + scale, zero_point, minval = extract_export_qparams(fake_quant_module) + + axis = fake_quant_module.qparams_calculator._resolved_axis + axis = axis if axis is not None else 0 + + # Create the replacement module sequence + # First: ActivationQuantizeParametrization + quant_module = ActivationQuantizeModule( + scale=scale, + output_dtype=fake_quant_module.dtype, + zero_point=zero_point.clone() if zero_point is not None else None, + minval=minval.clone() if minval is not None else None, + axis=axis, + ) + + # Pass input_dtype for integer quantization + # needed for determining n_bits for subbyte (eg. int4) quantization + input_dtype = fake_quant_module.dtype if not fake_quant_module.dtype.is_floating_point else None + + # Second: ActivationDequantizeParametrization + dequant_module = ActivationDequantizeModule( + scale=scale.clone(), # make one more copy for dequantize module buffers + zero_point=zero_point.clone() if zero_point is not None else None, + minval=minval.clone() if minval is not None else None, + axis=axis, + input_dtype=input_dtype, + output_dtype=dequant_output_dtype(fake_quant_module), + ) + + # Create a sequential module to combine both operations + return nn.Sequential( + OrderedDict( + [ + ("quantize", quant_module), + ("dequantize", dequant_module), + ] + ) + ) + + +def _process_activation_quantization(model: nn.Module): + """ + Replace FakeQuantizeImplBase modules with ActivationQuantizeParametrization + followed by ActivationDequantizeParametrization. + """ modules_to_replace = [] # Collect modules that need to be replaced # If there are duplicated fake quant modules they should ideally have duplicated # parent modules as well, unless manually inserted. Since we don't yet support # manual insertion, we aren't handling duplicates separately (remove_duplicate=True) + # + # Collect before mutating: named_modules() walks the live tree, so replacing a + # module mid-iteration would invalidate the traversal. for name, module in list(model.named_modules(remove_duplicate=True)): if isinstance(module, FakeQuantizeImplBase) and module.quantization_target in ( CompressionTargetTensor.ACTIVATION, ): - if is_float4_dtype(module.dtype): - raise ValueError("Core AI export does not support FP4 activation quantization.") - if isinstance(module.granularity, PerBlockGranularity): - raise ValueError( - "Core AI export does not support PerBlockGranularity on activations." - ) modules_to_replace.append((name, module)) # Replace each FakeQuantizeImplBase module for name, fake_quant_module in modules_to_replace: - # Extract quantization parameters - scale, zero_point, minval = extract_quantization_params(fake_quant_module) - - # Drop one of the offsets so that the export - # module / runtime selects the right dequant path. - zero_point, minval = select_export_qparams_by_formulation( - fake_quant_module, zero_point, minval - ) - - # Cast scale and minval to appropriate dtype for MLIR backend inference - _compute_dtype_for_export = fake_quant_module.qparams_calculator._compute_dtype_for_export - scale = scale.to(dtype=_compute_dtype_for_export) - if minval is not None: - minval = minval.to(dtype=_compute_dtype_for_export) - - if fake_quant_module.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: - scale = scale.to(torch.float8_e8m0fnu) - - # Canonicalize scale/zero_point/minval to 0-D (per-tensor) or 1-D (per-channel) - granularity = fake_quant_module.granularity - scale = canonicalize_qparam_shape(scale, granularity) - if zero_point is not None: - zero_point = canonicalize_qparam_shape(zero_point, granularity) - if minval is not None: - minval = canonicalize_qparam_shape(minval, granularity) - - axis = fake_quant_module.qparams_calculator._resolved_axis - axis = axis if axis is not None else 0 - - # Create the replacement module sequence - # First: ActivationQuantizeParametrization - quant_module = ActivationQuantizeModule( - scale=scale, - output_dtype=fake_quant_module.dtype, - zero_point=zero_point.clone() if zero_point is not None else None, - minval=minval.clone() if minval is not None else None, - axis=axis, - ) - - # Second: ActivationDequantizeParametrization - if fake_quant_module.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: - output_dtype = _compute_dtype_for_export - else: - output_dtype = None - - # Pass input_dtype for integer quantization - # needed for determining n_bits for subbyte (eg. int4) quantization - input_dtype = ( - fake_quant_module.dtype if not fake_quant_module.dtype.is_floating_point else None - ) - - dequant_module = ActivationDequantizeModule( - scale=scale.clone(), # make one more copy for dequantize module buffers - zero_point=zero_point.clone() if zero_point is not None else None, - minval=minval.clone() if minval is not None else None, - axis=axis, - input_dtype=input_dtype, - output_dtype=output_dtype, - ) - - # Create a sequential module to combine both operations - replacement_module = nn.Sequential( - OrderedDict( - [ - ("quantize", quant_module), - ("dequantize", dequant_module), - ] - ) - ) - - # Replace the module in the model + replacement_module = _build_activation_replacement_module(fake_quant_module) parent_module, attr_name = get_parent_module_and_attr_name(model, name) setattr(parent_module, attr_name, replacement_module) diff --git a/src/coreai_opt/quantization/_export_utils.py b/src/coreai_opt/quantization/_export_utils.py index c705819..6aebf34 100644 --- a/src/coreai_opt/quantization/_export_utils.py +++ b/src/coreai_opt/quantization/_export_utils.py @@ -12,11 +12,15 @@ from __future__ import annotations from collections import OrderedDict +from collections.abc import Callable import torch from torch import nn +from coreai_opt._utils.torch_utils import is_float4_dtype +from coreai_opt.common import ExportBackend from coreai_opt.config.spec import CompressionTargetTensor +from coreai_opt.quantization.config.quantization_config import ExecutionMode from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase from coreai_opt.quantization.spec.granularity import ( PerBlockGranularity, @@ -139,8 +143,8 @@ def create_mil_act_quant_seq( ), ), ("dequantize", _MILActivationDequantizeModule()), - ], - ), + ] + ) ) @@ -219,6 +223,176 @@ def select_export_qparams_by_formulation( raise NotImplementedError(f"Unknown qformulation: {fake_quant_mod.qformulation}") +def extract_export_qparams( + fake_quant_mod: FakeQuantizeImplBase, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Extract scale / zero_point / minval, cast and reshape them for Core AI export. + + Args: + fake_quant_mod (FakeQuantizeImplBase): The fake quantization module. + + Returns: + tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: The + ``(scale, zero_point, minval)`` triple, cast for export and canonicalized + to 0-D (per-tensor) or 1-D (per-channel). At most one of ``zero_point`` + and ``minval`` is not None. + """ + scale, zero_point, minval = extract_quantization_params(fake_quant_mod) + zero_point, minval = select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval) + + compute_dtype = fake_quant_mod.qparams_calculator._compute_dtype_for_export + scale = scale.to(dtype=compute_dtype) + if minval is not None: + minval = minval.to(dtype=compute_dtype) + + if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: + scale = scale.to(torch.float8_e8m0fnu) + + granularity = fake_quant_mod.granularity + scale = canonicalize_qparam_shape(scale, granularity) + if zero_point is not None: + zero_point = canonicalize_qparam_shape(zero_point, granularity) + if minval is not None: + minval = canonicalize_qparam_shape(minval, granularity) + + return scale, zero_point, minval + + +def dequant_output_dtype(fake_quant_mod: FakeQuantizeImplBase) -> torch.dtype | None: + """Return the dequantize op's ``output_dtype``, or None to let it infer one. + + An e8m0 scale carries no dtype the consumer can infer a compute dtype from, + so the dequantize op has to be told which one to produce. + """ + if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: + return fake_quant_mod.qparams_calculator._compute_dtype_for_export + return None + + +# ────────────────────────────────────────────────────────────────────── +# Activation export handlers +# ────────────────────────────────────────────────────────────────────── + +# Emits the quantize/dequantize pair for one fake-quant node and replaces it. +GraphActivationExportHandler = Callable[ + [torch.fx.GraphModule, torch.fx.Node, FakeQuantizeImplBase], + None, +] + +# Builds the replacement module for one fake-quant module. The caller swaps it in. +EagerActivationExportHandler = Callable[[FakeQuantizeImplBase], nn.Module] + +_GRAPH_ACTIVATION_EXPORT_HANDLERS: dict[ + type[QuantizationGranularity], + GraphActivationExportHandler, +] = {} +_EAGER_ACTIVATION_EXPORT_HANDLERS: dict[ + type[QuantizationGranularity], + EagerActivationExportHandler, +] = {} + + +def register_graph_activation_export_handler( + granularity_cls: type[QuantizationGranularity], + handler: GraphActivationExportHandler, +) -> None: + """Register a graph-mode activation export handler for a granularity. + + Args: + granularity_cls (type[QuantizationGranularity]): Granularity type the + handler applies to, matched exactly rather than by subclass. + handler (GraphActivationExportHandler): Called with + ``(model, node, fake_quant_mod)``; must replace ``node``. + """ + _GRAPH_ACTIVATION_EXPORT_HANDLERS[granularity_cls] = handler + + +def register_eager_activation_export_handler( + granularity_cls: type[QuantizationGranularity], + handler: EagerActivationExportHandler, +) -> None: + """Register an eager-mode activation export handler for a granularity. + + Args: + granularity_cls (type[QuantizationGranularity]): Granularity type the + handler applies to, matched exactly rather than by subclass. + handler (EagerActivationExportHandler): Called with + ``(fake_quant_mod)``; must return the replacement module. + """ + _EAGER_ACTIVATION_EXPORT_HANDLERS[granularity_cls] = handler + + +def get_activation_export_handler( + granularity: QuantizationGranularity, + execution_mode: ExecutionMode, +) -> GraphActivationExportHandler | EagerActivationExportHandler | None: + """Return the handler registered for a granularity in ``execution_mode``, if any. + + Args: + granularity (QuantizationGranularity): The activation's granularity. + execution_mode (ExecutionMode): Which registry to look in. + + Returns: + GraphActivationExportHandler | EagerActivationExportHandler | None: The + registered handler, or None when the built-in export path applies. + """ + if execution_mode == ExecutionMode.EAGER: + return _EAGER_ACTIVATION_EXPORT_HANDLERS.get(type(granularity)) + return _GRAPH_ACTIVATION_EXPORT_HANDLERS.get(type(granularity)) + + +def can_export_stateless_fake_quant( + fake_quant_mod: FakeQuantizeImplBase, + backend: ExportBackend, + execution_mode: ExecutionMode, +) -> bool: + """Return whether a recompute-every-forward calculator can still be exported. + + 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. + + Args: + fake_quant_mod (FakeQuantizeImplBase): The module whose calculator is stateless. + backend (ExportBackend): The export target. + execution_mode (ExecutionMode): The quantizer's execution mode. + + Returns: + bool: True if a registered handler covers this module. + """ + return ( + backend == ExportBackend.CoreAI + and fake_quant_mod.quantization_target == CompressionTargetTensor.ACTIVATION + and get_activation_export_handler(fake_quant_mod.granularity, execution_mode) is not None + ) + + +def validate_activation_export_supported(fake_quant_mod: FakeQuantizeImplBase) -> None: + """Reject activation configurations the built-in Core AI export path cannot express. + + Only called once no registered handler claimed the granularity, so both + messages can point at the extension hook as an alternative. + + Args: + fake_quant_mod (FakeQuantizeImplBase): The activation fake quantization module. + + Raises: + ValueError: If the dtype is FP4, or the granularity is per-block. + """ + if is_float4_dtype(fake_quant_mod.dtype): + raise ValueError("Core AI export does not support FP4 activation quantization.") + + if isinstance(fake_quant_mod.granularity, PerBlockGranularity): + raise ValueError( + "Core AI export does not support PerBlockGranularity on activations: a " + "per-block scale cannot be canonicalized to 0-D or 1-D. Use " + "PerTensorGranularity or PerChannelGranularity, export with " + "backend=ExportBackend._TORCH, or install and activate an extension " + "package that registers an activation export handler for this granularity." + ) + + def validate_qformulation_for_mil_export(fake_quant_mod: FakeQuantizeImplBase) -> None: """Reject CoreML export for non-ZP quantization formulations. @@ -352,3 +526,37 @@ def validate_fp4_export( f"the last axis, no blocking elsewhere). Got resolved block sizes " f"{resolved_block_size} from granularity={granularity!r}." ) + + +def validate_e8m0_int_export(model: nn.Module) -> None: + """Reject an integer dtype with an e8m0 scale on a weight, before anything mutates. + + ``constexpr_blockwise_shift_scale``, which the weight paths emit, leaves the scale in + e8m0 for an integer input, and the consuming matmul then sees a scale whose dtype does + not match its other operand. Activations are unaffected -- their path casts the scale. + + Runs over the whole model rather than per module: both weight paths free or erase as + they go, so raising mid-walk would leave the model half-rewritten. + + Args: + model (nn.Module): The prepared model to check. + + Raises: + ValueError: If any weight carries an integer dtype with an e8m0 scale. + """ + offenders = [ + name + for name, mod in model.named_modules() + if is_module_fake_quant_target(mod, CompressionTargetTensor.WEIGHT) + and not mod.dtype.is_floating_point + and mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu + ] + if offenders: + msg = ( + "An e8m0 scale_dtype on an integer dtype is not supported for weights during " + "Core AI export, only for activations. constexpr_blockwise_shift_scale does " + "not accept an e8m0 scale with an integer input. Affected FakeQuantize " + f"modules: {offenders}. Set scale_dtype=None for those weights, or use a " + "floating-point dtype." + ) + raise ValueError(msg) diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index 477b03f..10cc8a4 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -24,12 +24,15 @@ from coreai_opt._utils.torch_utils import is_float4_dtype, sanitize_module_name from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization._export_utils import ( - canonicalize_qparam_shape, convert_dtype_for_torch_quantize, create_mil_act_quant_seq, + dequant_output_dtype, + extract_export_qparams, extract_quantization_params, + get_activation_export_handler, pack_fp4_to_float4tensor, select_export_qparams_by_formulation, + validate_activation_export_supported, validate_fp4_export, validate_qformulation_for_mil_export, ) @@ -37,8 +40,8 @@ remove_fake_quant_module, resolve_attr, ) +from coreai_opt.quantization.config.quantization_config import ExecutionMode from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase -from coreai_opt.quantization.spec.granularity import PerBlockGranularity logger = logging.getLogger(__name__) @@ -332,11 +335,14 @@ def _process_mlir_activation_quantization( node: The fake quantization node to replace fake_quant_mod: The fake quantization module """ - if is_float4_dtype(fake_quant_mod.dtype): - raise ValueError("Core AI export does not support FP4 activation quantization.") + # A registered handler owns this granularity end to end, including its own + # dtype and shape validation. + handler = get_activation_export_handler(fake_quant_mod.granularity, ExecutionMode.GRAPH) + if handler is not None: + handler(model, node, fake_quant_mod) # type: ignore[call-arg] + return - if isinstance(fake_quant_mod.granularity, PerBlockGranularity): - raise ValueError("Core AI export does not support PerBlockGranularity on activations.") + validate_activation_export_supported(fake_quant_mod) def _import_coreai_custom_ops(): import coreai_torch._compression.custom_layers # noqa: PLC0415, F401 @@ -349,29 +355,7 @@ def _import_coreai_custom_ops(): if not node.args: raise ValueError(f"Node {node} has no input arguments") - # Extract and prepare quantization parameters - scale, zero_point, minval = extract_quantization_params(fake_quant_mod) - - # Drop the offset the active formulation doesn't consume so the runtime op - # selects the right dequant path. - zero_point, minval = select_export_qparams_by_formulation(fake_quant_mod, zero_point, minval) - - # Cast scale and minval to appropriate dtype for MLIR backend inference - _compute_dtype_for_export = fake_quant_mod.qparams_calculator._compute_dtype_for_export - scale = scale.to(dtype=_compute_dtype_for_export) - if minval is not None: - minval = minval.to(dtype=_compute_dtype_for_export) - - if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: - scale = scale.to(torch.float8_e8m0fnu) - - # Canonicalize scale/zero_point/minval to 0-D (per-tensor) or 1-D (per-channel) - granularity = fake_quant_mod.granularity - scale = canonicalize_qparam_shape(scale, granularity) - if zero_point is not None: - zero_point = canonicalize_qparam_shape(zero_point, granularity) - if minval is not None: - minval = canonicalize_qparam_shape(minval, granularity) + scale, zero_point, minval = extract_export_qparams(fake_quant_mod) # Register buffers and get buffer names base_name = node.name.replace(".", "_") @@ -383,10 +367,7 @@ def _import_coreai_custom_ops(): axis = fake_quant_mod.qparams_calculator._resolved_axis # Determine output_dtype for dequantize (needed for FP8 when scale is float8_e8m0fnu) - if fake_quant_mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu: - dequant_output_dtype = _compute_dtype_for_export - else: - dequant_output_dtype = None + output_dtype = dequant_output_dtype(fake_quant_mod) # Create graph nodes and replace fake quantization with model.graph.inserting_before(node): @@ -413,7 +394,7 @@ def _import_coreai_custom_ops(): # coreai.dequantize(input, scale, zero_point=, minval=, axis=, input_dtype=, output_dtype=) dequant_args = (quantize_node, scale_node) - dequant_kwargs = {"output_dtype": dequant_output_dtype} + dequant_kwargs = {"output_dtype": output_dtype} if zp_node is not None: # output = scale * (input - zero_point) dequant_kwargs["zero_point"] = zp_node diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py index 627cf14..88e9176 100644 --- a/src/coreai_opt/quantization/quantizer.py +++ b/src/coreai_opt/quantization/quantizer.py @@ -26,6 +26,10 @@ from coreai_opt._utils.torch_utils import get_module_name as _get_module_name 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, + validate_e8m0_int_export as _validate_e8m0_int_export, +) from coreai_opt.quantization._graph import GraphQuantizer as _GraphQuantizer from coreai_opt.quantization.base_quantizer import _BaseQuantizer from coreai_opt.quantization.config.quantization_config import ( @@ -414,6 +418,8 @@ def _validate_no_persistent_observer_calculators( ) -> None: """Reject CoreAI/CoreML export when any qparams calculator is a ``StatelessQParamsCalculatorBase`` (e.g. dynamic quantization). + + A registered activation export handler lifts the restriction. """ if backend == ExportBackend._TORCH: return @@ -423,6 +429,7 @@ def _validate_no_persistent_observer_calculators( for name, mod in model_to_check.named_modules() if isinstance(mod, FakeQuantizeImplBase) and isinstance(mod.qparams_calculator, StatelessQParamsCalculatorBase) + and not _can_export_stateless_fake_quant(mod, backend, self._execution_mode) ] if stateless_fq_names: raise NotImplementedError( @@ -481,6 +488,8 @@ def finalize( """ self._validate_mmap_dir_constraints(model, backend, mmap_dir) self._validate_no_persistent_observer_calculators(model, backend) + if backend == ExportBackend.CoreAI: + _validate_e8m0_int_export(model if model is not None else self._model) 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 5260302..fd1f8ac 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -19,8 +19,8 @@ ) from coreai_opt._utils.torch_utils import ( E8M0_EXPONENT_BIAS as _E8M0_EXPONENT_BIAS, + E8M0_TARGET_MAX_POW2 as _E8M0_TARGET_MAX_POW2, F32_MIN_NORMAL as _F32_MIN_NORMAL, - FP_DTYPE_TO_MAX_POW2 as _FP_DTYPE_TO_MAX_POW2, ) from .granularity import QuantizationGranularity @@ -174,11 +174,11 @@ def _compute_e8m0_scale(self, max_abs: torch.Tensor) -> torch.Tensor: - torchao implementation: https://github.com/pytorch/ao/blob/main/torchao/prototype/mx_formats/mx_tensor.py """ - target_max_pow2 = _FP_DTYPE_TO_MAX_POW2.get(self.dtype) + target_max_pow2 = _E8M0_TARGET_MAX_POW2.get(self.dtype) if target_max_pow2 is None: raise ValueError( f"Unsupported dtype for e8m0 scale computation: {self.dtype}. " - f"Supported: {list(_FP_DTYPE_TO_MAX_POW2.keys())}" + f"Supported: {list(_E8M0_TARGET_MAX_POW2.keys())}" ) max_abs_fp32 = max_abs.to(torch.float32) @@ -219,17 +219,23 @@ def _compute_scale_zero_point_minval( constrained to powers of 2 following the OCP Microscaling (MX) specification (FLOOR mode): scale = 2^(floor(log2(max_abs)) - target_max_pow2) - where ``target_max_pow2`` is the largest power-of-2 component of the - target dtype's maximum representable value: - - FP4 E2M1: max = 6.0 = 1.5 * 2^2, target_max_pow2 = 2 - - FP8 E4M3: max = 448.0 = 1.75 * 2^8, target_max_pow2 = 8 - - FP8 E5M2: max = 57344.0 = 1.75 * 2^15, target_max_pow2 = 15 + where ``target_max_pow2`` comes from ``E8M0_TARGET_MAX_POW2``, which + is the single source of truth for which dtypes support an e8m0 scale. """ # 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 + # An integer dtype with an e8m0 scale is symmetric with a zero offset, but + # the integer quantize path indexes both offsets, so return zeros not None. + return ( + scale, + torch.zeros_like(scale, dtype=torch.int32), + torch.zeros_like(scale), + ) # 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 c39f3a1..e3660f0 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -19,6 +19,7 @@ ) from coreai_opt._utils.torch_utils import ( + E8M0_TARGET_MAX_POW2 as _E8M0_TARGET_MAX_POW2, get_n_bits_from_dtype as _get_n_bits_from_dtype, is_float4_dtype as _is_float4_dtype, ) @@ -26,6 +27,7 @@ from .fake_quantize import FakeQuantizeImplBase from .granularity import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationGranularity, @@ -316,7 +318,9 @@ class type: MinMaxRangeCalculator or custom registered class type (defaults to e8m0) - FP8 (float8_e4m3fn, float8_e5m2): scale_dtype must be torch.float8_e8m0fnu or None (defaults to None) - - Integer dtypes: scale_dtype must be None (defaults to None) + - int8: scale_dtype may be torch.float8_e8m0fnu, but only together with + PerBlockGranularity, qscheme=SYMMETRIC and qformulation=ZP + - Every other integer dtype: scale_dtype must be None (defaults to None) Default: None @@ -537,21 +541,54 @@ def validate_scale_dtype(self) -> QuantizationSpec: Rules: - Only None or torch.float8_e8m0fnu are supported. - - Integer dtypes: scale_dtype must be None. - - FP8 dtypes: scale_dtype may be None or torch.float8_e8m0fnu. + - The dtype must have an e8m0 target exponent, i.e. be a key of + ``E8M0_TARGET_MAX_POW2``. + - Integer dtypes additionally require PerBlockGranularity, + qscheme=SYMMETRIC and qformulation=ZP. - FP4 dtypes: scale_dtype is resolved to torch.float8_e8m0fnu by resolve_scale_dtype (before validator). """ - if self.scale_dtype is not None and self.scale_dtype != torch.float8_e8m0fnu: + if self.scale_dtype is None: + return self + + if self.scale_dtype != torch.float8_e8m0fnu: raise ValueError( f"Unsupported scale_dtype: {self.scale_dtype}. " f"Only None or torch.float8_e8m0fnu are supported." ) - if not self.dtype.is_floating_point and self.scale_dtype is not None: + if self.dtype not in _E8M0_TARGET_MAX_POW2: + raise ValueError( + 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)}." + ) + + if self.dtype.is_floating_point: + return self + + # An integer dtype with an e8m0 scale carries no offsets, so it can only + # express a symmetric, block-wise, zero-point range. + if not isinstance(self.granularity, PerBlockGranularity): + raise ValueError( + f"An e8m0 scale_dtype on {self.dtype} requires PerBlockGranularity, " + f"got {type(self.granularity).__name__}. The per-tensor and " + f"per-channel export path divides by the scale without casting it, " + f"which torch cannot do for an fp8 type." + ) + if self.qscheme != QuantizationScheme.SYMMETRIC: + raise ValueError( + f"An e8m0 scale_dtype on {self.dtype} requires " + f"qscheme=QuantizationScheme.SYMMETRIC, got qscheme={self.qscheme}. " + f"This configuration carries no offsets, so it cannot express an " + f"asymmetric range." + ) + if self.qformulation != QuantizationFormulation.ZP: raise ValueError( - f"scale_dtype must be None for integer dtypes, " - f"got scale_dtype={self.scale_dtype} with dtype={self.dtype}." + f"An e8m0 scale_dtype on {self.dtype} requires " + f"qformulation=QuantizationFormulation.ZP, got " + f"qformulation={self.qformulation}. This configuration carries no " + f"offsets, so it cannot express a minimum-value bias." ) return self diff --git a/tests/export/export_utils.py b/tests/export/export_utils.py index 02d9a97..5932aad 100644 --- a/tests/export/export_utils.py +++ b/tests/export/export_utils.py @@ -24,6 +24,7 @@ from coremltools import ComputeUnit from coreai_opt import CoreMLExportError, ExportBackend +from coreai_opt.quantization import Quantizer, QuantizerConfig from tests.test_utils.general import verify_snr_psnr as _verify_snr_psnr if platform.system() == "Darwin": @@ -641,3 +642,67 @@ def convert_and_verify( ) return converted_model + + +def run_export_test( + model: torch.nn.Module, + input_data: torch.Tensor, + config: QuantizerConfig, + expected_ops: Mapping[str, int], + export_backend: ExportBackend, + model_dtype: torch.dtype | None = None, + calibrate: bool = False, + externalized_model: torch.nn.Module | None = None, + snr_thresh: float = 20.0, + psnr_thresh: float = 22.0, +) -> None: + """Quantize, finalize, export and verify a model against its prepared forward. + + The whole export test workflow, in the order it has to run: the prepared + forward happens before finalize, because finalize is what replaces the + fake-quantize modules with the export ops. The execution mode comes from + ``config``, so one function serves eager and graph mode. + + Args: + model: PyTorch model to quantize and export + input_data: Input tensor for model + config: Quantization configuration, carrying the execution mode + expected_ops: Expected operation counts in converted model + export_backend: Target inference stack (CoreML or CoreAI) + model_dtype: Model dtype (float16, float32, bfloat16, or None for no conversion) + calibrate: If True, run one calibration pass under + ``quantizer.calibration_mode()`` before the reference forward. + externalized_model: The model patched in place by + ``coreai_torch._patch_model_for_externalization``. Only supported by the + CoreAI backend. + snr_thresh: Minimum acceptable SNR value + psnr_thresh: Minimum acceptable PSNR value + + """ + if model_dtype is not None: + model = model.to(dtype=model_dtype) + input_data = input_data.to(dtype=model_dtype) + + model.eval() + quantizer = Quantizer(model, config) + prepared_model = quantizer.prepare((input_data,)) + + if calibrate: + with quantizer.calibration_mode(), torch.no_grad(): + prepared_model(input_data) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_data) + + finalized_model = quantizer.finalize(backend=export_backend) + + convert_and_verify( + finalized_model=finalized_model, + input_data=input_data, + expected_ops=expected_ops, + export_backend=export_backend, + prepared_model_output=prepared_model_output, + externalized_model=externalized_model, + snr_thresh=snr_thresh, + psnr_thresh=psnr_thresh, + ) diff --git a/tests/export/test_eager_mil_export.py b/tests/export/test_eager_mil_export.py index ffe9835..6398c7f 100644 --- a/tests/export/test_eager_mil_export.py +++ b/tests/export/test_eager_mil_export.py @@ -9,7 +9,7 @@ import torch from coreai_opt import CoreMLExportError, ExportBackend -from coreai_opt.quantization import Quantizer, QuantizerConfig +from coreai_opt.quantization import QuantizerConfig from tests.fixtures.quantization import ( COREML_ACT_REJECT_DTYPES, COREML_WEIGHT_REJECT_DTYPES, @@ -27,35 +27,18 @@ def _run_eager_mil_export_test_ex( expected_ops: Mapping[str, int], model_dtype: torch.dtype | None = None, ) -> None: - """Run eager CoreML export test with expanded configuration parameters. - - Args: - model: PyTorch model to quantize and export - input_data: Input tensor for model - config: Eager quantization configuration - model_dtype: Model dtype (float16, float32, bfloat16, or None for no conversion) - expected_ops: Expected operation counts in converted model + """Run the shared export test workflow against the CoreML backend. + CoreML quantizes to a coarser representation than Core AI, so the SNR/PSNR + thresholds are looser than :func:`export_utils.run_export_test`'s defaults. """ - if model_dtype is not None: - model = model.to(dtype=model_dtype) - input_data = input_data.to(dtype=model_dtype) - - model.eval() - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare((input_data,)) - - with torch.no_grad(): - prepared_model_output = prepared_model(input_data) - - finalized_model = quantizer.finalize(backend=ExportBackend.CoreML) - - export_utils.convert_and_verify( - finalized_model=finalized_model, + export_utils.run_export_test( + model=model, input_data=input_data, + config=config, expected_ops=expected_ops, export_backend=ExportBackend.CoreML, - prepared_model_output=prepared_model_output, + model_dtype=model_dtype, snr_thresh=18.0, psnr_thresh=35.0, ) diff --git a/tests/export/test_eager_mlir_export.py b/tests/export/test_eager_mlir_export.py index 6671bee..14624f2 100644 --- a/tests/export/test_eager_mlir_export.py +++ b/tests/export/test_eager_mlir_export.py @@ -36,34 +36,14 @@ def _run_eager_mlir_export_test_ex( expected_ops: Mapping[str, int], model_dtype: torch.dtype | None = None, ) -> None: - """Run eager Core AI export test with expanded configuration parameters. - - Args: - model: PyTorch model to quantize and export - input_data: Input tensor for model - config: Eager quantization configuration - model_dtype: Model dtype (float16, float32, bfloat16, or None for no conversion) - expected_ops: Expected operation counts in converted model - """ - if model_dtype is not None: - model = model.to(dtype=model_dtype) - input_data = input_data.to(dtype=model_dtype) - - model.eval() - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare((input_data,)) - - with torch.no_grad(): - prepared_model_output = prepared_model(input_data) - - finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI) - - export_utils.convert_and_verify( - finalized_model=finalized_model, + """Run the shared export test workflow against the Core AI backend.""" + export_utils.run_export_test( + model=model, input_data=input_data, + config=config, expected_ops=expected_ops, export_backend=ExportBackend.CoreAI, - prepared_model_output=prepared_model_output, + model_dtype=model_dtype, ) diff --git a/tests/export/test_graph_mode_mlir_export.py b/tests/export/test_graph_mode_mlir_export.py index 7de0585..a172db8 100644 --- a/tests/export/test_graph_mode_mlir_export.py +++ b/tests/export/test_graph_mode_mlir_export.py @@ -51,42 +51,15 @@ def _run_graph_mode_mlir_export_test_ex( calibrate: bool = False, externalized_model: torch.nn.Module | None = None, ) -> None: - """Run graph-mode Core AI export test with expanded configuration parameters. - - Args: - model: PyTorch model to quantize and export - input_data: Input tensor for model - config: graph-mode quantization configuration - model_dtype: Model dtype (float16, float32, bfloat16, or None for no conversion) - expected_ops: Expected operation counts in converted model - calibrate: If True, run one calibration pass under - ``quantizer.calibration_mode()`` before the reference forward. - externalized_model: The model patched in place by - ``coreai_torch._patch_model_for_externalization``. - """ - if model_dtype is not None: - model = model.to(dtype=model_dtype) - input_data = input_data.to(dtype=model_dtype) - - model.eval() - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare((input_data,)) - - if calibrate: - with quantizer.calibration_mode(), torch.no_grad(): - prepared_model(input_data) - - with torch.no_grad(): - prepared_model_output = prepared_model(input_data) - - finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI) - - export_utils.convert_and_verify( - finalized_model=finalized_model, + """Run the shared export test workflow against the Core AI backend.""" + export_utils.run_export_test( + model=model, input_data=input_data, + config=config, expected_ops=expected_ops, export_backend=ExportBackend.CoreAI, - prepared_model_output=prepared_model_output, + model_dtype=model_dtype, + calibrate=calibrate, externalized_model=externalized_model, ) diff --git a/tests/quantization/test_activation_export_handlers.py b/tests/quantization/test_activation_export_handlers.py new file mode 100644 index 0000000..22f7c3e --- /dev/null +++ b/tests/quantization/test_activation_export_handlers.py @@ -0,0 +1,108 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the activation export handler registries.""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +from torch import nn + +from coreai_opt.quantization import ( + _export_utils, +) +from coreai_opt.quantization._export_utils import ( + canonicalize_qparam_shape, + get_activation_export_handler, + register_eager_activation_export_handler, + register_graph_activation_export_handler, + validate_activation_export_supported, +) +from coreai_opt.quantization.config.quantization_config import ExecutionMode +from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase +from coreai_opt.quantization.spec.granularity import ( + PerBlockGranularity, + PerChannelGranularity, + PerTensorGranularity, +) + + +@pytest.fixture(autouse=True) +def _restore_registries(monkeypatch: pytest.MonkeyPatch) -> None: + """Give each test empty registries, and restore the originals afterwards. + + The registries are module-level dicts, so a test that registers a handler + would otherwise change what every later test sees. Starting from empty rather + than from the live dicts also keeps these assertions independent of what + earlier files registered. + """ + monkeypatch.setattr(_export_utils, "_GRAPH_ACTIVATION_EXPORT_HANDLERS", {}) + monkeypatch.setattr(_export_utils, "_EAGER_ACTIVATION_EXPORT_HANDLERS", {}) + + +def test_graph_handler_is_returned_for_its_granularity_only() -> None: + """Lookup matches on granularity type and leaves other granularities alone.""" + + def handler(model: Any, node: Any, fake_quant_mod: Any) -> None: + raise AssertionError("should not be called") + + register_graph_activation_export_handler(PerBlockGranularity, handler) + + graph = ExecutionMode.GRAPH + assert get_activation_export_handler(PerBlockGranularity(block_size=32), graph) is handler + assert get_activation_export_handler(PerTensorGranularity(), graph) is None + assert get_activation_export_handler(PerChannelGranularity(axis=1), graph) is None + + +def test_eager_handler_is_returned_for_its_granularity_only() -> None: + """The eager registry is independent of the graph one.""" + + def handler(fake_quant_mod: Any) -> nn.Module: + return nn.Identity() + + register_eager_activation_export_handler(PerBlockGranularity, handler) + + per_block = PerBlockGranularity(block_size=32) + assert get_activation_export_handler(per_block, ExecutionMode.EAGER) is handler + assert get_activation_export_handler(PerTensorGranularity(), ExecutionMode.EAGER) is None + # Registering for eager must not register for graph. + assert get_activation_export_handler(per_block, ExecutionMode.GRAPH) is None + + +def test_per_block_export_error_names_the_caller_alternatives() -> None: + """With no handler registered, per-block activations fail with actionable options. + + This is reached from ``finalize``, after a model is prepared and calibrated, so the + message has to tell a caller what to do -- switch granularity, switch backend, or + activate an extension -- not only name the developer-facing registration hook. + """ + # validate_activation_export_supported only reads dtype and granularity. + fake_quant_mod = cast( + FakeQuantizeImplBase, + SimpleNamespace(dtype=torch.int8, granularity=PerBlockGranularity(block_size=32)), + ) + + with pytest.raises(ValueError, match="does not support PerBlockGranularity") as e: + validate_activation_export_supported(fake_quant_mod) + + msg = str(e.value) + for alternative in ("PerChannelGranularity", "ExportBackend._TORCH", "extension"): + assert alternative in msg, f"the error should offer {alternative}: {msg}" + + +@pytest.mark.parametrize( + ("qparam", "granularity", "expected_shape"), + [ + (torch.ones(1, 1), PerTensorGranularity(), ()), + (torch.ones(1, 4), PerChannelGranularity(axis=1), (4,)), + ], +) +def test_canonicalize_still_handles_supported_granularities( + qparam: torch.Tensor, granularity: Any, expected_shape: tuple[int, ...] +) -> None: + """The built-in path is unchanged by the registry.""" + assert canonicalize_qparam_shape(qparam, granularity).shape == expected_shape