feat(quantization): support float16 scale_dtype

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`.

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

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]] = {

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,

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]:

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)

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,13 +231,26 @@ 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
         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
 
+        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.