diff --git a/changelog.d/169.float16-scale-dtype.added b/changelog.d/169.float16-scale-dtype.added new file mode 100644 index 00000000..69a05f14 --- /dev/null +++ b/changelog.d/169.float16-scale-dtype.added @@ -0,0 +1 @@ +Add `scale_dtype=torch.float16` to round quantization scales to FP16 in PyTorch simulation; Core AI and Core ML export raise `NotImplementedError`. diff --git a/docs/src/quantization/config.md b/docs/src/quantization/config.md index bef83410..0796aba2 100644 --- a/docs/src/quantization/config.md +++ b/docs/src/quantization/config.md @@ -27,8 +27,9 @@ We will first take a look at the {class}`~coreai_opt.quantization.spec.Quantizat - `qscheme`: symmetric or asymmetric quantization scheme (default: {class}`~coreai_opt.quantization.spec.QuantizationScheme`.SYMMETRIC). See [Symmetric vs asymmetric quantization](advanced.md#symmetric-vs-asymmetric-quantization) for how these schemes differ. - `granularity`: granularity for quantization scales/zp, per-block, per-channel, per-tensor (default: {class}`~coreai_opt.quantization.spec.PerTensorGranularity()`) - `qparam_calculator_cls`: algorithm used for calculating scales/zp, e.g. static, global min/max, moving average min/max etc. (defaults: {class}`~coreai_opt.quantization.spec.MovingAverageQParamsCalculator` for activations) -- `scale_dtype`: dtype for the quantization scale to be used. (Default: None, picked automatically based on quantization `dtype`) - Example: +- `scale_dtype`: how the quantization scale is stored. `None` (default) keeps the computed scale, except for FP4, which defaults to `torch.float8_e8m0fnu`. `torch.float8_e8m0fnu` rounds scales to powers of two (OCP MX), for FP4 and FP8. `torch.float16` rounds the scale to FP16 whatever the model dtype, for integer and FP8 dtypes; Core AI and Core ML export of it is not supported yet. With `torch.float16`, cast the model before `Quantizer.prepare`: a bf16 scale buffer would round the scales again, so casting the prepared model to bf16 raises `RuntimeError`. + +Example: ```python from coreai_opt.quantization import QuantizationSpec diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 0b86f051..02e67320 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -33,6 +33,9 @@ E8M0_EXPONENT_BIAS: Final[int] = 127 F32_MIN_NORMAL: Final[float] = 2**-126 # ~1.175e-38 +# Scale dtypes that the scale is rounded to before element codes are chosen. +ROUNDED_SCALE_DTYPES: Final[frozenset[torch.dtype]] = frozenset({torch.float16}) + # Maps a graph-mode aten ``OpOverload`` to the ``nn.Module`` type it corresponds # to. ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[torch.nn.Module]] = { diff --git a/src/coreai_opt/inspection/_bpw_utils.py b/src/coreai_opt/inspection/_bpw_utils.py index 75d96301..6c56b63b 100644 --- a/src/coreai_opt/inspection/_bpw_utils.py +++ b/src/coreai_opt/inspection/_bpw_utils.py @@ -11,6 +11,7 @@ import torch from torch.nn.utils.parametrize import ParametrizationList as _ParametrizationList +from coreai_opt._utils.torch_utils import ROUNDED_SCALE_DTYPES from coreai_opt.config.spec import ( CompressionSimulatorBase as _CompressionSimulatorBase, CompressionTargetTensor as _CompressionTargetTensor, @@ -42,7 +43,7 @@ torch.float4_e2m1fn_x2, } ) -_SUPPORTED_SCALE_DTYPES = frozenset({None, torch.float8_e8m0fnu}) +_SUPPORTED_SCALE_DTYPES = frozenset({None, torch.float8_e8m0fnu, *ROUNDED_SCALE_DTYPES}) _SUPPORTED_QSCHEMES = frozenset( { _QuantizationScheme.SYMMETRIC, diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index 5f79e093..6d4a4535 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -15,6 +15,7 @@ model_validator, ) +from coreai_opt._utils.torch_utils import ROUNDED_SCALE_DTYPES as _ROUNDED_SCALE_DTYPES from coreai_opt.common import CompressionType from coreai_opt.config.spec import CompressionSpec from coreai_opt.quantization.spec import ( @@ -49,8 +50,8 @@ class PalettizationSpec(CompressionSpec): If None, no quantization is applied to the LUT. When specified, only ``torch.int8``, ``torch.uint8``, ``torch.float8_e4m3fn``, and ``torch.float8_e5m2`` dtypes are supported, and granularity must be - ``PerTensorGranularity``. FP8 dtypes require symmetric quantization. - Default: None. + ``PerTensorGranularity``. FP8 dtypes require symmetric quantization, and + a float16 ``scale_dtype`` is not supported. Default: None. granularity: Defines how palettization is applied - per-tensor applies a single palette to the entire tensor, per-channel applies separate palettes to each channel. Default: PerTensorGranularity(). @@ -143,6 +144,13 @@ def validate_lut_qspec(self) -> "PalettizationSpec": "Use lut_qspec.qformulation=ZP instead." ) + if self.lut_qspec.scale_dtype in _ROUNDED_SCALE_DTYPES: + msg = ( + f"lut_qspec.scale_dtype={self.lut_qspec.scale_dtype} is not supported for " + f"palettization." + ) + raise ValueError(msg) + return self def model_dump_preserve_objects(self) -> dict[str, Any]: diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py index 7764eae8..da19c81b 100644 --- a/src/coreai_opt/quantization/quantizer.py +++ b/src/coreai_opt/quantization/quantizer.py @@ -24,7 +24,10 @@ from coreai_opt._utils.export_utils import ( validate_mmap_backend_and_device as _validate_mmap_backend_and_device, ) -from coreai_opt._utils.torch_utils import get_module_name as _get_module_name +from coreai_opt._utils.torch_utils import ( + ROUNDED_SCALE_DTYPES as _ROUNDED_SCALE_DTYPES, + 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 ( @@ -357,6 +360,11 @@ def prepare( The prepared model with fake quantization modules inserted, ready for calibration or training. This is a data-free PTQ compressed model. + Raises: + RuntimeError: If a spec sets ``scale_dtype=torch.float16`` and its scale buffers + cannot hold FP16 scales exactly, e.g. bf16 buffers created under + ``torch.set_default_dtype(torch.bfloat16)``. + Note: In graph mode, the returned ``fx.GraphModule`` supports calling ``.train()`` and ``.eval()``, but with limited effect: only dropout @@ -462,6 +470,33 @@ def _validate_export_supported_dtypes( ) raise NotImplementedError(msg) + def _validate_export_supported_qparams( + self, + model: nn.Module | fx.GraphModule | None, + backend: ExportBackend, + ) -> None: + """Reject CoreAI/CoreML export of qparams those backends cannot store yet. + + Export stores each scale at the model dtype, so it would store a float16 ``scale_dtype`` + at the model dtype instead of FP16 (rounding it again in a bf16 model). + """ + if backend == ExportBackend.Torch: + return + model_to_check = model if model is not None else self._model + unsupported_fq_names = [ + name + for name, mod in model_to_check.named_modules() + if isinstance(mod, FakeQuantizeImplBase) + and mod.qparams_calculator.scale_dtype in _ROUNDED_SCALE_DTYPES + ] + if unsupported_fq_names: + msg = ( + f"backend={backend} does not yet support a float16 scale_dtype. " + f"Affected FakeQuantize modules: {unsupported_fq_names}. Use " + f"backend=ExportBackend.Torch for torch-only inference." + ) + raise NotImplementedError(msg) + def finalize( self, model: nn.Module | fx.GraphModule | None = None, @@ -503,6 +538,8 @@ def finalize( NotImplementedError: If ``backend`` is CoreAI or CoreML and a quantization spec uses ``torch.int6`` or ``torch.uint6``, which are only supported for ``ExportBackend.Torch``. + NotImplementedError: If ``backend`` is CoreAI or CoreML and the model holds qparams + that backend cannot export yet, or scales rounded to a float16 ``scale_dtype``. Note: In graph mode, the returned ``fx.GraphModule`` supports calling @@ -516,6 +553,7 @@ def finalize( finalize frees the original dense weights. """ self._validate_mmap_dir_constraints(model, backend, mmap_dir) + self._validate_export_supported_qparams(model, backend) self._validate_no_persistent_observer_calculators(model, backend) self._validate_export_supported_dtypes(model, backend) return self._quantizer.finalize(model, backend, mmap_dir=mmap_dir) diff --git a/src/coreai_opt/quantization/spec/qparams_calculator.py b/src/coreai_opt/quantization/spec/qparams_calculator.py index 28aa49a7..96ff5b58 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -7,6 +7,8 @@ import warnings from abc import abstractmethod +from collections.abc import Callable +from typing import Self import torch import torch.nn as nn @@ -21,6 +23,7 @@ E8M0_EXPONENT_BIAS as _E8M0_EXPONENT_BIAS, E8M0_TARGET_MAX_POW2 as _E8M0_TARGET_MAX_POW2, F32_MIN_NORMAL as _F32_MIN_NORMAL, + ROUNDED_SCALE_DTYPES as _ROUNDED_SCALE_DTYPES, ) from .granularity import QuantizationGranularity @@ -53,7 +56,7 @@ def __init__( quant_max: int, range_calculator: RangeCalculatorBase, float_range: tuple[float | None, float | None], - scale_dtype: torch.dtype | None = None, # primarily for FP4 + scale_dtype: torch.dtype | None = None, # e8m0 or float16; see QuantizationSpec **kwargs, ): super().__init__() @@ -161,8 +164,13 @@ def _get_min_and_max_val(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torc # A one-sided float_range pairs a float_range bound (built at float32) # with a range-calculator bound at the input dtype. Cast both to the # input dtype so this mixed case has matching dtypes for the qparams op. - min_val = min_val.to(tensor.dtype) - max_val = max_val.to(tensor.dtype) + range_dtype = tensor.dtype + if self.scale_dtype in _ROUNDED_SCALE_DTYPES: + # A rounded scale is computed in at least fp32, so keep the float_range + # bound unrounded rather than round it to a half-precision input dtype. + range_dtype = torch.promote_types(range_dtype, torch.float32) + min_val = min_val.to(range_dtype) + max_val = max_val.to(range_dtype) return min_val, max_val def _compute_e8m0_scale(self, max_abs: torch.Tensor) -> torch.Tensor: @@ -209,7 +217,9 @@ def _compute_e8m0_scale(self, max_abs: torch.Tensor) -> torch.Tensor: def _compute_scale_zero_point_minval( self, tensor: torch.Tensor, min_val: torch.Tensor, max_val: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - """Supports two scale computation modes based on ``scale_dtype``: + """Return ``(scale, zero_point, minval)`` for the observed range. + + Supports three scale computation modes based on ``scale_dtype``: 1. **Default (``scale_dtype=None``)**: Uses torchao's ``choose_qparams_affine_with_min_max`` custom op to compute scale and @@ -221,6 +231,11 @@ def _compute_scale_zero_point_minval( scale = 2^(floor(log2(max_abs)) - target_max_pow2) where ``target_max_pow2`` comes from ``E8M0_TARGET_MAX_POW2``, which is the single source of truth for which dtypes support an e8m0 scale. + + 3. **Rounded scales** (``scale_dtype`` in ``ROUNDED_SCALE_DTYPES``): The + default-path scale, computed in fp32 whatever the input dtype, is clamped to the + largest finite ``scale_dtype`` value and rounded to it. An asymmetric zero point + is then derived from the rounded scale. """ # e8m0 path: power-of-2 scales @@ -228,6 +243,14 @@ def _compute_scale_zero_point_minval( max_abs = torch.maximum(torch.abs(min_val), torch.abs(max_val)) return self._compute_e8m0_scale(max_abs), None, None + rounds_scale = self.scale_dtype in _ROUNDED_SCALE_DTYPES + if rounds_scale: + # Start from an fp32 scale, so a half-precision input is not rounded twice. An + # fp64 input stays fp64, so its minval is not rounded to fp32. + range_dtype = torch.promote_types(min_val.dtype, torch.float32) + min_val = min_val.to(range_dtype) + max_val = max_val.to(range_dtype) + # Default path: torchao handles integer and FP8 dtypes scale, zero_point = choose_qparams_affine_with_min_max( min_val=min_val, @@ -241,6 +264,19 @@ def _compute_scale_zero_point_minval( zero_point_dtype=torch.int32, ) + if rounds_scale: + # Clamp first, so a large scale does not round to inf. + scale_dtype_max = torch.finfo(self.scale_dtype).max + scale = scale.clamp(max=scale_dtype_max).to(self.scale_dtype).to(torch.float32) + if rounds_scale and self.qscheme == QuantizationScheme.ASYMMETRIC: + # The zero point above came from the unrounded scale; recompute it from the + # rounded scale so the two match. + zero_point = torch.clamp( + self.quant_min - torch.round(min_val / scale), + self.quant_min, + self.quant_max, + ).to(torch.int32) + # ``minval`` is the minimum representable float value for the # observed range. The fake-quantize / export layer decides whether # the ZP or MINVAL formulation is in use and selects between @@ -336,8 +372,13 @@ def forward( On the first forward pass, initializes internal buffers using the observed tensor shape and device. Delegates the actual qparams calculation to ``compute_qparams``. + + Raises: + RuntimeError: If the ``scale`` buffer cannot hold rounded ``scale_dtype`` scales + exactly (see ``QuantizationSpec.scale_dtype``). """ self._resolve_axis(tensor.ndim) + self._validate_scale_buffer_dtype() scale = self.scale.clone() zero_point = self.zero_point.clone() if self.zero_point is not None else None @@ -393,7 +434,12 @@ def get_qparams(self) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | """ Return the computed scale, zero point and minval. For FP4/FP8/floating-point quantization, zero_point and minval are None. + + Raises: + RuntimeError: If the ``scale`` buffer cannot hold rounded ``scale_dtype`` scales + exactly (see ``QuantizationSpec.scale_dtype``). """ + self._validate_scale_buffer_dtype() if not self._initialized: warnings.warn( "Using default scale and zero point. Call forward pass to compute actual values.", @@ -410,6 +456,27 @@ def extra_repr(self) -> str: f"minval={self.minval}, export_mode={self._export_mode}" ) + def _apply(self, fn: Callable[[torch.Tensor], torch.Tensor], recurse: bool = True) -> Self: + """Cast or move the buffers, then raise if the cast made the ``scale`` buffer lossy.""" + super()._apply(fn, recurse) + self._validate_scale_buffer_dtype() + return self + + def _validate_scale_buffer_dtype(self) -> None: + """Raise if the ``scale`` buffer cannot hold a rounded ``scale_dtype`` scale exactly.""" + if self.scale_dtype in _ROUNDED_SCALE_DTYPES and self.scale.dtype not in ( + torch.float32, + torch.float64, + self.scale_dtype, + ): + msg = ( + f"The scale buffer is {self.scale.dtype}, which cannot hold " + f"scale_dtype={self.scale_dtype} scales exactly. Cast the model before " + f"Quantizer.prepare, not after, and keep torch.get_default_dtype() float32 " + f"while preparing." + ) + raise RuntimeError(msg) + class StatelessQParamsCalculatorBase(QParamsCalculatorBase): """Stateless base: no cached qparams; recomputed every forward. diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index 065670be..60c45d0e 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -20,6 +20,7 @@ from coreai_opt._utils.torch_utils import ( E8M0_TARGET_MAX_POW2 as _E8M0_TARGET_MAX_POW2, + ROUNDED_SCALE_DTYPES as _ROUNDED_SCALE_DTYPES, get_n_bits_from_dtype as _get_n_bits_from_dtype, is_float4_dtype as _is_float4_dtype, ) @@ -156,6 +157,7 @@ class QuantizationSpec(CompressionSpec): power-of-2 scale per OCP MX spec — ``scale = 2^(floor(log2(max_abs)) - target_max_pow2)``, with ``target_max_pow2`` of 2 for FP4 E2M1, 8 for FP8 E4M3, 15 for FP8 E5M2. + - ``scale_dtype=torch.float16`` (FP8): the ``None`` scale, rounded to FP16. **MINVAL formulation**, e.g. with 8 bit signed and unsigned fixed point types: @@ -303,8 +305,9 @@ class type: MinMaxRangeCalculator or custom registered class type - float_min < float_max scale_dtype (torch.dtype | None): Data type for quantization scale factors. - Controls whether scales are constrained to power-of-2 values (e8m0 format) - or allowed to be arbitrary floating-point values. + Controls whether scales are constrained to power-of-2 values (e8m0 format), + rounded to a floating-point storage dtype, or allowed to be arbitrary + floating-point values. Valid inputs: - None: Use default scale computation via torchao's @@ -312,14 +315,24 @@ class type: MinMaxRangeCalculator or custom registered class type For FP4, None is resolved to torch.float8_e8m0fnu automatically. - torch.float8_e8m0fnu: Power-of-2 scales following OCP Microscaling (MX) spec. Required for FP4 quantization, optional for FP8. + - torch.float16: Each scale is rounded to the nearest FP16 value, whatever + the model dtype. The tensor is then quantized with the rounded scale (an + asymmetric zero point is recomputed from it too), so the result matches a + format that stores its scales in FP16. Constraints: - FP4 (float4_e2m1fn_x2): scale_dtype must be torch.float8_e8m0fnu or None (defaults to e8m0) - FP8 (float8_e4m3fn, float8_e5m2): scale_dtype must be - torch.float8_e8m0fnu or None (defaults to None) - - Integer dtypes: scale_dtype must be None (defaults to None) + torch.float8_e8m0fnu, torch.float16 or None (defaults to None) + - Integer dtypes: scale_dtype must be torch.float16 or None (defaults to None) + - Export of torch.float16 to CoreAI or CoreML is not supported yet: + ``Quantizer.finalize`` raises ``NotImplementedError`` for those backends. + - With torch.float16, cast the model before ``Quantizer.prepare``, and keep + ``torch.get_default_dtype()`` float32 while preparing: a bf16 scale buffer + would round the cached scales again, so ``Quantizer.prepare``, a bf16 cast + of the prepared model, or the next forward raises ``RuntimeError``. Default: None @@ -541,28 +554,40 @@ def validate_scale_dtype(self) -> QuantizationSpec: Validate scale_dtype based on element dtype. Rules: - - Only None or torch.float8_e8m0fnu are supported. - - When scale_dtype is set, dtype must be a key of ``E8M0_TARGET_MAX_POW2`` - (the dtypes with an e8m0 target exponent); for any other dtype, including - all integer dtypes, scale_dtype must be None. + - Only None, torch.float8_e8m0fnu or torch.float16 are supported. + - torch.float16 is rejected for FP4 dtypes. + - torch.float8_e8m0fnu requires dtype to be a key of + ``E8M0_TARGET_MAX_POW2`` (the dtypes with an e8m0 target exponent); for + any other dtype, including all integer dtypes, it is rejected. - FP4 dtypes: scale_dtype is resolved to torch.float8_e8m0fnu by resolve_scale_dtype (before validator). """ if self.scale_dtype is None: return self + if self.scale_dtype in _ROUNDED_SCALE_DTYPES: + if _is_float4_dtype(self.dtype): + msg = ( + f"scale_dtype={self.scale_dtype} is not supported for dtype={self.dtype}. " + f"FP4 scales must be power-of-two (scale_dtype=torch.float8_e8m0fnu)." + ) + raise ValueError(msg) + return self + if self.scale_dtype != torch.float8_e8m0fnu: - raise ValueError( + msg = ( f"Unsupported scale_dtype: {self.scale_dtype}. " - f"Only None or torch.float8_e8m0fnu are supported." + f"Only None, torch.float8_e8m0fnu or torch.float16 are supported." ) + raise ValueError(msg) 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)}." + msg = ( + f"scale_dtype={self.scale_dtype} is not supported for dtype={self.dtype}. " + f"No e8m0 target exponent is defined for it; supported dtypes are " + f"{list(_E8M0_TARGET_MAX_POW2)}." ) + raise ValueError(msg) return self diff --git a/tests/inspection/test_bits_per_weight.py b/tests/inspection/test_bits_per_weight.py index 9bd9575d..f5783eff 100644 --- a/tests/inspection/test_bits_per_weight.py +++ b/tests/inspection/test_bits_per_weight.py @@ -53,6 +53,7 @@ _FP32_BITS = 32 _INT64_BITS = 64 _E8M0_BITS = 8 +_FP16_BITS = 16 # Sanity envelope for the qparam / LUT / bias overhead a sane config adds on top of the # nominal bit width, in bpw. @@ -465,6 +466,22 @@ def test_e8m0_scale_width(): _assert_bpw_is_plausible(result.bpw, 8) +def test_rounded_scale_dtype_width(): + """A float16 ``scale_dtype`` charges one 16-bit scale per block.""" + prepared = _prepare_eager_quant( + SimpleLinearModel(bias=False), + torch.int8, + granularity=_PER_BLOCK_32, + scale_dtype=torch.float16, + ) + result = bits_per_weight(prepared) + + expected_bits = _expected_quant_bits( + 8, _PER_BLOCK_32_BLOCKS, "symmetric", bias=False, scale_dtype_bits=_FP16_BITS + ) + assert result.total_bits == expected_bits + + def test_supported_quantization_settings_cover_the_library(): """Every dtype, scheme, and formulation the library defines must have a known storage cost.""" assert _SUPPORTED_QUANT_DTYPES == frozenset(QuantizationSpec.SUPPORTED_DTYPES) diff --git a/tests/palettization/test_palettization_spec.py b/tests/palettization/test_palettization_spec.py index 2d6b98a7..2fc6c6c8 100644 --- a/tests/palettization/test_palettization_spec.py +++ b/tests/palettization/test_palettization_spec.py @@ -267,6 +267,10 @@ def test_lut_qspec_validation(): ) ) + # Invalid: FP16 LUT scales, which palettization export does not store as FP16 + with pytest.raises(ValidationError, match="lut_qspec.scale_dtype=torch.float16 is not"): + PalettizationSpec(lut_qspec=QuantizationSpec(dtype=torch.int8, scale_dtype=torch.float16)) + @pytest.mark.parametrize( "granularity", [PerTensorGranularity(), PerGroupedChannelGranularity(axis=1, group_size=8)] diff --git a/tests/quantization/test_fake_quantize.py b/tests/quantization/test_fake_quantize.py index 9dbd0878..e3f66938 100644 --- a/tests/quantization/test_fake_quantize.py +++ b/tests/quantization/test_fake_quantize.py @@ -404,6 +404,43 @@ def test_float8_matches_torchao_primitives(dtype, granularity): torch.testing.assert_close(custom_dequant, torchao_dequant) +@pytest.mark.parametrize("weight_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + ("dtype", "qscheme"), + [ + (torch.int8, QuantizationScheme.SYMMETRIC), + (torch.int8, QuantizationScheme.ASYMMETRIC), + (torch.float8_e4m3fn, QuantizationScheme.SYMMETRIC), + ], +) +def test_float16_scale_dtype_matches_fp32_reference(dtype, qscheme, weight_dtype): + """The output matches an fp32 decode with the FP16 scale, cast to the weight dtype.""" + spec = QuantizationSpec( + dtype=dtype, + qscheme=qscheme, + granularity=PerChannelGranularity(axis=0), + scale_dtype=torch.float16, + ) + fq = QuantizationComponentFactory.create_fake_quantizer(spec, CompressionTargetTensor.WEIGHT) + weight = (torch.randn(8, 64) * 0.37 + 0.11).to(weight_dtype) + + out = fq(weight) + + scale, zero_point, _minval = fq.calculate_qparams() + weight_fp32 = weight.float() + if dtype.is_floating_point: + max_abs = weight_fp32.abs().amax(dim=1, keepdim=True) + assert torch.equal(scale, (max_abs / torch.finfo(dtype).max).half().float()) + codes = _quantize_affine_float8(weight_fp32, scale, dtype) + expected = _dequantize_affine_float8(codes, scale) + else: + codes = torch.round(weight_fp32 / scale) + zero_point + codes = torch.clamp(codes, spec.quant_min, spec.quant_max) + expected = (codes - zero_point) * scale + assert out.dtype == weight_dtype + assert torch.equal(out, expected.to(weight_dtype)) + + class TestFusedFakeQuantizeIntSTEMatchesTorchBuiltin: """Verify FusedFakeQuantizeIntSTE produces identical results to torch.fake_quantize_per_{tensor,channel}_affine for forward and backward.""" diff --git a/tests/quantization/test_qparams_calculator.py b/tests/quantization/test_qparams_calculator.py index fa424e30..43c3e051 100644 --- a/tests/quantization/test_qparams_calculator.py +++ b/tests/quantization/test_qparams_calculator.py @@ -19,6 +19,7 @@ DynamicQParamsCalculator, GlobalMinMaxQParamsCalculator, MovingAverageQParamsCalculator, + QParamsCalculatorBase, StatelessQParamsCalculatorBase, StaticQParamsCalculator, ) @@ -1078,3 +1079,136 @@ def test_scale_clamping(precision_dtype, dtype, granularity, compression_target) f"Scale should be less than {precision_dtype} eps floor {precision_eps}, " f"max scale={scale.max().item()}" ) + + +@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("qscheme", [QuantizationScheme.SYMMETRIC, QuantizationScheme.ASYMMETRIC]) +def test_float16_scale_dtype(qscheme, input_dtype): + """The scale is rounded to FP16 whatever the input dtype. + + An asymmetric zero point follows the rounded scale. + """ + spec = QuantizationSpec( + dtype=torch.int8, + qscheme=qscheme, + granularity=PerChannelGranularity(axis=0), + scale_dtype=torch.float16, + ) + calculator = QuantizationComponentFactory.create_qparams_calculator( + spec, CompressionTargetTensor.WEIGHT + ) + tensor = (torch.randn(1024, 16) * 0.37 + 0.11).to(input_dtype) + + scale, zero_point, _minval = calculator(tensor) + + min_val = tensor.float().amin(dim=1, keepdim=True).clamp(max=0) + max_val = tensor.float().amax(dim=1, keepdim=True).clamp(min=0) + if qscheme == QuantizationScheme.ASYMMETRIC: + unrounded_scale = (max_val - min_val) / 255 + expected_zero_point = torch.clamp(-128 - torch.round(min_val / scale), -128, 127) + else: + unrounded_scale = torch.maximum(-min_val, max_val) / 127.5 + expected_zero_point = torch.zeros_like(scale) + assert scale.dtype == torch.float32 + assert torch.equal(scale, unrounded_scale.half().float()) + assert torch.equal(zero_point, expected_zero_point.to(torch.int32)) + + +def test_float16_scale_dtype_clamps_to_the_largest_finite_value(): + """A scale beyond the FP16 range is clamped to the FP16 max, not rounded to inf.""" + spec = QuantizationSpec( + dtype=torch.int8, + granularity=PerTensorGranularity(), + scale_dtype=torch.float16, + ) + calculator = QuantizationComponentFactory.create_qparams_calculator( + spec, CompressionTargetTensor.WEIGHT + ) + + scale, _zero_point, _minval = calculator(torch.tensor([[1e9, -1e9]])) + + assert scale.item() == torch.finfo(torch.float16).max + + +def test_float16_scale_dtype_keeps_an_fp64_minval(): + """Only the scale is rounded, so an fp64 input keeps its fp64 minval.""" + spec = QuantizationSpec( + dtype=torch.uint8, + qscheme=QuantizationScheme.ASYMMETRIC, + granularity=PerTensorGranularity(), + scale_dtype=torch.float16, + ) + calculator = QuantizationComponentFactory.create_qparams_calculator( + spec, CompressionTargetTensor.WEIGHT + ) + # Neither bound is representable in fp32. + tensor = torch.tensor([[-0.1234567890123, 0.9876543210123]], dtype=torch.float64) + + _scale, _zero_point, minval = calculator(tensor) + + assert minval.dtype == torch.float64 + assert minval.item() == tensor.min().item() + + +def test_float16_scale_dtype_keeps_a_one_sided_float_range_bound_unrounded(): + """A float_range bound is not rounded to a bf16 input dtype before the FP16 rounding.""" + spec = QuantizationSpec( + dtype=torch.int8, + granularity=PerTensorGranularity(), + float_range=[None, 0.1], + scale_dtype=torch.float16, + ) + calculator = QuantizationComponentFactory.create_qparams_calculator( + spec, CompressionTargetTensor.WEIGHT + ) + # Non-negative values below the bound, so the scale comes from the 0.1 bound alone. + tensor = torch.rand(8, 16, dtype=torch.bfloat16) * 0.05 + + scale, _zero_point, _minval = calculator(tensor) + + # int8 symmetric: max_abs / 127.5, with max_abs the fp32 bound 0.1. + assert scale.item() == (torch.tensor(0.1) / 127.5).half().item() + + +def _create_calculator_under_default_dtype( + spec: QuantizationSpec, default_dtype: torch.dtype +) -> QParamsCalculatorBase: + """Build a weight calculator while ``torch.get_default_dtype()`` is ``default_dtype``.""" + previous_default_dtype = torch.get_default_dtype() + torch.set_default_dtype(default_dtype) + try: + return QuantizationComponentFactory.create_qparams_calculator( + spec, CompressionTargetTensor.WEIGHT + ) + finally: + torch.set_default_dtype(previous_default_dtype) + + +def test_float16_scale_dtype_rejects_a_bf16_default_dtype_buffer(): + """A bf16 default dtype makes the scale buffer bf16, which would round FP16 scales again.""" + spec = QuantizationSpec(dtype=torch.int8, scale_dtype=torch.float16) + calculator = _create_calculator_under_default_dtype(spec, torch.bfloat16) + + with pytest.raises(RuntimeError, match="torch.get_default_dtype"): + calculator(torch.randn(8, 16)) + + +def test_float16_scale_dtype_rejects_a_bf16_buffer_set_without_apply(): + """A bf16 buffer assigned directly, as FSDP buffer casting does, raises on the next read.""" + calculator = QuantizationComponentFactory.create_qparams_calculator( + QuantizationSpec(dtype=torch.int8, scale_dtype=torch.float16), + CompressionTargetTensor.WEIGHT, + ) + calculator(torch.randn(8, 16)) + + calculator.scale = calculator.scale.bfloat16() + + with pytest.raises(RuntimeError, match="Cast the model before Quantizer.prepare"): + calculator.get_qparams() + + +def test_default_scale_dtype_accepts_a_bf16_scale_buffer(): + """The buffer check applies only to rounded scale dtypes, so default scales still run.""" + calculator = _create_calculator_under_default_dtype(QuantizationSpec(), torch.bfloat16) + + calculator(torch.randn(8, 16)) diff --git a/tests/quantization/test_quantization.py b/tests/quantization/test_quantization.py index 5011b8fc..a7d23f34 100644 --- a/tests/quantization/test_quantization.py +++ b/tests/quantization/test_quantization.py @@ -27,6 +27,7 @@ from coreai_opt.quantization.config.quantization_config import QATSchedule from coreai_opt.quantization.spec import ( PerBlockGranularity, + PerChannelGranularity, PerTensorGranularity, QuantizationScheme, default_activation_quantization_spec, @@ -665,6 +666,84 @@ def test_finalize_supports_6bit_only_for_torch_backend(self, execution_mode, dty torch.testing.assert_close(finalized(example_input), prepared_output) +class TestFloat16ScaleDtype: + """Simulation and export of a float16 ``scale_dtype``.""" + + @staticmethod + def _make_config(execution_mode: str) -> QuantizerConfig: + weight_spec = QuantizationSpec( + dtype=torch.int8, + granularity=PerChannelGranularity(axis=0), + scale_dtype=torch.float16, + ) + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_spec}, + op_input_spec=None, + op_output_spec=None, + ), + ).set_execution_mode(execution_mode) + + @staticmethod + def _weight_scales(prepared_model: nn.Module) -> list[torch.Tensor]: + return [ + m.calculate_qparams()[0] + for m in prepared_model.modules() + if isinstance(m, FakeQuantizeImplBase) + ] + + @pytest.mark.parametrize("backend", [ExportBackend.CoreAI, ExportBackend.CoreML]) + def test_finalize_rejects_non_torch_backends(self, execution_mode, backend): + quantizer = Quantizer(SimpleLinearModel(), self._make_config(execution_mode)) + quantizer.prepare((torch.randn(4, 64),)) + + with pytest.raises(NotImplementedError, match="float16 scale_dtype"): + quantizer.finalize(backend=backend) + + @pytest.mark.parametrize("model_dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_simulation_in_any_model_dtype(self, execution_mode, model_dtype): + """Cached scales are fp32 scales rounded to FP16, and Torch folding keeps the outputs.""" + model = SimpleLinearModel().to(model_dtype) + # int8 symmetric per-channel: max_abs / 127.5, computed in fp32 and rounded once. + expected_scales = [ + (layer.weight.float().abs().amax(dim=1, keepdim=True) / 127.5).half().float() + for layer in (model.l1, model.l2) + ] + inputs = torch.randn(4, 64, dtype=model_dtype) + quantizer = Quantizer(model, self._make_config(execution_mode)) + prepared_model = quantizer.prepare((inputs,)) + + out = prepared_model(inputs) + + assert out.dtype == model_dtype + scales = self._weight_scales(prepared_model) + for scale, expected_scale in zip(scales, expected_scales, strict=True): + assert torch.equal(scale, expected_scale) + assert torch.equal(quantizer.finalize(backend=ExportBackend.Torch)(inputs), out) + + def test_bf16_cast_of_the_prepared_model_raises(self, execution_mode): + """The cast itself raises, since a bf16 scale buffer would round the FP16 scales.""" + quantizer = Quantizer(SimpleLinearModel(), self._make_config(execution_mode)) + prepared_model = quantizer.prepare((torch.randn(4, 64),)) + + with pytest.raises(RuntimeError, match="Cast the model before Quantizer.prepare"): + prepared_model.to(torch.bfloat16) + + @pytest.mark.parametrize("cast_dtype", [torch.float16, torch.float64]) + def test_fp16_or_fp64_cast_keeps_the_scales(self, execution_mode, cast_dtype): + """Casting the prepared model to fp16 or fp64 keeps the cached FP16 scales exactly.""" + inputs = torch.randn(4, 64) + quantizer = Quantizer(SimpleLinearModel(), self._make_config(execution_mode)) + prepared_model = quantizer.prepare((inputs,)) + scales = [scale.clone() for scale in self._weight_scales(prepared_model)] + + prepared_model.to(cast_dtype) + + for cast_scale, scale in zip(self._weight_scales(prepared_model), scales, strict=True): + assert torch.equal(cast_scale.float(), scale) + prepared_model(inputs.to(cast_dtype)) + + 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 db4922d3..4d6cedaa 100644 --- a/tests/quantization/test_quantization_spec.py +++ b/tests/quantization/test_quantization_spec.py @@ -884,6 +884,9 @@ def test_partial_defaults(): ), (torch.float4_e2m1fn_x2, None, torch.float8_e8m0fnu), (torch.float4_e2m1fn_x2, torch.float8_e8m0fnu, torch.float8_e8m0fnu), + (torch.int8, torch.float16, torch.float16), + (torch.float8_e4m3fn, torch.float16, torch.float16), + (torch.uint8, "float16", torch.float16), ], ) def test_valid_scale_dtype(dtype, scale_dtype, expected_scale_dtype): @@ -900,9 +903,10 @@ def test_valid_scale_dtype(dtype, scale_dtype, expected_scale_dtype): "dtype,scale_dtype", [ (torch.float8_e4m3fn, torch.float32), - (torch.int8, torch.float16), + (torch.int8, torch.bfloat16), (torch.float8_e4m3fn, "float8_typo"), (torch.float4_e2m1fn_x2, torch.float32), + (torch.float4_e2m1fn_x2, torch.float16), ( torch.int4, torch.float8_e8m0fnu,