diff --git a/src/coreai_opt/quantization/_eager/_prepare_for_export.py b/src/coreai_opt/quantization/_eager/_prepare_for_export.py index 0b394ae..bedf246 100644 --- a/src/coreai_opt/quantization/_eager/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_eager/_prepare_for_export.py @@ -28,6 +28,7 @@ validate_fp4_export, ) from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase +from coreai_opt.quantization.spec.granularity import PerBlockGranularity logger = logging.getLogger(__name__) @@ -198,7 +199,11 @@ def _import_coreai_torch_modules(): CompressionTargetTensor.ACTIVATION, ): if is_float4_dtype(module.dtype): - raise ValueError("FP4 activation quantization is not supported for MLIR export.") + 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 diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index d7f3aca..477b03f 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -38,6 +38,7 @@ resolve_attr, ) from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase +from coreai_opt.quantization.spec.granularity import PerBlockGranularity logger = logging.getLogger(__name__) @@ -332,7 +333,10 @@ def _process_mlir_activation_quantization( fake_quant_mod: The fake quantization module """ if is_float4_dtype(fake_quant_mod.dtype): - raise ValueError("FP4 activation quantization is not supported for MLIR export.") + 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.") def _import_coreai_custom_ops(): import coreai_torch._compression.custom_layers # noqa: PLC0415, F401 diff --git a/src/coreai_opt/quantization/_graph/_utils.py b/src/coreai_opt/quantization/_graph/_utils.py index c5988d8..c02a1df 100644 --- a/src/coreai_opt/quantization/_graph/_utils.py +++ b/src/coreai_opt/quantization/_graph/_utils.py @@ -11,6 +11,7 @@ from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase from coreai_opt.quantization.spec.granularity import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationGranularity, @@ -267,8 +268,8 @@ def _shared_granularity_axis_is_safe( """Return True only if it's proven safe to keep ``fake_quant``'s granularity shared across ``op_node``; unproven cases default to unsafe. - Per-channel activation quantization on a shared observer is assumed - unsafe by default. Each op category below has exactly one condition + Per-channel and per-block activation quantization on a shared observer are + assumed unsafe by default. Each op category below has exactly one condition under which it stops being safe, checked directly against that category rather than composing generic checks that apply to every op: @@ -284,10 +285,6 @@ def _shared_granularity_axis_is_safe( # No axis to violate — trivially safe. if isinstance(granularity, PerTensorGranularity): return True - # Anything other than PerChannelGranularity (e.g. PerBlockGranularity) - # has no condition checked below, so it's not safe. - if not isinstance(granularity, PerChannelGranularity): - return False output_shape = op_node.meta["val"].shape input_shape = input_fq_node.all_input_nodes[0].meta["val"].shape @@ -296,14 +293,19 @@ def _shared_granularity_axis_is_safe( # index, so there's no condition to check here. if len(input_shape) != len(output_shape): return False - axis = QuantizationGranularity._resolve_axis(granularity, len(input_shape)) - if axis is None: - return False + + if isinstance(granularity, PerBlockGranularity): + axes: tuple[int, ...] = tuple(range(len(input_shape))) + else: + axis = QuantizationGranularity._resolve_axis(granularity, len(input_shape)) + if axis is None: + return False + axes = (axis,) if op_node.target in _AXIS_RESIZING_ATEN_OPS: - return input_shape[axis] == output_shape[axis] + return all(input_shape[a] == output_shape[a] for a in axes) if op_node.target in _AXIS_REORDERING_ATEN_OPS: - return _op_preserves_axis_identity(op_node, axis) + return all(_op_preserves_axis_identity(op_node, a) for a in axes) # flatten/reshape/view/unsqueeze, or an unrecognized future op: no known # single condition to prove safety, so default to unsafe. return False diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index fbd1ac2..408e419 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -35,12 +35,17 @@ class QuantizationComponentFactory(CompressionComponentFactoryBase): """ @classmethod - def create_range_calculator(cls, spec: QuantizationSpec) -> RangeCalculatorBase: + def create_range_calculator( + cls, + spec: QuantizationSpec, + quantization_target: CompressionTargetTensor = CompressionTargetTensor.WEIGHT, + ) -> RangeCalculatorBase: """ Create a RangeCalculatorBase instance from a QuantizationSpec. Args: spec: QuantizationSpec instance containing configuration + quantization_target: The target tensor for quantization (weight/activation). Returns: RangeCalculatorBase instance configured from the spec @@ -48,6 +53,7 @@ def create_range_calculator(cls, spec: QuantizationSpec) -> RangeCalculatorBase: # Standard arguments for range calculator common_args = { "granularity": spec.granularity, + "quantization_target": quantization_target, } # Automatically detect and include any extra arguments @@ -96,7 +102,7 @@ def create_qparams_calculator( ) # Create range calculator first - range_calculator = cls.create_range_calculator(spec) + range_calculator = cls.create_range_calculator(spec, quantization_target) # Standard arguments for qparams calculator common_args = { @@ -202,7 +208,6 @@ def create_fake_quantizer( "quant_min": spec.quant_min, "quant_max": spec.quant_max, "qparams_calculator": qparams_calculator, - "quantization_target": quantization_target, "n_bits": spec.n_bits, } @@ -283,7 +288,6 @@ def create_fake_quantizer_partial( "target_dtype": spec.target_dtype, "quant_min": spec.quant_min, "quant_max": spec.quant_max, - "quantization_target": quantization_target, "n_bits": spec.n_bits, } diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index 1c0bfaa..9ff72aa 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/src/coreai_opt/quantization/spec/fake_quantize.py @@ -53,7 +53,6 @@ def __init__( quant_min: int | float, quant_max: int | float, qparams_calculator: QParamsCalculatorBase, - quantization_target: CompressionTargetTensor, n_bits: int | None = None, **kwargs, ): @@ -65,7 +64,6 @@ def __init__( self.quant_min = quant_min self.quant_max = quant_max self.qparams_calculator = qparams_calculator - self.quantization_target = quantization_target self.register_buffer("_disabled", torch.tensor(False)) # Infer n_bits from dtype if not provided @@ -78,6 +76,11 @@ def qscheme(self) -> QuantizationScheme: """The quantization scheme, delegated to the qparams_calculator.""" return self.qparams_calculator.qscheme + @property + def quantization_target(self) -> CompressionTargetTensor: + """Getter for quantization target.""" + return self.qparams_calculator.quantization_target + @property def granularity(self) -> QuantizationGranularity: """Getter for granularity.""" @@ -359,7 +362,7 @@ def _quantize_int( This function quantizes the values in tensor but keeps the quantized tensor dtype in FP. """ - block_size = self.granularity.get_block_size(tensor.shape) + block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target) original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes( tensor, block_size ) @@ -384,7 +387,7 @@ def _dequantize_int( output_dtype: torch.dtype, ) -> torch.Tensor: """Integer dequantization. See :func:`_dequantize_int` for the math.""" - block_size = self.granularity.get_block_size(tensor.shape) + block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target) original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes( tensor, block_size ) @@ -406,7 +409,7 @@ def _quantize_float( """ Floating-point quantization: cast_to_low_precision(clamp(input / scale, min, max)) """ - block_size = self.granularity.get_block_size(tensor.shape) + block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target) original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes( tensor, block_size ) @@ -427,7 +430,7 @@ def _dequantize_float( output_dtype: torch.dtype, ) -> torch.Tensor: """Floating-point dequantization: input * scale""" - block_size = self.granularity.get_block_size(tensor.shape) + block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target) original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes( tensor, block_size ) @@ -449,7 +452,7 @@ def _fused_fake_quant_dequant( Dispatches to the int or float fused STE class based on self.dtype. """ - block_size = self.granularity.get_block_size(tensor.shape) + block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target) original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes( tensor, block_size ) diff --git a/src/coreai_opt/quantization/spec/granularity.py b/src/coreai_opt/quantization/spec/granularity.py index 859bac3..30a6225 100644 --- a/src/coreai_opt/quantization/spec/granularity.py +++ b/src/coreai_opt/quantization/spec/granularity.py @@ -13,6 +13,7 @@ from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis +from coreai_opt.config.spec import CompressionTargetTensor as _CompressionTargetTensor from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError @@ -50,7 +51,11 @@ def _serialize_model(self) -> dict[str, Any]: return data @abstractmethod - def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: + def _get_block_size( + self, + block_sizes_list: list[int], + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> list[int]: """ Given an initial list of the tensor shape, return a list of block sizes corresponding to each axis: @@ -62,6 +67,11 @@ def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: - if per-block structuring is being done for a certain axis, set the block size for that specific axis + ``quantization_target`` distinguishes weight from activation tensors. + Only per-block granularity uses it, since the two targets collapse + different sets of non-blocked axes (see :class:`PerBlockGranularity`); + the other granularities ignore it. + Example: - ``[10, 5, 2]`` with per-channel structuring on axis 1 results in ``[10, 1, 2]`` @@ -71,11 +81,20 @@ def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: """ pass - def get_block_size(self, tensor_shape: torch.Size) -> tuple[int, ...]: + def get_block_size( + self, + tensor_shape: torch.Size, + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> tuple[int, ...]: """ Get a list of block sizes based on the granularity. + + Args: + tensor_shape: Shape of the tensor being quantized. + quantization_target: Whether the tensor is a weight or an activation. + Defaults to ``WEIGHT``, which preserves the historical behavior. """ - return tuple(self._get_block_size(list(tensor_shape))) + return tuple(self._get_block_size(list(tensor_shape), quantization_target)) # The axis resolution logic lives here because it is granularity-specific. # Currently only PerChannelGranularity has a meaningful axis to resolve, but @@ -118,7 +137,11 @@ class PerTensorGranularity(QuantizationGranularity): axis: Literal[None] = None - def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: + def _get_block_size( + self, + block_sizes_list: list[int], + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> list[int]: return block_sizes_list @@ -137,7 +160,11 @@ class PerChannelGranularity(QuantizationGranularity): axis: int | None = None - def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: + def _get_block_size( + self, + block_sizes_list: list[int], + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> list[int]: if self.axis is None: raise ValueError( @@ -167,9 +194,11 @@ class PerBlockGranularity(QuantizationGranularity): This applies quantization to blocks of values within the tensor. Supports two modes: 1. Single-axis mode: Quantize blocks along one specific axis - (no blocking for axis>=2) - - ``axis``: The axis to create blocks (0 or 1) + - ``axis``: The axis to create blocks. May be negative (Python-style + indexing). For weight quantization this is typically ``0`` or ``1`` + (the channel axes); for activation quantization it is commonly the + last / reduction axis (e.g. ``-1``). - ``block_size``: Integer specifying block size for that axis 2. Multi-axis mode: Create blocks across multiple axes simultaneously @@ -182,39 +211,68 @@ class PerBlockGranularity(QuantizationGranularity): ``Quantizer.prepare()`` automatically resolves the axis based on the module type for weight quantization. + Single-axis mode treats weights and activations differently: + + - ``WEIGHT``: only the two leading channel axes take part. Whichever of them + is not the block axis collapses to ``1`` (one scale per slice), while + trailing dimensions — e.g. conv kernel dims — keep their full size, so each + block spans the whole kernel. + - ``ACTIVATION``: every axis other than the block axis collapses to ``1``, so + the scale holds one entry per block *and* per position along all the other + axes. + .. list-table:: :header-rows: 1 - * - Weight tensor shape (input) + * - Tensor shape (input) + - target - axis - block_size - - Weight shape of each block (output) + - Shape of each block (output) * - [C_out, C_in] + - weight - 1 - 32 - [1, 32] * - [C_out, C_in] + - weight - None - (4, 8) - [4, 8] * - [C_out, C_in, KH, KW] + - weight - 0 - 16 - [16, 1, KH, KW] * - [C_out, C_in, KH, KW] + - weight - None - (4, 16, 3, -1) - [4, 16, 3, KW] + * - [B, S, D] + - activation + - -1 + - 16 + - [1, 1, 16] + * - [B, C, H, W] + - activation + - 1 + - 16 + - [1, 16, 1, 1] """ - axis: Annotated[int, Field(ge=0, le=1)] | None = None + axis: int | None = None block_size: Annotated[int, Field(gt=0)] | tuple[Annotated[int, Field(gt=0)] | Literal[-1], ...] - def _get_block_size(self, block_sizes_list: list[int]) -> list[int]: + def _get_block_size( + self, + block_sizes_list: list[int], + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> list[int]: if isinstance(self.block_size, tuple): return self._handle_multi_axis_block_size(block_sizes_list) else: - return self._handle_single_axis_block_size(block_sizes_list) + return self._handle_single_axis_block_size(block_sizes_list, quantization_target) def _handle_multi_axis_block_size(self, block_sizes_list: list[int]) -> list[int]: """Handle blocking when self.block_size is a tuple""" @@ -244,31 +302,45 @@ def _handle_multi_axis_block_size(self, block_sizes_list: list[int]) -> list[int return block_sizes_list - def _handle_single_axis_block_size(self, block_sizes_list: list[int]) -> list[int]: - """Handle blocking when self.block_size is an integer""" + def _handle_single_axis_block_size( + self, + block_sizes_list: list[int], + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + ) -> list[int]: + """Handle blocking when ``block_size`` is an integer. + + ``axis`` may be negative and is resolved against the tensor rank. + + ``quantization_target`` decides which of the non-blocked axes collapse to + ``1``: weights keep their trailing (e.g. kernel) dimensions whole, while + activations collapse every axis but the block axis. See the class + docstring for examples. + """ # TODO: Logic to be added where if self.axis is None, # we can figure out the optimal axis for the user if self.axis is None: raise ValueError("axis must be specified when block_size is an int") - if self.axis >= len(block_sizes_list): + rank = len(block_sizes_list) + axis = self.axis + rank if self.axis < 0 else self.axis + + if axis < 0 or axis >= rank: raise ValueError( - f"axis {self.axis} is out of bounds for tensor of rank {len(block_sizes_list)}" + f"axis {self.axis} is out of bounds for tensor of rank {rank}. " + f"Allowed axis range is [{-rank}, {rank})" ) - if block_sizes_list[self.axis] % self.block_size != 0: + if block_sizes_list[axis] % self.block_size != 0: raise _BlockSizeMismatchError( - f"Tensor size {block_sizes_list[self.axis]} along axis {self.axis} " + f"Tensor size {block_sizes_list[axis]} along axis {axis} " f"is not divisible by block size {self.block_size}" ) - # For integer block_size, only process the first two dimensions - # (which would be input and output channel axis in no particular order) - # Set the specified axis to block_size, set the other dimension (0 or 1) to 1 - # Leave all higher dimensions (index 2+) unchanged - block_sizes_list[self.axis] = self.block_size - for axis, _ in enumerate(block_sizes_list[:2]): - if axis != self.axis: - block_sizes_list[axis] = 1 + collapse_upto = rank if quantization_target == _CompressionTargetTensor.ACTIVATION else 2 + + block_sizes_list[axis] = self.block_size + for i, _ in enumerate(block_sizes_list[:collapse_upto]): + if i != axis: + block_sizes_list[i] = 1 return block_sizes_list diff --git a/src/coreai_opt/quantization/spec/qparams_calculator.py b/src/coreai_opt/quantization/spec/qparams_calculator.py index b110292..5260302 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -82,6 +82,11 @@ def granularity(self) -> QuantizationGranularity: """Getter for granularity.""" return self._granularity + @property + def quantization_target(self): + """Whether the quantized tensor is a weight or an activation.""" + return self.range_calculator.quantization_target + @granularity.setter def granularity(self, granularity: QuantizationGranularity) -> None: """Update granularity for this calculator and its range calculator. @@ -117,7 +122,9 @@ def _get_tensor_with_granularity_from_scalar( Return a tensor with dimensions equal to num blocks in each dimension, comprised of values equal to scalar. """ - block_size_list = self.granularity.get_block_size(input_tensor.shape) + block_size_list = self.granularity.get_block_size( + input_tensor.shape, self.quantization_target + ) num_blocks_list = [ inp_size // block_size for inp_size, block_size in zip(input_tensor.shape, block_size_list, strict=True) @@ -229,7 +236,7 @@ def _compute_scale_zero_point_minval( min_val=min_val, max_val=max_val, mapping_type=QuantizationScheme._to_mapping_type(self.qscheme), - block_size=self.granularity.get_block_size(tensor.shape), + block_size=self.granularity.get_block_size(tensor.shape, self.quantization_target), target_dtype=self.target_dtype, quant_min=self.quant_min, quant_max=self.quant_max, diff --git a/src/coreai_opt/quantization/spec/range_calculator.py b/src/coreai_opt/quantization/spec/range_calculator.py index d021c99..73daf86 100644 --- a/src/coreai_opt/quantization/spec/range_calculator.py +++ b/src/coreai_opt/quantization/spec/range_calculator.py @@ -10,6 +10,7 @@ from torchao.quantization.quant_primitives import _get_reduction_params from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin +from coreai_opt.config.spec import CompressionTargetTensor as _CompressionTargetTensor from .granularity import QuantizationGranularity @@ -20,16 +21,22 @@ class RangeCalculatorBase(_ClassRegistryMixin, nn.Module): of a given tensor. """ - def __init__(self, granularity: QuantizationGranularity, **kwargs): + def __init__( + self, + granularity: QuantizationGranularity, + quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT, + **kwargs, + ): super().__init__() self.granularity = granularity + self.quantization_target = quantization_target def _reshape_min_max(self, range_tensor: torch.Tensor, input_shape: torch.Size): """ Reshape range_tensor to have the same number of dimensions as input shape, taking block size into account. """ - block_size_list = self.granularity.get_block_size(input_shape) + block_size_list = self.granularity.get_block_size(input_shape, self.quantization_target) # While reducing, each dimension with block size other than 1 or the original # dimension size will be split into 2 dimensions of num_blocks and block_size. @@ -77,7 +84,7 @@ class MinMaxRangeCalculator(RangeCalculatorBase): """ def _generate_min_max(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - block_size_list = self.granularity.get_block_size(tensor.shape) + block_size_list = self.granularity.get_block_size(tensor.shape, self.quantization_target) shape_for_reduction, reduction_dims = _get_reduction_params(block_size_list, tensor.size()) # If tensor is already the shape required, no minmaxing is needed. diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index 857549e..9cefda0 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -37,6 +37,7 @@ ) from coreai_opt.quantization.config import OpQuantizerConfig from coreai_opt.quantization.spec import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationScheme, @@ -228,29 +229,36 @@ def forward(self, x): } -def _build_pool(pool_type: str, dim: int) -> nn.Module: +def _build_pool(pool_type: str, dim: int, preserve_shape: bool = False) -> nn.Module: pool_cls = _POOL_CLS[pool_type][dim] if pool_type == "adaptive_avg": # AdaptiveAvgPool takes an output size, not a kernel/stride; 4 is half # of _POOL_SPATIAL_SHAPE's 8, matching stride-2 max/avg pool's shrink. - return pool_cls((4,) * dim) + # Asking for the input's own size keeps every dimension intact instead. + output_size = _POOL_SPATIAL_SHAPE[dim] if preserve_shape else (4,) * dim + return pool_cls(output_size) + if preserve_shape: + # A 3-wide window with stride 1 and padding 1 leaves every spatial dimension the same size. + return pool_cls(3, stride=1, padding=1) return pool_cls(2, stride=2) class SimplePoolModel(nn.Module): - """Simple model with a pool that shrinks its spatial dimensions. + """Simple model with a pool between two convs. Parametrized over ``pool_type`` ("max", "avg", "adaptive_avg") and ``dim`` - (1, 2, or 3) to exercise MaxPool/AvgPool/AdaptiveAvgPool 1d/2d/3d. + (1, 2, or 3) to exercise MaxPool/AvgPool/AdaptiveAvgPool 1d/2d/3d. The pool + shrinks its spatial dimensions by default; ``preserve_shape=True`` keeps them + the same size. """ - def __init__(self, pool_type: str = "max", dim: int = 2): + def __init__(self, pool_type: str = "max", dim: int = 2, preserve_shape: bool = False): super().__init__() conv_cls = _POOL_CONV_CLS[dim] kernel = (3,) * dim padding = (1,) * dim self.conv = conv_cls(3, 4, kernel, padding=padding, bias=False) - self.pool = _build_pool(pool_type, dim) + self.pool = _build_pool(pool_type, dim, preserve_shape=preserve_shape) self.conv2 = conv_cls(4, 4, kernel, padding=padding, bias=False) def forward(self, x): @@ -1393,20 +1401,29 @@ def test_shared_observer_no_qspec_set(self): for node in prepared_model.graph.nodes: assert "activation_post_process" not in node.name - def test_shared_observer_forces_per_tensor_for_flatten(self): + @pytest.mark.parametrize( + "granularity", + [ + PerChannelGranularity(axis=0), + PerBlockGranularity(axis=0, block_size=2), + ], + ids=["per_channel", "per_block"], + ) + def test_shared_observer_forces_per_tensor_for_flatten(self, granularity): """ - When per-channel activation granularity is used globally, shared observer - ops that alter channel semantics (e.g., flatten) should have their shared - fake quantize modules forced to per-tensor granularity. The conv and linear - activation quantizers (which are separate objects) should remain per-channel. + When an axis-carrying activation granularity is used globally, shared + observer ops that alter channel semantics (e.g., flatten) should have their + shared fake quantize modules forced to per-tensor granularity. The conv and + linear activation quantizers (which are separate objects) should keep their + configured granularity. """ model = SimpleFlattenModel() inp = torch.randn(2, 3, 2, 2) - per_channel_activation_spec = QuantizationSpec( + activation_spec = QuantizationSpec( dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC, - granularity=PerChannelGranularity(axis=0), + granularity=granularity, fake_quantize_cls="default", qparam_calculator_cls="default", range_calculator_cls="minmax", @@ -1415,8 +1432,8 @@ def test_shared_observer_forces_per_tensor_for_flatten(self): config = QuantizerConfig( global_config=ModuleQuantizerConfig( op_state_spec=None, - op_input_spec={"*": per_channel_activation_spec}, - op_output_spec={"*": per_channel_activation_spec}, + op_input_spec={"*": activation_spec}, + op_output_spec={"*": activation_spec}, ), ) @@ -1449,14 +1466,15 @@ def test_shared_observer_forces_per_tensor_for_flatten(self): prepared_model(inp) assert input_fq.qparams_calculator.scale.numel() == 1 - # Conv input and linear output fake quantizers should remain per-channel - # since they are separate objects (not shared across a shape-destroying op). + # Conv input and linear output fake quantizers should keep their configured + # granularity since they are separate objects (not shared across a + # shape-destroying op). conv_node = [node for node in prepared_model.graph.nodes if node.name == "conv2d"][0] conv_input_fq_node = [ n for n in conv_node.all_input_nodes if "activation_post_process" in n.name ][0] conv_input_fq = getattr(prepared_model, conv_input_fq_node.name) - assert isinstance(conv_input_fq.granularity, PerChannelGranularity) + assert isinstance(conv_input_fq.granularity, type(granularity)) assert conv_input_fq.qparams_calculator.scale.numel() > 1 linear_node = [node for node in prepared_model.graph.nodes if node.name == "linear"][0] @@ -1464,32 +1482,40 @@ def test_shared_observer_forces_per_tensor_for_flatten(self): n for n in linear_node.users if "activation_post_process" in n.name ][0] linear_output_fq = getattr(prepared_model, linear_output_fq_node.name) - assert isinstance(linear_output_fq.granularity, PerChannelGranularity) + assert isinstance(linear_output_fq.granularity, type(granularity)) assert linear_output_fq.qparams_calculator.scale.numel() > 1 _POOL_ATEN_PREFIX = {"max": "max_pool", "avg": "avg_pool", "adaptive_avg": "adaptive_avg_pool"} @pytest.mark.parametrize("dim", [1, 2, 3], ids=["1d", "2d", "3d"]) @pytest.mark.parametrize("pool_type", ["max", "avg", "adaptive_avg"]) + @pytest.mark.parametrize( + "granularity", + [ + PerChannelGranularity(axis=-1), + PerBlockGranularity(axis=-1, block_size=2), + ], + ids=["per_channel", "per_block"], + ) def test_shared_observer_forces_per_tensor_for_pool_axis_that_shrinks( - self, caplog, pool_type, dim + self, caplog, pool_type, dim, granularity ): """ - When per-channel activation granularity is configured on an axis whose + When activation granularity is configured on an axis whose size a pool actually shrinks (e.g. the last spatial axis under a stride-2/output-shrinking pool), the shared fake quantize module spanning the pool's input and output should be forced to per-tensor granularity, with a warning naming the op and axis. Conv activation - quantizers (separate objects) should remain per-channel. + quantizers (separate objects) should keep their configured granularity. """ model = SimplePoolModel(pool_type=pool_type, dim=dim) inp = SimplePoolModel.example_input(dim=dim) pool_target_name = f"{self._POOL_ATEN_PREFIX[pool_type]}{dim}d" - per_channel_activation_spec = QuantizationSpec( + activation_spec = QuantizationSpec( dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC, - granularity=PerChannelGranularity(axis=-1), + granularity=granularity, fake_quantize_cls="default", qparam_calculator_cls="default", range_calculator_cls="minmax", @@ -1498,8 +1524,8 @@ def test_shared_observer_forces_per_tensor_for_pool_axis_that_shrinks( config = QuantizerConfig( global_config=ModuleQuantizerConfig( op_state_spec=None, - op_input_spec={"*": per_channel_activation_spec}, - op_output_spec={"*": per_channel_activation_spec}, + op_input_spec={"*": activation_spec}, + op_output_spec={"*": activation_spec}, ), ) @@ -1536,13 +1562,13 @@ def test_shared_observer_forces_per_tensor_for_pool_axis_that_shrinks( assert input_fq.qparams_calculator.scale.numel() == 1 # Conv activation quantizers (separate objects, not shared across the - # pool) should remain per-channel. + # pool) should keep their configured granularity. conv_node = [node for node in prepared_model.graph.nodes if node.name == f"conv{dim}d"][0] conv_input_fq_node = [ n for n in conv_node.all_input_nodes if "activation_post_process" in n.name ][0] conv_input_fq = getattr(prepared_model, conv_input_fq_node.name) - assert isinstance(conv_input_fq.granularity, PerChannelGranularity) + assert isinstance(conv_input_fq.granularity, type(granularity)) assert conv_input_fq.qparams_calculator.scale.numel() > 1 @pytest.mark.parametrize("dim", [1, 2, 3], ids=["1d", "2d", "3d"]) @@ -1607,6 +1633,67 @@ def test_shared_observer_preserves_per_channel_for_pool_axis_that_is_invariant( prepared_model(inp) assert input_fq.qparams_calculator.scale.numel() > 1 + @pytest.mark.parametrize("dim", [1, 2, 3], ids=["1d", "2d", "3d"]) + @pytest.mark.parametrize("pool_type", ["max", "avg", "adaptive_avg"]) + def test_shared_observer_preserves_per_block_when_pool_preserves_shape( + self, caplog, pool_type, dim + ): + """ + A per-block scale stays valid across a shared observer when the pool leaves + every dimension the same size, so output should be not forced to per-tensor. + """ + model = SimplePoolModel(pool_type=pool_type, dim=dim, preserve_shape=True) + inp = SimplePoolModel.example_input(dim=dim) + pool_target_name = f"{self._POOL_ATEN_PREFIX[pool_type]}{dim}d" + + # Block along the last axis (size 8) so every activation in the model, + # including the 3-channel input, is divisible by the block size. + per_block_activation_spec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerBlockGranularity(axis=-1, block_size=2), + fake_quantize_cls="default", + qparam_calculator_cls="default", + range_calculator_cls="minmax", + ) + + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec=None, + op_input_spec={"*": per_block_activation_spec}, + op_output_spec={"*": per_block_activation_spec}, + ), + ) + + quantizer = Quantizer(model, config) + with caplog.at_level(logging.WARNING): + prepared_model = quantizer.prepare(example_inputs=(inp,)) + + pool_node = [ + node for node in prepared_model.graph.nodes if pool_target_name in str(node.target) + ][0] + input_fq = getattr(prepared_model, pool_node.all_input_nodes[0].name) + output_fq = getattr(prepared_model, list(pool_node.users.keys())[0].name) + + # Shared observer: input and output should be the same object + assert input_fq is output_fq + + # Nothing was resized or moved, so per-block must survive. + assert isinstance(input_fq.granularity, PerBlockGranularity) + + # No warning should have been logged for this safe op. + assert not any(pool_target_name in record.message for record in caplog.records) + + prepared_model.eval() + with torch.no_grad(): + prepared_model(inp) + + # One scale entry per position on every axis except the blocked last one, + # which holds 8 / 2 blocks. The pooled tensor is (2, 4, *spatial). + spatial = _POOL_SPATIAL_SHAPE[dim] + expected_scale_shape = (2, 4, *spatial[:-1], spatial[-1] // 2) + assert tuple(input_fq.qparams_calculator.scale.shape) == expected_scale_shape + def test_cat(self): """ Given a model with @@ -1671,9 +1758,19 @@ def test_concat(self): # Final linear output quantizer is not associated with the others assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 - def test_shared_observer_forces_per_tensor_when_transpose_swaps_equal_size_axes(self, caplog): + @pytest.mark.parametrize( + "granularity", + [ + PerChannelGranularity(axis=2), + PerBlockGranularity(axis=2, block_size=2), + ], + ids=["per_channel", "per_block"], + ) + def test_shared_observer_forces_per_tensor_when_transpose_swaps_equal_size_axes( + self, caplog, granularity + ): """ - A per-channel axis that transpose swaps with another axis of the SAME + An axis that transpose swaps with another axis of the SAME size must be forced to per-tensor even though the sizes match on both sides - a size match is not proof the axis is untouched, since transpose can relabel which physical dimension sits at that index. @@ -1685,10 +1782,10 @@ def test_shared_observer_forces_per_tensor_when_transpose_swaps_equal_size_axes( model = SimpleConcatTransposeModel() inp = torch.randn(2, 3, 8, 8) - per_channel_activation_spec = QuantizationSpec( + activation_spec = QuantizationSpec( dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC, - granularity=PerChannelGranularity(axis=2), + granularity=granularity, fake_quantize_cls="default", qparam_calculator_cls="default", range_calculator_cls="minmax", @@ -1697,8 +1794,8 @@ def test_shared_observer_forces_per_tensor_when_transpose_swaps_equal_size_axes( config = QuantizerConfig( global_config=ModuleQuantizerConfig( op_state_spec=None, - op_input_spec={"*": per_channel_activation_spec}, - op_output_spec={"*": per_channel_activation_spec}, + op_input_spec={"*": activation_spec}, + op_output_spec={"*": activation_spec}, ), ) @@ -1721,15 +1818,26 @@ def test_shared_observer_forces_per_tensor_when_transpose_swaps_equal_size_axes( assert isinstance(input_fq.granularity, PerTensorGranularity) @pytest.mark.parametrize( - ("axis", "expected_granularity"), + ("granularity", "expected_granularity"), [ - (1, PerChannelGranularity), # channel: permute(0,1,3,2) leaves it alone - (2, PerTensorGranularity), # height: swapped with width by the permute - (3, PerTensorGranularity), # width: swapped with height by the permute + # channel: permute(0,1,3,2) leaves it alone + (PerChannelGranularity(axis=1), PerChannelGranularity), + # height: swapped with width by the permute + (PerChannelGranularity(axis=2), PerTensorGranularity), + # width: swapped with height by the permute + (PerChannelGranularity(axis=3), PerTensorGranularity), + # Same untouched channel axis as the first row, but a per-block scale + # also spans axes 2 and 3, which the permute swaps. + (PerBlockGranularity(axis=1, block_size=2), PerTensorGranularity), + ], + ids=[ + "per_channel_axis_untouched_by_permute", + "per_channel_axis_swapped_to_3", + "per_channel_axis_swapped_to_2", + "per_block_axis_untouched_but_scale_spans_swapped_axes", ], - ids=["axis_untouched_by_permute", "axis_swapped_to_3", "axis_swapped_to_2"], ) - def test_shared_observer_permute_axis_identity(self, axis, expected_granularity): + def test_shared_observer_permute_axis_identity(self, granularity, expected_granularity): """ permute(0, 1, 3, 2) leaves axes 0 and 1 mapped to themselves but swaps axes 2 and 3 with each other. A per-channel axis on the @@ -1738,14 +1846,18 @@ def test_shared_observer_permute_axis_identity(self, axis, expected_granularity) be forced to per-tensor, even though both have equal size (8 == 8) on a square conv output - a size match alone doesn't prove the axis wasn't relabeled. + + Per-block granularity on that same untouched channel axis must still be + forced, because its scale carries one entry per position on every other + axis - including the two the permute swaps. """ model = SimpleConcatPermuteModel() inp = torch.randn(2, 3, 8, 8) - per_channel_activation_spec = QuantizationSpec( + activation_spec = QuantizationSpec( dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC, - granularity=PerChannelGranularity(axis=axis), + granularity=granularity, fake_quantize_cls="default", qparam_calculator_cls="default", range_calculator_cls="minmax", @@ -1754,8 +1866,8 @@ def test_shared_observer_permute_axis_identity(self, axis, expected_granularity) config = QuantizerConfig( global_config=ModuleQuantizerConfig( op_state_spec=None, - op_input_spec={"*": per_channel_activation_spec}, - op_output_spec={"*": per_channel_activation_spec}, + op_input_spec={"*": activation_spec}, + op_output_spec={"*": activation_spec}, ), ) diff --git a/tests/quantization/test_eager_quant.py b/tests/quantization/test_eager_quant.py index 942f139..17049bb 100644 --- a/tests/quantization/test_eager_quant.py +++ b/tests/quantization/test_eager_quant.py @@ -3379,7 +3379,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerBlockGranularity(axis=1, block_size=32), True, - "FP4 activation quantization is not supported for MLIR export", + "Core AI export does not support FP4 activation quantization", id="fp4_activation_rejected", ), pytest.param( diff --git a/tests/quantization/test_factory.py b/tests/quantization/test_factory.py index 1afb627..e17f112 100644 --- a/tests/quantization/test_factory.py +++ b/tests/quantization/test_factory.py @@ -38,7 +38,14 @@ class TestQuantizationComponentFactory: """Test the QuantizationComponentFactory class""" - def test_create_range_calculator(self): + @pytest.mark.parametrize( + "quantization_target", + [ + CompressionTargetTensor.WEIGHT, + CompressionTargetTensor.ACTIVATION, + ], + ) + def test_create_range_calculator(self, quantization_target): """Test creating range calculator from spec""" spec = QuantizationSpec( dtype=torch.int8, @@ -49,10 +56,11 @@ def test_create_range_calculator(self): range_calculator_cls=MinMaxRangeCalculator, ) - range_calc = QuantizationComponentFactory.create_range_calculator(spec) + range_calc = QuantizationComponentFactory.create_range_calculator(spec, quantization_target) assert isinstance(range_calc, MinMaxRangeCalculator) assert range_calc.granularity == spec.granularity + assert range_calc.quantization_target == quantization_target @pytest.mark.parametrize( "range", diff --git a/tests/quantization/test_fake_quantize.py b/tests/quantization/test_fake_quantize.py index 16e3ada..c953201 100644 --- a/tests/quantization/test_fake_quantize.py +++ b/tests/quantization/test_fake_quantize.py @@ -69,7 +69,6 @@ def test_fake_quant_dequant_no_reduction(qscheme, granularity, qformulation): qparam_calculator = StaticQParamsCalculator(range_calculator=range_calculator, **kwargs) fq = _DefaultFakeQuantizeImpl( qparams_calculator=qparam_calculator, - quantization_target=CompressionTargetTensor.WEIGHT, **kwargs, ) @@ -105,7 +104,6 @@ def test_set_granularity(): qparam_calculator = StaticQParamsCalculator(range_calculator=range_calculator, **kwargs) fq = _DefaultFakeQuantizeImpl( qparams_calculator=qparam_calculator, - quantization_target=CompressionTargetTensor.WEIGHT, **kwargs, ) x = torch.randn(2, 5) @@ -189,7 +187,6 @@ def fq(self, dtype, qscheme, qformulation, granularity): qparam_calculator = StaticQParamsCalculator(range_calculator=range_calculator, **kwargs) return _DefaultFakeQuantizeImpl( qparams_calculator=qparam_calculator, - quantization_target=CompressionTargetTensor.WEIGHT, **kwargs, ) diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index b04589d..901cb24 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -1369,7 +1369,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerBlockGranularity(axis=1, block_size=32), True, - "FP4 activation quantization is not supported for MLIR export", + "Core AI export does not support FP4 activation quantization", id="fp4_activation_rejected", ), pytest.param( @@ -1423,29 +1423,52 @@ def test_fp4_invalid_config_rejected( class TestBlockSizeMismatchSkipGraphMode: """Test that non-divisible block sizes produce a warning and skip quantization in graph mode.""" - def test_non_divisible_block_size_warns_and_skips(self, caplog): - """Graph mode: non-divisible layer gets FQ disabled and removed, divisible stays.""" - # out_features=1000: 1000 % 32 != 0 → FQ disabled and removed - # out_features=1024: 1024 % 32 == 0 → FQ stays enabled + @pytest.mark.parametrize( + ("target", "expected_fq_count"), + [ + # Weights: only Linear(1000, 1024)'s weight is divisible on axis 0. + ("weight", 1), + # Activations: the 768-wide input and 1024-wide output survive; the + # 1000-wide tensor between the two linears does not. + ("activation", 2), + ], + ) + def test_non_divisible_block_size_warns_and_skips(self, caplog, target, expected_fq_count): + """Graph mode: non-divisible tensors get their FQ disabled and removed, + divisible ones stay. Applies to weights and activations alike.""" + # 768 % 32 == 0 and 1024 % 32 == 0, but 1000 % 32 != 0 model = torch.nn.Sequential( torch.nn.Linear(768, 1000), torch.nn.Linear(1000, 1024), ) example_inputs = (torch.randn(1, 768),) - weight_spec = QuantizationSpec( - dtype="int8", - qscheme="symmetric", - granularity=PerBlockGranularity(axis=0, block_size=32), - ) - config = QuantizerConfig( - global_config=ModuleQuantizerConfig( - op_state_spec={"weight": weight_spec}, + if target == "weight": + # axis 0 is out_features: 1000 is not divisible, 1024 is. + spec = QuantizationSpec( + dtype="int8", + qscheme="symmetric", + granularity=PerBlockGranularity(axis=0, block_size=32), + ) + module_config = ModuleQuantizerConfig( + op_state_spec={"weight": spec}, op_input_spec=None, op_output_spec=None, - ), - execution_mode="graph", - ) + ) + else: + # The last axis carries the feature dim for every activation here. + spec = QuantizationSpec( + dtype="int8", + qscheme="symmetric", + granularity=PerBlockGranularity(axis=-1, block_size=32), + ) + module_config = ModuleQuantizerConfig( + op_state_spec=None, + op_input_spec={"*": spec}, + op_output_spec={"*": spec}, + ) + + config = QuantizerConfig(global_config=module_config, execution_mode="graph") quantizer = Quantizer(model, config) @@ -1455,15 +1478,18 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog): assert prepared_model is not None assert any("Skipping quantization" in msg for msg in caplog.messages) - # Only the second Linear (out_features=1024) is compatible with block_size=32 fq_modules = [m for m in prepared_model.modules() if isinstance(m, FakeQuantizeImplBase)] assert all(not m.is_disabled() for m in fq_modules), ( "Disabled FQ modules should be removed during prepare()" ) - assert len(fq_modules) == 1, ( - f"Expected 1 enabled FQ module (for divisible layer), got {len(fq_modules)}" + assert len(fq_modules) == expected_fq_count, ( + f"Expected {expected_fq_count} enabled FQ module(s) for the divisible " + f"tensors, got {len(fq_modules)}" ) + # The graph must still be runnable after the disabled nodes were removed. + assert prepared_model(*example_inputs).shape == (1, 1024) + @pytest.mark.parametrize( "backend", [ diff --git a/tests/quantization/test_quantization.py b/tests/quantization/test_quantization.py index 8e3b917..f369401 100644 --- a/tests/quantization/test_quantization.py +++ b/tests/quantization/test_quantization.py @@ -16,6 +16,7 @@ import torch.nn as nn from coreai_opt import ExportBackend +from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization import ( ModuleQuantizerConfig, QuantizationSpec, @@ -25,6 +26,7 @@ from coreai_opt.quantization._graph.quantizer import GraphQuantizer from coreai_opt.quantization.config.quantization_config import QATSchedule from coreai_opt.quantization.spec import ( + PerBlockGranularity, PerTensorGranularity, QuantizationScheme, default_activation_quantization_spec, @@ -535,6 +537,103 @@ def test_qat_schedule_does_not_disable_dynamic_observer(self, execution_mode): assert dynamic_fq.observer_enabled.item() == 1 +class TestPerBlockActivationQuantization: + """Per-block activation quantization is supported for prepare/simulation but + is not exportable.""" + + class _TwoConvModel(nn.Module): + def __init__(self, channels: int = 32) -> None: + super().__init__() + self.conv1 = nn.Conv2d(channels, channels, 3, padding=1) + self.conv2 = nn.Conv2d(channels, channels, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv2(self.conv1(x)) + + @staticmethod + def _make_per_block_activation_config( + execution_mode: str, weight_axis: int = 1, activation_axis: int = -1 + ) -> QuantizerConfig: + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerBlockGranularity(axis=weight_axis, block_size=16), + ) + }, + op_input_spec={ + "*": QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerBlockGranularity(axis=activation_axis, block_size=16), + ) + }, + op_output_spec=None, + ), + ).set_execution_mode(execution_mode) + + @pytest.mark.parametrize( + "weight_axis,expected_weight_scale_shape,activation_axis,expected_activation_scale_shape", + [ + (1, (32, 2, 1, 1), 1, (1, 2, 16, 16)), + (0, (2, 32, 1, 1), -1, (1, 32, 16, 1)), + ], + ) + def test_per_block_scale_shape( + self, + execution_mode, + weight_axis, + expected_weight_scale_shape, + activation_axis, + expected_activation_scale_shape, + ): + """Weight and Activations should use respective logic for block_size""" + config = self._make_per_block_activation_config( + execution_mode, weight_axis, activation_axis + ) + + example_input = torch.randn(1, 32, 16, 16) + quantizer = Quantizer(self._TwoConvModel(), config) + prepared_model = quantizer.prepare((example_input,)) + prepared_model(example_input) + + act_scale_shapes = { + tuple(m.calculate_qparams()[0].shape) + for m in prepared_model.modules() + if isinstance(m, FakeQuantizeImplBase) + and m.quantization_target == CompressionTargetTensor.ACTIVATION + and not m.is_disabled() + } + assert act_scale_shapes == {expected_activation_scale_shape}, ( + f"Expected activation scale shape {expected_activation_scale_shape}, " + f"got {act_scale_shapes}" + ) + + weight_scale_shapes = { + tuple(m.calculate_qparams()[0].shape) + for m in prepared_model.modules() + if isinstance(m, FakeQuantizeImplBase) + and m.quantization_target == CompressionTargetTensor.WEIGHT + and not m.is_disabled() + } + assert weight_scale_shapes == {expected_weight_scale_shape}, ( + f"Expected weight scale shape {expected_weight_scale_shape}, got {weight_scale_shapes}" + ) + + @pytest.mark.parametrize("backend", [ExportBackend.CoreAI, ExportBackend.CoreML]) + def test_finalize_rejects_per_block_activation(self, execution_mode, backend): + config = self._make_per_block_activation_config(execution_mode) + quantizer = Quantizer(SimpleLinearModel(), config) + + prepared_model = quantizer.prepare((torch.randn(4, 64),)) + prepared_model(torch.randn(4, 64)) + + with pytest.raises(Exception, match="does not support PerBlockGranularity on activation"): + quantizer.finalize(prepared_model, backend=backend) + + 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 8dc9efe..da122ea 100644 --- a/tests/quantization/test_quantization_spec.py +++ b/tests/quantization/test_quantization_spec.py @@ -9,6 +9,7 @@ from pydantic import ValidationError from torchao.quantization import MappingType as TorchAOMappingType +from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.quantization import QuantizationSpec from coreai_opt.quantization.spec import ( PerBlockGranularity, @@ -253,8 +254,6 @@ def test_invalid_qformulation_string(): }, ), ("per_block", {"axis": 1, "block_size": 0}), - ("per_block", {"axis": 2, "block_size": 5}), - ("per_block", {"axis": 3, "block_size": 5}), ( "per_block", { @@ -302,6 +301,14 @@ def test_invalid_axis_block_size(granularity_type, granularity_params): ("per_block", {"axis": 1, "block_size": 3}, (7, 3), (1, 3)), ("per_block", {"axis": 0, "block_size": 4}, (8, 16, 3), (4, 1, 3)), ("per_block", {"axis": 1, "block_size": 8}, (7, 16, 3, 3), (1, 8, 3, 3)), + # Per block - block axis >= 2 and negative axes. Default target is WEIGHT, + # so only the two leading channel axes collapse to 1; trailing dims are + # left at full size. See test_get_block_size_per_block_target for how the + # ACTIVATION target differs. + ("per_block", {"axis": 2, "block_size": 16}, (1, 10, 32), (1, 1, 16)), + ("per_block", {"axis": -1, "block_size": 16}, (1, 10, 32), (1, 1, 16)), + ("per_block", {"axis": -1, "block_size": 4}, (10, 20), (1, 4)), + ("per_block", {"axis": -2, "block_size": 5}, (10, 20), (5, 1)), ("per_block", {"axis": None, "block_size": (2,)}, (10,), (2,)), ( "per_block", @@ -329,6 +336,37 @@ def test_get_block_size_valid_conditions( assert result == expected_block_size +@pytest.mark.parametrize( + "axis,block_size,tensor_shape,weight_block_size,activation_block_size", + [ + (1, 16, (1, 32, 10), (1, 16, 10), (1, 16, 1)), + (-1, 16, (1, 10, 32), (1, 1, 16), (1, 1, 16)), + (1, 16, (1, 32, 10, 64), (1, 16, 10, 64), (1, 16, 1, 1)), + (0, 16, (32, 64, 3, 3), (16, 1, 3, 3), (16, 1, 1, 1)), + (3, 16, (1, 32, 10, 64), (1, 1, 10, 16), (1, 1, 1, 16)), + (-1, 16, (1, 32, 10, 64), (1, 1, 10, 16), (1, 1, 1, 16)), + ], +) +def test_get_block_size_per_block_target( + axis, block_size, tensor_shape, weight_block_size, activation_block_size +): + """Per-block single-axis blocking resolves differently per quantization target. + + Weights block across the channel axes only; activations give every non-block + axis its own scale. + """ + granularity = PerBlockGranularity(axis=axis, block_size=block_size) + shape = torch.Size(tensor_shape) + + assert granularity.get_block_size(shape, CompressionTargetTensor.WEIGHT) == weight_block_size + assert ( + granularity.get_block_size(shape, CompressionTargetTensor.ACTIVATION) + == activation_block_size + ) + # Omitting the target keeps the default weight behavior. + assert granularity.get_block_size(shape) == weight_block_size + + @pytest.mark.parametrize( "granularity_type,granularity_params,tensor_shape", [ @@ -338,6 +376,9 @@ def test_get_block_size_valid_conditions( ("per_channel", {"axis": -3}, (5, 10)), # Per block - axis out of bounds ("per_block", {"axis": 1, "block_size": 3}, (5,)), + ("per_block", {"axis": 2, "block_size": 5}, (10, 20)), + ("per_block", {"axis": 3, "block_size": 5}, (8, 16, 3)), + ("per_block", {"axis": -3, "block_size": 4}, (10, 20)), # Per block - None axis with integer block_size ("per_block", {"axis": None, "block_size": 5}, (10, 20)), # Per block - integer axis with list block_size