diff --git a/changelog.d/161.added b/changelog.d/161.added new file mode 100644 index 00000000..3c586c1f --- /dev/null +++ b/changelog.d/161.added @@ -0,0 +1 @@ +Add `QuantizationSpec.qparam_qspec` to simulate double-scale quantization in PyTorch, such as NVFP4 and GGUF-style double scales, and `scale_dtype=torch.float16` to round scales to FP16. diff --git a/docs/architecture_notes/graph_annotation.md b/docs/architecture_notes/graph_annotation.md index 15689295..1cbcc8d6 100644 --- a/docs/architecture_notes/graph_annotation.md +++ b/docs/architecture_notes/graph_annotation.md @@ -44,12 +44,12 @@ The fields are exactly the settable inputs of `coreai_opt...QuantizationSpec`, p `_qspec_constraints._FIELD_POLICY` is the authority. The reasoning: -| Field | Policy | Why | -| ----------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `DTYPE` | Highest-priority proposal wins | Two tensors sharing an observer can't have different dtypes, and there's no safe join — int4 and int8 aren't compatible in either direction. Config precedence decides. | -| `QFORMULATION`, `FAKE_QUANTIZE_CLS`, `QPARAM_CALCULATOR_CLS`, `RANGE_CALCULATOR_CLS`, `SCALE_DTYPE`, `QSCHEME`, `GRANULARITY` | Highest-priority proposal wins | Ordinary config choices. | -| `FLOAT_RANGE` | Widest of the proposals, per bound | A shared observer must cover every member's data, so a pin survives only while every member agrees on it. Priority is deliberately not consulted: covering the data is a correctness constraint, not a preference. This is also where "fixed qparams" lives — a fixed observer is one whose `float_range` is pinned on both bounds. | -| `QUANTIZATION_TARGET` | Highest-priority proposal wins | This is not actually set as part of the config, but comes from whether the tensor is a weight or an activation. If we share a qspec between a weight and an activation, using the higher-priority one creates a coherent set of fields that matches at least one target's original provisional qspec. | +| Field | Policy | Why | +| --------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `DTYPE` | Highest-priority proposal wins | Two tensors sharing an observer can't have different dtypes, and there's no safe join — int4 and int8 aren't compatible in either direction. Config precedence decides. | +| `QFORMULATION`, `FAKE_QUANTIZE_CLS`, `QPARAM_CALCULATOR_CLS`, `RANGE_CALCULATOR_CLS`, `SCALE_DTYPE`, `QPARAM_QSPEC`, `QSCHEME`, `GRANULARITY` | Highest-priority proposal wins | Ordinary config choices. | +| `FLOAT_RANGE` | Widest of the proposals, per bound | A shared observer must cover every member's data, so a pin survives only while every member agrees on it. Priority is deliberately not consulted: covering the data is a correctness constraint, not a preference. This is also where "fixed qparams" lives — a fixed observer is one whose `float_range` is pinned on both bounds. | +| `QUANTIZATION_TARGET` | Highest-priority proposal wins | This is not actually set as part of the config, but comes from whether the tensor is a weight or an activation. If we share a qspec between a weight and an activation, using the higher-priority one creates a coherent set of fields that matches at least one target's original provisional qspec. | Some properties are deliberately *not* fields. `quant_min` / `quant_max` / `n_bits` / `target_dtype` are computed by `QuantizationSpec` from `dtype` and `qscheme`, so they fall out of reassembly; reconciling them alongside their own inputs would let them contradict those inputs. `ch_axis` lives inside `GRANULARITY`. diff --git a/docs/src/quantization/config.md b/docs/src/quantization/config.md index a884b6e7..d09e2474 100644 --- a/docs/src/quantization/config.md +++ b/docs/src/quantization/config.md @@ -27,8 +27,10 @@ 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. +- `qparam_qspec`: a nested `QuantizationSpec` that quantizes the per-block scales themselves (and the per-block minvals, with `qformulation="minval"`). See [Quantizing the quantization parameters](#example-quantizing-the-quantization-parameters-qparam_qspec). (Default: None) + +Example: ```python from coreai_opt.quantization import QuantizationSpec @@ -267,6 +269,133 @@ flowchart LR style output fill:#d4edda,stroke:#28a745 ``` +### Example: Quantizing the quantization parameters (`qparam_qspec`) + +`qparam_qspec` quantizes the quantization parameters themselves. Quantizing a tensor per block produces a grid of qparams with one scale per block, plus one minval per block with `qformulation="minval"`. The nested `QuantizationSpec` quantizes that grid like any other tensor, and the tensor is then quantized against the decoded qparams, as a format that stores its qparams at low precision reconstructs them. + +Every `QuantizationSpec` field works in the nested spec as usual: + +- `granularity`: `PerBlockGranularity(block_size=8)` shares one set of nested qparams across 8 consecutive blocks. Left at `axis=None`, it follows the outer spec's block axis. +- `scale_dtype` is the dtype of the super-scales; for example, `torch.float16` gives FP16 super-scales like GGUF Q4_K's `d` and `dmin`. +- Calibration, observer toggles and QAT schedules apply to the nested qparams the same way as to any other qparams. + +For example, int4 blocks whose scales are int8 codes under one scale per output channel: + +```python +from coreai_opt.quantization import QuantizationSpec +from coreai_opt.quantization.spec import PerBlockGranularity, PerChannelGranularity + +int4_weight = QuantizationSpec( + dtype=torch.int4, + granularity=PerBlockGranularity(block_size=32), + qparam_qspec=QuantizationSpec( + dtype=torch.int8, + granularity=PerChannelGranularity(axis=0), + ), +) +``` + +Layers whose blocks do not divide into the nested groups are left unquantized, with a warning. Core AI and Core ML export of a spec with `qparam_qspec` is not supported yet: `finalize()` raises `NotImplementedError`. Use the prepared model for torch inference. + +Double scale can be used to specify several popular quantization formats, such as GGUF's Q4_K and NVFP4. + +#### GGUF double scale + +GGUF k-quants such as Q4_K store 4-bit weights in blocks of 32, each with a scale and a min (`qformulation="minval"`). The scales and the mins are quantized in turn, under separate super-scales shared by 8 blocks (256 weights), so layers whose input channels are not a multiple of 256 are left unquantized. The super-scales are FP16 (`scale_dtype=torch.float16` in the nested spec), whatever the model dtype. Q4_K itself stores 6-bit codes, which needs `torch.uint6`; it is not supported yet, so this example uses uint8 codes. + +```python +# programmatic + +from coreai_opt.quantization.config import QuantizerConfig, ModuleQuantizerConfig +from coreai_opt.quantization import QuantizationSpec +from coreai_opt.quantization.spec import PerBlockGranularity + +gguf_weight = QuantizationSpec( + dtype=torch.uint4, + qscheme="asymmetric", + qformulation="minval", + granularity=PerBlockGranularity(block_size=32), + qparam_qspec=QuantizationSpec( + dtype=torch.uint8, + qscheme="asymmetric", + granularity=PerBlockGranularity(block_size=8), + scale_dtype=torch.float16, + ), +) + +config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_input_spec=None, + op_output_spec=None, + op_state_spec={"weight": gguf_weight}, + ) +) +``` + +```yaml +# yaml config +quantization_spec: + spec1: &gguf_weight + dtype: uint4 + qscheme: asymmetric + qformulation: minval + granularity: { type: per_block, block_size: 32 } + qparam_qspec: + dtype: uint8 + qscheme: asymmetric + granularity: { type: per_block, block_size: 8 } + scale_dtype: float16 + +quantization_config: + global_config: + op_input_spec: null + op_output_spec: null + op_state_spec: + weight: *gguf_weight +``` + +#### NVFP4 + +NVFP4 stores FP4 E2M1 values in blocks of 16, each block with an FP8 E4M3 scale, under one FP32 scale per tensor (4.5 bits per weight, plus 32 bits per tensor). For weights and activations: + +```python +# programmatic + +from coreai_opt.quantization.config import QuantizerConfig, ModuleQuantizerConfig +from coreai_opt.quantization import QuantizationSpec +from coreai_opt.quantization.spec import PerBlockGranularity, PerTensorGranularity + +nvfp4_weight = QuantizationSpec( + dtype=torch.float4_e2m1fn_x2, + granularity=PerBlockGranularity(block_size=16), + # FP8 E4M3 block scales, under one FP32 scale per tensor + qparam_qspec=QuantizationSpec( + dtype=torch.float8_e4m3fn, + granularity=PerTensorGranularity(), + ), +) + +# NVFP4 activations: block scales recomputed every forward, under a tensor scale +# that is calibrated and then frozen +nvfp4_activation = QuantizationSpec( + dtype=torch.float4_e2m1fn_x2, + granularity=PerBlockGranularity(axis=-1, block_size=16), + qparam_calculator_cls="dynamic", + qparam_qspec=QuantizationSpec( + dtype=torch.float8_e4m3fn, + granularity=PerTensorGranularity(), + ), +) + +config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_input_spec={"*": nvfp4_activation}, + op_output_spec=None, + op_state_spec={"weight": nvfp4_weight}, + ) +) +``` + ### Example: Skip quantization for a specific layer type Let's say we want to skip all `linear` layers in the model from being quantized. Since linear is commonly used as `nn.Module` in model definition, there are two equivalent ways to do this: diff --git a/docs/src/quantization/overview.md b/docs/src/quantization/overview.md index fb2231da..81fbb5da 100644 --- a/docs/src/quantization/overview.md +++ b/docs/src/quantization/overview.md @@ -164,6 +164,8 @@ Broadly speaking, the following dtypes can be used for quantization: A few other dtypes (e.g. FP8_E5M2) may also be available to experiment with; however, they may have limited support on the Core AI runtime. +The quantization parameters can themselves be quantized with a nested spec (`qparam_qspec`), for weights and activations. These are available for simulation in PyTorch; Core AI and Core ML export is not supported yet. See [Quantizing the quantization parameters](config.md#example-quantizing-the-quantization-parameters-qparam_qspec). + ## Two Execution Modes: graph and eager The {class}`~coreai_opt.quantization.Quantizer` has two internal implementations: diff --git a/docs/src/utils/mixed_precision.md b/docs/src/utils/mixed_precision.md index ba73e900..8c08d985 100644 --- a/docs/src/utils/mixed_precision.md +++ b/docs/src/utils/mixed_precision.md @@ -103,6 +103,7 @@ The utility supports: - **Eager-mode integer weight quantization**: any spec-supported integer dtype, any quantization scheme (symmetric or asymmetric), at any granularity. - **Eager-mode floating-point weight quantization**: FP8 (`e4m3`, `e5m2`) and FP4 (`e2m1`). +- **Quantized qparams** (`qparam_qspec`): each block's scale (and minval) costs its code width, plus the nested spec's own qparams, counted the same way. For example, NVFP4 (FP4 elements, FP8 E4M3 scales per block of 16, one FP32 tensor scale) comes to 4.5 BPW, plus 32 bits per tensor. - **Palettization** at any spec-supported `n_bits`, including a quantized LUT. Sub-byte payloads and zero-points are packed at `n_bits` with no padding. diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 0b86f051..d948a1e4 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]] = { @@ -305,6 +308,11 @@ def is_tensor_positive(tensor: torch.Tensor) -> bool: return bool(tensor.gt(0).all()) +def nonzero_scale(scale: torch.Tensor) -> torch.Tensor: + """Return ``scale`` with zeros replaced by 1, as a safe divisor.""" + return torch.where(scale == 0, 1.0, scale) + + def is_float_quant_dtype(dtype: torch.dtype) -> bool: """Check if dtype is a floating-point (FP4/FP8) quantization dtype. diff --git a/src/coreai_opt/inspection/_bpw_utils.py b/src/coreai_opt/inspection/_bpw_utils.py index 8c91df73..7ca70230 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, @@ -40,7 +41,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, @@ -186,7 +187,7 @@ def tensor_storage_bits(weight: torch.Tensor, compressor: WeightCompressor | Non return full_precision_bits(weight) if isinstance(compressor, _FakePalettizeImplBase): return _palettized_bits(weight, compressor) - return _quantized_bits(weight, compressor) + return _quantized_tensor_bits(weight.shape, compressor) def full_precision_bits(tensor: torch.Tensor) -> int: @@ -194,37 +195,44 @@ def full_precision_bits(tensor: torch.Tensor) -> int: return int(tensor.numel() * tensor.element_size() * _BITS_PER_BYTE) -def _quantized_bits(weight: torch.Tensor, quantization_fq: _FakeQuantizeImplBase) -> int: - """Return the storage cost of a quantized weight including scale / offset overhead. +def _quantized_tensor_bits( + shape: torch.Size, fq: _FakeQuantizeImplBase, qparams_nonnegative: bool = False +) -> int: + """Return the storage cost of a tensor quantized by ``fq``, including its qparams. + + Qparams quantized per ``qparam_qspec`` cost their own quantized storage, recursively. Args: - weight (torch.Tensor): The dense original weight tensor. - quantization_fq (FakeQuantizeImplBase): The weight fake-quantize parametrization. + shape (torch.Size): Shape of the quantized tensor. + fq (FakeQuantizeImplBase): The fake quantizer of the tensor. + qparams_nonnegative (bool): Whether the tensor holds non-negative qparams + quantized per ``qparam_qspec``, for :func:`_offset_bits`. Returns: - int: ``payload_bits + scale_bits + offset_bits`` in bits, where the payload - is ``numel * n_bits`` and the per-block scale and offset overhead is - amortized across the weight. + int: ``payload_bits + scale_bits + offset_bits`` in bits. Note: - ``num_blocks`` is read directly from the materialized - ``qparams_calculator.scale`` buffer. When that buffer - is not yet materialized,``num_blocks`` is - derived analytically from ``granularity.get_block_size``. + The per-block qparam grid shape is read directly from the materialized + ``qparams_calculator.scale`` buffer. When that buffer is not yet materialized, + it is derived analytically from ``fq.get_block_size``. """ - num_elements = weight.numel() - - scale = quantization_fq.qparams_calculator.scale + scale = fq.qparams_calculator.scale if scale is not None and scale.numel() > 0: - num_blocks = scale.numel() + qparam_shape = scale.shape else: - block_size = quantization_fq.granularity.get_block_size(weight.shape) - num_blocks = num_elements // math.prod(block_size) + block_size = fq.get_block_size(shape) + qparam_shape = torch.Size( + dim // block for dim, block in zip(shape, block_size, strict=True) + ) - payload_bits = num_elements * quantization_fq.n_bits - scale_bits = num_blocks * _scale_storage_bits(quantization_fq) + payload_bits = math.prod(shape) * fq.n_bits + scale_fq = fq.qparams_calculator.scale_fake_quantize + if scale_fq is not None: + scale_bits = _quantized_tensor_bits(qparam_shape, scale_fq, qparams_nonnegative=True) + else: + scale_bits = math.prod(qparam_shape) * _scale_storage_bits(fq) - return int(payload_bits + scale_bits + _offset_bits(quantization_fq, num_blocks)) + return int(payload_bits + scale_bits + _offset_bits(fq, qparam_shape, qparams_nonnegative)) def _scale_storage_bits(fq: _FakeQuantizeImplBase) -> int: @@ -240,14 +248,31 @@ def _float_qparam_bits(fq: _FakeQuantizeImplBase) -> int: return int(dtype.itemsize * _BITS_PER_BYTE) -def _offset_bits(fq: _FakeQuantizeImplBase, num_blocks: int) -> int: - """Return the per-block dequantization offset cost of a quantized weight, in bits. +def _offset_bits( + fq: _FakeQuantizeImplBase, qparam_shape: torch.Size, qparams_nonnegative: bool +) -> int: + """Return the per-block dequantization offset cost of a quantized tensor, in bits. - - ``ZP``: an asymmetric weight stores a per-block ``zero_point``, packed at ``n_bits``. + - ``ZP``: an asymmetric tensor stores a per-block ``zero_point``, packed at ``n_bits``. - ``MINVAL``: the formulation stores a per-block ``minval`` float instead of a - zero-point, so it costs a full float per block. + zero-point, so it costs a full float per block, unless ``qparam_qspec`` quantizes + it. + + Qparams quantized per ``qparam_qspec`` are non-negative, so an asymmetric nested spec + whose range minimum is not pinned below 0 by ``float_range`` always gets the offset + ``quant_min`` (ZP) or 0 (MINVAL), which needs no storage. """ + if ( + qparams_nonnegative + and fq.qscheme == _QuantizationScheme.ASYMMETRIC + and fq.qparams_calculator.float_range[0] in (None, 0) + ): + return 0 + num_blocks = int(math.prod(qparam_shape)) if fq.qformulation == _QuantizationFormulation.MINVAL: + minval_fq = fq.qparams_calculator.minval_fake_quantize + if minval_fq is not None: + return _quantized_tensor_bits(qparam_shape, minval_fq, qparams_nonnegative=True) return num_blocks * _float_qparam_bits(fq) if fq.qscheme == _QuantizationScheme.ASYMMETRIC: return num_blocks * fq.n_bits diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index 24baeb99..f8aa17f7 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -14,6 +14,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 ( @@ -48,7 +49,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. + ``PerTensorGranularity``. FP8 dtypes require symmetric quantization, and + neither ``qparam_qspec`` nor a float16 ``scale_dtype`` is supported. Default: None. granularity: Defines how palettization is applied - per-tensor applies a single palette to the entire tensor, per-channel applies separate @@ -132,6 +134,17 @@ def validate_lut_qspec(self) -> "PalettizationSpec": "Use lut_qspec.qformulation=ZP instead." ) + if self.lut_qspec.qparam_qspec is not None: + msg = "lut_qspec.qparam_qspec is not supported for palettization." + raise ValueError(msg) + + 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/_graph/_provisional_qspec_generation.py b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py index e19deabe..2e971c27 100644 --- a/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py +++ b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py @@ -65,6 +65,7 @@ FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls", FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls", FieldName.SCALE_DTYPE: "scale_dtype", + FieldName.QPARAM_QSPEC: "qparam_qspec", } diff --git a/src/coreai_opt/quantization/_graph/_qspec_constraints.py b/src/coreai_opt/quantization/_graph/_qspec_constraints.py index 8b19d419..f4bb227f 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_constraints.py +++ b/src/coreai_opt/quantization/_graph/_qspec_constraints.py @@ -318,6 +318,7 @@ def _policy_float_range_union(proposals: Sequence[FieldValue]) -> FieldValue: FieldName.QPARAM_CALCULATOR_CLS: _policy_priority_wins, FieldName.RANGE_CALCULATOR_CLS: _policy_priority_wins, FieldName.SCALE_DTYPE: _policy_priority_wins, + FieldName.QPARAM_QSPEC: _policy_priority_wins, # Covering every member's values is a correctness constraint, not a # preference, so this one unions instead of deferring to priority. FieldName.FLOAT_RANGE: _policy_float_range_union, diff --git a/src/coreai_opt/quantization/_graph/_qspec_resolution.py b/src/coreai_opt/quantization/_graph/_qspec_resolution.py index 65c1e4f8..9f6cc462 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_resolution.py +++ b/src/coreai_opt/quantization/_graph/_qspec_resolution.py @@ -200,6 +200,7 @@ def _shared_spec_pointing_at(anchor: NodeSlot) -> _SharedQuantizationSpec: FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls", FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls", FieldName.SCALE_DTYPE: "scale_dtype", + FieldName.QPARAM_QSPEC: "qparam_qspec", } diff --git a/src/coreai_opt/quantization/_graph/_qspec_types.py b/src/coreai_opt/quantization/_graph/_qspec_types.py index 199849a9..b91fa9eb 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_types.py +++ b/src/coreai_opt/quantization/_graph/_qspec_types.py @@ -63,6 +63,7 @@ class FieldName(enum.Enum): QPARAM_CALCULATOR_CLS = enum.auto() RANGE_CALCULATOR_CLS = enum.auto() SCALE_DTYPE = enum.auto() + QPARAM_QSPEC = enum.auto() # nested spec that quantizes the scale / minval # Weight or activation, set by which config dict the spec came from. Not a # QuantizationSpec attribute, but an input to construct_partial alongside # them, so resolution needs it to rebuild the observer. diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py index 8bd7aa8d..71cd2625 100644 --- a/src/coreai_opt/quantization/quantizer.py +++ b/src/coreai_opt/quantization/quantizer.py @@ -23,7 +23,11 @@ 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, + get_parent_module_and_attr_name as _get_parent_module_and_attr_name, +) from coreai_opt.common import ExportBackend from coreai_opt.quantization._eager import EagerQuantizer as _EagerQuantizer from coreai_opt.quantization._export_utils import ( @@ -38,6 +42,7 @@ QuantizerConfig, ) from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase +from coreai_opt.quantization.spec.qparams_calculator import QParamsCalculatorBase class Quantizer(_BaseQuantizer): @@ -221,8 +226,12 @@ def _maybe_apply_qat_schedule(self) -> None: """Apply observer/fake-quant state for the current step count.""" for fq_module, schedule in self._fq_to_schedule.items(): state = schedule._compute_state(self._step_count) - fq_module.enable_observer(state.obs_on) - fq_module.enable_fake_quant(state.fq_on) + # ``apply`` also reaches the fake quantizers a ``qparam_qspec`` nests in the qparams + # calculator, so a nested scale that a dynamic calculator keeps using freezes too. + fq_module.apply(_torchao_enable_observer if state.obs_on else _torchao_disable_observer) + fq_module.apply( + _torchao_enable_fake_quant if state.fq_on else _torchao_disable_fake_quant + ) def _validate_no_schedule_configured(self) -> None: """Raise RuntimeError if any FQ modules have a qat_schedule.""" @@ -435,6 +444,41 @@ def _validate_no_persistent_observer_calculators( f"backend=ExportBackend.Torch for torch-only inference." ) + 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 drop the codes and super-scales of + a ``qparam_qspec``, and 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_fake_quantize is not None + or mod.qparams_calculator.scale_dtype in _ROUNDED_SCALE_DTYPES + ) + and not isinstance( + _get_parent_module_and_attr_name(model_to_check, name)[0], + QParamsCalculatorBase, + ) + ] + if unsupported_fq_names: + msg = ( + f"backend={backend} does not yet support quantized qparams (qparam_qspec) or a " + f"float16 scale_dtype. 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, @@ -472,6 +516,12 @@ def finalize( Returns: The finalized quantized model ready for deployment on the target backend. + Raises: + NotImplementedError: If ``backend`` is CoreAI or CoreML and the model holds qparams + that backend cannot export yet: qparams calculators that recompute every + forward (e.g. dynamic quantization), qparams quantized by a + ``qparam_qspec``, or scales rounded to a float16 ``scale_dtype``. + Note: In graph mode, the returned ``fx.GraphModule`` supports calling ``.train()`` and ``.eval()``, but with limited effect: only dropout @@ -484,6 +534,9 @@ def finalize( finalize frees the original dense weights. """ self._validate_mmap_dir_constraints(model, backend, mmap_dir) + # First, so a nested fake quantizer with a dynamic calculator is reported as part of its + # qparam_qspec rather than by name under the dynamic-quantization error. + self._validate_export_supported_qparams(model, backend) self._validate_no_persistent_observer_calculators(model, backend) return self._quantizer.finalize(model, backend, mmap_dir=mmap_dir) diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index 408e4192..8853a0e4 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -14,6 +14,7 @@ ) from .fake_quantize import FakeQuantizeImplBase +from .qformulation import QuantizationFormulation from .qparams_calculator import ( DynamicQParamsCalculator, MovingAverageQParamsCalculator, @@ -116,6 +117,15 @@ def create_qparams_calculator( "float_range": spec.float_range, "scale_dtype": spec.scale_dtype, } + if spec.qparam_qspec is not None: + common_args["scale_fake_quantize"] = cls.create_fake_quantizer( + spec.qparam_qspec, quantization_target + ) + if spec.qformulation == QuantizationFormulation.MINVAL: + # A separate instance, so minval gets its own super-scales. + common_args["minval_fake_quantize"] = cls.create_fake_quantizer( + spec.qparam_qspec, quantization_target + ) # Automatically detect and include any extra arguments extra_args = spec.get_extra_args() diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index 49a47954..7a007002 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/src/coreai_opt/quantization/spec/fake_quantize.py @@ -24,6 +24,7 @@ is_float4_dtype as _is_float4_dtype, is_float8_dtype as _is_float8_dtype, is_float_quant_dtype as _is_float_quant_dtype, + nonzero_scale as _nonzero_scale, ) from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes @@ -131,10 +132,9 @@ def disable_observer(self) -> None: super().disable_observer() def enable_observer(self, enabled: bool = True) -> None: - """Inverse of ``disable_observer``: ignore ``enabled=False`` when the - qparams calculator is stateless. Covers callers that invoke - ``enable_observer(False)`` directly (e.g. the QAT scheduler at - ``quantizer.py:_maybe_apply_qat_schedule``); ``disable_observer()`` + """Enable the observer; ignore ``enabled=False`` for a stateless qparams calculator. + + Covers callers that invoke ``enable_observer(False)`` directly; ``disable_observer()`` itself routes through the override above. """ if not enabled and self.is_stateless: @@ -688,11 +688,14 @@ def _quantize_int( - ZP: ``quant_offset = zero_point``, ``float_offset = 0`` - MINVAL: ``quant_offset = quant_min``, ``float_offset = minval`` + A block whose scale is zero, which a quantized scale can decode to, is divided by 1 + instead, so :func:`_dequantize_int` returns its offset rather than NaN. + The quantized tensor remains in FP dtype. The quantized tensor as well as a mask tensor marking positions in the tensor which were clamped are both returned. """ - result = (tensor - float_offset) / scale + result = (tensor - float_offset) / _nonzero_scale(scale) result.round_() result.add_(quant_offset) mask = result >= quant_min @@ -727,10 +730,13 @@ def _quantize_float( """ Float quantization: cast_decast(clamp(tensor / scale, min, max)) + A block whose scale is zero, which a quantized scale can decode to, is divided by 1 + instead, so :func:`_dequantize_float` returns zeros rather than NaN. + The quantized tensor as well as a mask tensor marking positions in the tensor which were clamped are both returned. """ - result = tensor / scale + result = tensor / _nonzero_scale(scale) mask = result >= quant_min mask &= result <= quant_max result.clamp_(min=quant_min, max=quant_max) diff --git a/src/coreai_opt/quantization/spec/qparams_calculator.py b/src/coreai_opt/quantization/spec/qparams_calculator.py index 28aa49a7..3664d2bf 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -7,6 +7,7 @@ import warnings from abc import abstractmethod +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -21,12 +22,18 @@ 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, + nonzero_scale as _nonzero_scale, ) +from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError -from .granularity import QuantizationGranularity +from .granularity import PerBlockGranularity, PerTensorGranularity, QuantizationGranularity from .qscheme import QuantizationScheme from .range_calculator import RangeCalculatorBase +if TYPE_CHECKING: + from .fake_quantize import FakeQuantizeImplBase + class QParamsCalculatorBase(_ClassRegistryMixin, nn.Module): """Abstract base for qparams calculators — common configuration and helpers. @@ -54,10 +61,17 @@ def __init__( range_calculator: RangeCalculatorBase, float_range: tuple[float | None, float | None], scale_dtype: torch.dtype | None = None, # primarily for FP4 + scale_fake_quantize: FakeQuantizeImplBase | None = None, + minval_fake_quantize: FakeQuantizeImplBase | None = None, **kwargs, ): super().__init__() self.scale_dtype = scale_dtype + # Quantize the scale and minval as the spec's ``qparam_qspec`` describes. The + # factory sets them only when ``qparam_qspec`` is set, the minval one only under + # the MINVAL formulation. + self.scale_fake_quantize = scale_fake_quantize + self.minval_fake_quantize = minval_fake_quantize self.dtype = dtype self.qscheme = qscheme self._granularity = granularity @@ -77,6 +91,8 @@ def __init__( # Ellipsis sentinel means unresolved; resolved on first forward() call. self._resolved_axis: int | None = ... # type: ignore[assignment] + self._resolve_nested_granularity() + @property def granularity(self) -> QuantizationGranularity: """Getter for granularity.""" @@ -100,6 +116,32 @@ def granularity(self, granularity: QuantizationGranularity) -> None: ) self._granularity = granularity self.range_calculator.granularity = granularity + self._resolve_nested_granularity() + + def _resolve_nested_granularity(self) -> None: + """Fit the ``qparam_qspec`` granularities to this calculator's grid of qparams. + + Weight axis defaults only reach the fake quantizers inserted in the model, so a nested + single-axis per-block granularity left at ``axis=None`` follows this calculator's block + axis whenever it is set. A per-tensor calculator, such as a shared observer that graph + mode forces to per-tensor, has a single qparam, so the nested fake quantizers quantize + it per tensor. + """ + granularity = self.granularity + for fake_quantize in (self.scale_fake_quantize, self.minval_fake_quantize): + if fake_quantize is None: + continue + nested = fake_quantize.granularity + if isinstance(granularity, PerTensorGranularity): + fake_quantize.granularity = PerTensorGranularity() + elif ( + isinstance(granularity, PerBlockGranularity) + and granularity.axis is not None + and isinstance(nested, PerBlockGranularity) + and isinstance(nested.block_size, int) + and nested.axis is None + ): + fake_quantize.granularity = nested.model_copy(update={"axis": granularity.axis}) def _resolve_axis(self, tensor_ndim: int) -> None: """Resolve axis to non-negative on first call, delegating to granularity. @@ -209,11 +251,13 @@ 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 - zero point. Supports integer dtypes and FP8 dtypes. + zero point. Supports integer dtypes, FP8 dtypes, and FP4 dtypes with ``qparam_qspec``. 2. **e8m0 scales** (``scale_dtype=torch.float8_e8m0fnu``): Scales are constrained to powers of 2 following the OCP Microscaling (MX) @@ -221,6 +265,17 @@ 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. + + With ``scale_fake_quantize`` / ``minval_fake_quantize`` set (the spec's + ``qparam_qspec``), the default-path scale and minval are then quantized, and + the decoded values are returned, so elements are quantized against the qparams + a low-precision qparam format reconstructs. + + In both cases, an asymmetric zero point is derived from the rounded or decoded scale. """ # e8m0 path: power-of-2 scales @@ -228,7 +283,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 - # Default path: torchao handles integer and FP8 dtypes + rounds_scale = self.scale_dtype in _ROUNDED_SCALE_DTYPES + adjusts_scale = rounds_scale or self.scale_fake_quantize is not None + if adjusts_scale: + # Start from an fp32 scale, so a half-precision input is not rounded twice. + min_val = min_val.to(torch.float32) + max_val = max_val.to(torch.float32) + + # Default path: torchao handles integer and FP8 dtypes, and FP4 with qparam_qspec scale, zero_point = choose_qparams_affine_with_min_max( min_val=min_val, max_val=max_val, @@ -241,6 +303,22 @@ 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 self.scale_fake_quantize is not None: + scale = self._quantize_qparam(scale, self.scale_fake_quantize, "scale") + if adjusts_scale and self.qscheme == QuantizationScheme.ASYMMETRIC: + # The zero point above came from the unrounded scale; recompute it from the + # rounded or decoded scale so the two match. Where that scale is 0, every element + # dequantizes to 0 whatever the zero point, so divide by 1 only to keep it finite. + zero_point = torch.clamp( + self.quant_min - torch.round(min_val / _nonzero_scale(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 @@ -260,6 +338,10 @@ def _compute_scale_zero_point_minval( # ``_get_min_and_max_val``. minval = torch.min(min_val, torch.zeros_like(min_val)) + if self.minval_fake_quantize is not None: + # The nested spec sees non-negative values, since minval is never positive. + minval = -self._quantize_qparam(-minval, self.minval_fake_quantize, "minval") + # For FP dtypes, neither zero_point nor minval is used (symmetric # quantization with no offset). if self.dtype.is_floating_point: @@ -280,6 +362,42 @@ def _initialize_state( subclasses override this to resize their buffers (e.g. ``running_min``/``running_max``) to match the observed input.""" + @staticmethod + def _quantize_qparam( + values: torch.Tensor, fake_quantize: FakeQuantizeImplBase, qparam_name: str + ) -> torch.Tensor: + """Return per-block qparams after quantization by ``fake_quantize``, decoded. + + The nested qparams follow ``fake_quantize``'s observer state, like + ``FakeQuantizeImplBase.forward``: recomputed while it observes, frozen otherwise. + Its fake-quant toggle is ignored, since the qparams must always be quantized. + + Args: + values (torch.Tensor): Non-negative qparams in fp32, one entry per block. + fake_quantize (FakeQuantizeImplBase): Fake quantizer built from + ``qparam_qspec``. + qparam_name (str): Name of the qparam, for error messages. + + Returns: + torch.Tensor: The decoded qparams, in fp32 and with the shape of ``values``. + + Raises: + _BlockSizeMismatchError: If the blocks do not divide into the groups of + ``qparam_qspec``. + """ + if fake_quantize.observer_enabled[0] == 1: + try: + scale, zero_point, minval = fake_quantize.qparams_calculator(values) + except _BlockSizeMismatchError as e: + msg = ( + f"qparam_qspec cannot group the per-block {qparam_name} grid of shape " + f"{tuple(values.shape)}: {e}" + ) + raise _BlockSizeMismatchError(msg) from e + else: + scale, zero_point, minval = fake_quantize.qparams_calculator.get_qparams() + return fake_quantize._fused_fake_quant_dequant(values, scale, zero_point, minval) + def compute_qparams( self, tensor: torch.Tensor, diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index c721cf3f..40c6bd73 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, ) @@ -27,6 +28,7 @@ from .fake_quantize import FakeQuantizeImplBase from .granularity import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationGranularity, @@ -148,13 +150,16 @@ class QuantizationSpec(CompressionSpec): And for FP4/FP8 dtypes, zero-point is always set to 0 (FP supports only the symmetric qscheme). The scale formula depends on ``scale_dtype``: - - ``scale_dtype=None`` (FP8 only): ``scale = max_abs / fp_max``, where - ``fp_max`` is the largest representable value for the target FP dtype - (448.0 for FP8 E4M3, 57344.0 for FP8 E5M2). + - ``scale_dtype=None`` (FP8, or FP4 with ``qparam_qspec``): + ``scale = max_abs / fp_max``, where ``fp_max`` is the largest + representable value for the target FP dtype (6.0 for FP4 E2M1, 448.0 for + FP8 E4M3, 57344.0 for FP8 E5M2). ``qparam_qspec`` then quantizes this + scale. - ``scale_dtype=torch.float8_e8m0fnu`` (FP4 and FP8): 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: @@ -302,23 +307,67 @@ 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 choose_qparams_affine_with_min_max (integer and FP8 dtypes). - For FP4, None is resolved to torch.float8_e8m0fnu automatically. + For FP4, None is resolved to torch.float8_e8m0fnu automatically, + unless ``qparam_qspec`` is set. - torch.float8_e8m0fnu: Power-of-2 scales following OCP Microscaling (MX) - spec. Required for FP4 quantization, optional for FP8. + spec. Required for FP4 quantization without ``qparam_qspec``, optional + for FP8. + - torch.float16: The scale is rounded to FP16 before element codes are + chosen, whatever the model dtype, and an asymmetric zero point is derived + from the rounded scale. Inside ``qparam_qspec``, it rounds the + super-scales, as GGUF Q4_K stores ``d`` and ``dmin``. Constraints: - FP4 (float4_e2m1fn_x2): scale_dtype must be torch.float8_e8m0fnu or None - (defaults to e8m0) + (defaults to e8m0; stays None when ``qparam_qspec`` is set) - 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) + - scale_dtype must be None when ``qparam_qspec`` is set + - Export of torch.float16 to CoreAI or CoreML is not supported yet: + ``Quantizer.finalize`` raises ``NotImplementedError`` for those backends. + + Default: None + + qparam_qspec (coreai_opt.quantization.QuantizationSpec | None): Quantization + specification for this spec's quantization parameters. The per-block qparams + form a grid with one entry per block, and ``qparam_qspec`` quantizes that + grid like any other tensor: the scale always, and under + ``qformulation="minval"`` the minval too, each with its own qparams of the + nested spec. Elements are then quantized against the decoded qparams, as a + format that stores the qparams at low precision reconstructs them. All + ``QuantizationSpec`` fields apply to the nested spec, with the grid as the + quantized tensor: + + - ``granularity`` counts qparam entries, not elements: for example, + ``PerBlockGranularity(block_size=8)`` shares one set of nested qparams + across 8 consecutive blocks. A single-axis ``PerBlockGranularity`` with + ``axis=None`` follows this spec's block axis, and requires this spec to use + single-axis ``PerBlockGranularity``. A ``PerChannelGranularity`` needs an + explicit axis. + - ``scale_dtype`` of the nested spec is the dtype of its own scales, and its + own ``qparam_qspec`` quantizes those in turn. + - ``qparam_calculator_cls`` and observer state control when the nested + qparams update. For example, a dynamic activation spec can recompute block + scales every forward, while its nested spec's tensor scale is calibrated, + then frozen when observers are disabled. + + When this spec is an asymmetric integer spec with ``qformulation="zp"``, it + derives its zero point from the decoded scale. + + Constraints: + + - scale_dtype must be None when ``qparam_qspec`` is set + - Export to CoreAI or CoreML is not supported yet: ``Quantizer.finalize`` + raises ``NotImplementedError`` for those backends. Default: None @@ -346,6 +395,26 @@ class type: MinMaxRangeCalculator or custom registered class type ... range_calculator_cls="minmax", ... float_range=[-1.0, 1.0] ... ) + >>> + >>> # int4 blocks whose scales are int8 codes, with one scale per output channel + >>> spec = QuantizationSpec( + ... dtype=torch.int4, + ... granularity=PerBlockGranularity(axis=1, block_size=32), + ... qparam_qspec=QuantizationSpec( + ... dtype=torch.int8, + ... granularity=PerChannelGranularity(axis=0), + ... ), + ... ) + >>> + >>> # NVFP4: FP8 E4M3 block scales under one FP32 scale per tensor + >>> spec = QuantizationSpec( + ... dtype=torch.float4_e2m1fn_x2, + ... granularity=PerBlockGranularity(axis=1, block_size=16), + ... qparam_qspec=QuantizationSpec( + ... dtype=torch.float8_e4m3fn, + ... granularity=PerTensorGranularity(), + ... ), + ... ) Notes: - All fields have defaults and are optional @@ -371,6 +440,7 @@ class type: MinMaxRangeCalculator or custom registered class type range_calculator_cls: type[RangeCalculatorBase] = Field(default="minmax", validate_default=True) float_range: list[float | int | None] = Field(default_factory=lambda: [None, None]) scale_dtype: torch.dtype | None = None + qparam_qspec: QuantizationSpec | None = None # Private attribute for compression type _compression_type: CompressionType = PrivateAttr(default=CompressionType.QUANTIZATION) @@ -489,7 +559,11 @@ def _strip_computed_fields(cls, data: Any) -> Any: @model_validator(mode="before") @classmethod def resolve_scale_dtype(cls, data: Any) -> Any: - """Resolve scale_dtype: convert string to torch.dtype and default to e8m0 for FP4.""" + """Resolve scale_dtype: convert string to torch.dtype and default to e8m0 for FP4. + + FP4 keeps ``scale_dtype=None`` when ``qparam_qspec`` is set, since the nested + spec then defines how the scale is stored. + """ if isinstance(data, dict): dtype = data.get("dtype") scale_dtype = data.get("scale_dtype") @@ -498,7 +572,11 @@ def resolve_scale_dtype(cls, data: Any) -> Any: data["dtype"] = dtype if isinstance(scale_dtype, str): data["scale_dtype"] = cls._resolve_str_dtype(scale_dtype) - if _is_float4_dtype(dtype) and data.get("scale_dtype") is None: + if ( + _is_float4_dtype(dtype) + and data.get("scale_dtype") is None + and data.get("qparam_qspec") is None + ): data["scale_dtype"] = torch.float8_e8m0fnu return data @@ -538,21 +616,33 @@ 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). + by resolve_scale_dtype (before validator), unless ``qparam_qspec`` is set. """ 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), or " + f"quantized by a nested spec (qparam_qspec, e.g. FP8 E4M3 scales as in NVFP4)." + ) + 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( @@ -578,6 +668,58 @@ def validate_qscheme_for_dynamic_quant(self) -> QuantizationSpec: raise ValueError(error_msg) return self + @model_validator(mode="after") + def validate_qparam_qspec(self) -> QuantizationSpec: + """Validate ``qparam_qspec`` against this spec. + + Rules: + - ``scale_dtype`` must be None, since ``qparam_qspec`` defines how the scale + is stored. + - The nested granularity is applied to the grid of per-block qparams, so an + axis it needs must be known: a single-axis ``PerBlockGranularity`` left at + ``axis=None`` follows this spec's block axis, which requires this spec to + use single-axis ``PerBlockGranularity``; a ``PerChannelGranularity`` needs + an explicit axis. + """ + nested = self.qparam_qspec + if nested is None: + return self + + if self.scale_dtype is not None: + msg = ( + f"scale_dtype must be None when qparam_qspec is set, got {self.scale_dtype}; " + f"qparam_qspec defines how the scale is stored." + ) + raise ValueError(msg) + + nested_granularity = nested.granularity + if ( + isinstance(nested_granularity, PerChannelGranularity) + and nested_granularity.axis is None + ): + msg = ( + "qparam_qspec.granularity=PerChannelGranularity needs an explicit axis of the " + "per-block qparam grid." + ) + raise ValueError(msg) + if ( + isinstance(nested_granularity, PerBlockGranularity) + and isinstance(nested_granularity.block_size, int) + and nested_granularity.axis is None + and not ( + isinstance(self.granularity, PerBlockGranularity) + and isinstance(self.granularity.block_size, int) + ) + ): + msg = ( + "qparam_qspec.granularity=PerBlockGranularity with axis=None follows this " + "spec's block axis, so this spec needs single-axis PerBlockGranularity (an int " + "block_size). Otherwise, set the qparam_qspec axis explicitly." + ) + raise ValueError(msg) + + return self + def get_extra_args(self) -> dict[str, Any]: """ Automatically detect and return fields beyond base QuantizationSpec. diff --git a/tests/inspection/test_bits_per_weight.py b/tests/inspection/test_bits_per_weight.py index 98656c9c..5f964923 100644 --- a/tests/inspection/test_bits_per_weight.py +++ b/tests/inspection/test_bits_per_weight.py @@ -36,6 +36,12 @@ ) from coreai_opt.quantization.spec.qformulation import QuantizationFormulation from tests.models.simple import LinearBatchNormModel, SharedParamsModel, SimpleLinearModel +from tests.test_utils.general import ( + minval_fp16_super_scales_spec, + minval_spec, + nvfp4_spec, + twice_quantized_spec, +) # SimpleLinearModel and LinearBatchNormModel are both Linear(64, 128) -> Linear(128, 64), # so their weights are 128 x 64 and 64 x 128. @@ -53,6 +59,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. @@ -463,6 +470,84 @@ 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 _prepare_eager_weights(model: nn.Module, weight_spec: QuantizationSpec) -> nn.Module: + """Prepare ``model`` in eager mode with only its weights quantized, per ``weight_spec``.""" + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_spec}, op_input_spec=None + ), + execution_mode="eager", + ) + return Quantizer(model, config).prepare(example_inputs=(torch.rand(4, 256),)) + + +@pytest.mark.parametrize( + "spec, expected_bpw, tensor_overhead_bits", + [ + # One fp32 tensor scale per weight tensor on top of 4 + 8/16 bits per weight. + pytest.param(nvfp4_spec(axis=None), 4.5, 2 * _FP32_BITS, id="nvfp4"), + # 4 + (8 + 8) / 32 + (32 + 32) / 256 = 4.75 bits per weight: uint8 scale and minval + # codes per block, under fp32 super-scales per 8 blocks, with no tensor overhead. + pytest.param(minval_spec(axis=None), 4.75, 0, id="minval"), + # The same with FP16 super-scales, as GGUF Q4_K stores its d and dmin: + # 4 + (8 + 8) / 32 + (16 + 16) / 256. + pytest.param( + minval_fp16_super_scales_spec(axis=None), + 4.625, + 0, + id="minval_fp16_super_scales", + ), + ], +) +def test_quantized_qparams_cost( + spec: QuantizationSpec, expected_bpw: float, tensor_overhead_bits: int +): + """Quantized qparams cost their codes plus one super-scale per group of blocks.""" + model = nn.Sequential(nn.Linear(256, 512, bias=False), nn.Linear(512, 16, bias=False)) + prepared = _prepare_eager_weights(model, spec) + + result = bits_per_weight(prepared) + + num_weights = 256 * 512 + 512 * 16 + assert result.total_weights == num_weights + assert result.total_bits == int(expected_bpw * num_weights) + tensor_overhead_bits + + +@pytest.mark.parametrize( + "nested_float_range", + [pytest.param([None, None], id="unpinned"), pytest.param([0.0, None], id="min_pinned_at_0")], +) +def test_nested_qparam_qspec_cost(nested_float_range: list[float | None]): + """Quantized qparams whose own scales are quantized in turn cost each level's codes.""" + model = nn.Sequential(nn.Linear(256, 32, bias=False)) + prepared = _prepare_eager_weights(model, twice_quantized_spec(float_range=nested_float_range)) + + result = bits_per_weight(prepared) + + num_weights = 256 * 32 + num_blocks = num_weights // 32 # one uint8 scale code each + num_nested_scales = num_blocks // 4 # one fp8 code each; the uint8 zero point is implied + assert result.total_bits == ( + num_weights * 4 + num_blocks * 8 + num_nested_scales * 8 + _FP32_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..4dafd045 100644 --- a/tests/palettization/test_palettization_spec.py +++ b/tests/palettization/test_palettization_spec.py @@ -267,6 +267,19 @@ def test_lut_qspec_validation(): ) ) + # Invalid: quantized LUT qparams + with pytest.raises(ValidationError, match="lut_qspec.qparam_qspec is not supported"): + PalettizationSpec( + lut_qspec=QuantizationSpec( + dtype=torch.int8, + qparam_qspec=QuantizationSpec(dtype=torch.float8_e4m3fn), + ) + ) + + # 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_qparam_qspec.py b/tests/quantization/test_qparam_qspec.py new file mode 100644 index 00000000..b8479e63 --- /dev/null +++ b/tests/quantization/test_qparam_qspec.py @@ -0,0 +1,679 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for quantized qparams (``QuantizationSpec.qparam_qspec``), e.g. NVFP4. + +The references below re-implement each format from its definition: quantize the per-block +qparams first, then choose the element codes against the decoded qparams. +""" + +import logging +from typing import Any + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchao.quantization.pt2e import disable_observer, enable_observer + +from coreai_opt import ExportBackend +from coreai_opt._utils.torch_utils import nonzero_scale +from coreai_opt.config.spec import CompressionTargetTensor +from coreai_opt.quantization import ( + ModuleQuantizerConfig, + QuantizationSpec, + Quantizer, + QuantizerConfig, +) +from coreai_opt.quantization.config import ExecutionMode, QATSchedule +from coreai_opt.quantization.spec import ( + PerBlockGranularity, + PerChannelGranularity, + PerTensorGranularity, + QuantizationFormulation, + QuantizationScheme, + fp4_forward, +) +from coreai_opt.quantization.spec.factory import QuantizationComponentFactory +from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase +from coreai_opt.quantization.spec.qparams_calculator import MovingAverageQParamsCalculator +from tests.test_utils.general import ( + minval_fp16_super_scales_spec, + minval_spec, + nvfp4_spec, + twice_quantized_spec, +) + + +def _weight_fake_quantizer(spec: QuantizationSpec) -> FakeQuantizeImplBase: + """Return the weight fake quantizer that ``spec`` builds.""" + return QuantizationComponentFactory.create_fake_quantizer(spec, CompressionTargetTensor.WEIGHT) + + +def _nvfp4_reference(weight: torch.Tensor) -> torch.Tensor: + """Fake-quantize a 2-D weight as NVFP4, with blocks of 16 along the last dim.""" + blocks = weight.float().reshape(weight.shape[0], -1, 16) + ideal_scale = blocks.abs().amax(dim=-1, keepdim=True) / 6.0 + tensor_scale = ideal_scale.amax() / 448.0 + block_scale = (ideal_scale / tensor_scale).clamp(max=448.0).to(torch.float8_e4m3fn).float() + scale = block_scale * tensor_scale + elements = fp4_forward((blocks / nonzero_scale(scale)).clamp(-6.0, 6.0)) + return (elements * scale).reshape(weight.shape) + + +def _decode_uint8_qparams(values: torch.Tensor, super_scale_dtype: torch.dtype) -> torch.Tensor: + """Round trip non-negative qparams, grouped 8 along the last dim, through uint8 codes. + + The super-scales are rounded to ``super_scale_dtype`` first, as a float16 nested + ``scale_dtype`` does. + """ + super_scale = (values.amax(dim=-1, keepdim=True) / 255.0).to(super_scale_dtype).float() + return (values / super_scale).round().clamp(0, 255) * super_scale + + +def _minval_reference( + weight: torch.Tensor, super_scale_dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Fake-quantize a 2-D weight as ``minval_spec``, with groups of 8x32 along the last dim. + + The super-scales are rounded to ``super_scale_dtype``, as a float16 nested ``scale_dtype`` does. + """ + blocks = weight.float().reshape(weight.shape[0], -1, 8, 32) + low = blocks.amin(dim=-1).clamp(max=0) + high = blocks.amax(dim=-1).clamp(min=0) + scale = _decode_uint8_qparams((high - low) / 15.0, super_scale_dtype).unsqueeze(-1) + minval = -_decode_uint8_qparams(-low, super_scale_dtype).unsqueeze(-1) + codes = ((blocks - minval) / nonzero_scale(scale)).round().clamp(0, 15) + return (codes * scale + minval).reshape(weight.shape) + + +def _minval_fp16_super_scales_reference(weight: torch.Tensor) -> torch.Tensor: + return _minval_reference(weight, super_scale_dtype=torch.float16) + + +_REFERENCE_CASES = [ + (nvfp4_spec, _nvfp4_reference), + (minval_spec, _minval_reference), + (minval_fp16_super_scales_spec, _minval_fp16_super_scales_reference), +] +_SPEC_FACTORIES = [spec_factory for spec_factory, _ in _REFERENCE_CASES] +_EXECUTION_MODES = pytest.mark.parametrize( + "execution_mode", [ExecutionMode.EAGER, ExecutionMode.GRAPH] +) + + +def _nvfp4_activation_spec(tensor_scale_calculator: str = "default") -> QuantizationSpec: + """NVFP4 activations: block scales recomputed every forward, under a tensor scale.""" + return nvfp4_spec( + axis=-1, + qparam_calculator_cls="dynamic", + qparam_qspec=QuantizationSpec( + dtype=torch.float8_e4m3fn, + granularity=PerTensorGranularity(), + qparam_calculator_cls=tensor_scale_calculator, + ), + ) + + +def _mlp() -> nn.Module: + return nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 16)).eval() + + +def _config( + weight_spec: QuantizationSpec, execution_mode: ExecutionMode, **module_config_kwargs: Any +) -> QuantizerConfig: + """Return a config for ``weight_spec``; extra kwargs go to ``ModuleQuantizerConfig``.""" + module_config_kwargs = {"op_input_spec": None, "op_output_spec": None} | module_config_kwargs + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_spec}, **module_config_kwargs + ), + execution_mode=execution_mode, + ) + + +def _worked_example_spec(qformulation: QuantizationFormulation) -> QuantizationSpec: + """uint4 blocks of 4 whose qparams are uint4 codes under one super-scale per 2 blocks. + + A 1 x 8 weight is then 2 blocks forming 1 group, small enough to work out by hand, and the + coarse nested codes make the effect of quantizing the qparams visible. + """ + return QuantizationSpec( + dtype=torch.uint4, + qscheme=QuantizationScheme.ASYMMETRIC, + qformulation=qformulation, + granularity=PerBlockGranularity(axis=1, block_size=4), + qparam_qspec=QuantizationSpec( + dtype=torch.uint4, + qscheme=QuantizationScheme.ASYMMETRIC, + granularity=PerBlockGranularity(block_size=2), + ), + ) + + +def _assert_values(actual: torch.Tensor, expected: list[float]) -> None: + # Stricter than the fp32 defaults; still allows fp32 rounding of the decimal literals. + torch.testing.assert_close(actual.flatten(), torch.tensor(expected), rtol=0, atol=1e-6) + + +def _minval_with_block_1( + block_1: torch.Tensor | float, +) -> tuple[FakeQuantizeImplBase, torch.Tensor]: + """Fake-quantize one group of 8 blocks with ``minval_spec``, varying only block 1. + + Block 0 spans [-1, 2] and blocks 2 to 7 are zero. + """ + weight = torch.zeros(1, 256) + weight[0, :32] = torch.linspace(-1.0, 2.0, 32) + weight[0, 32:64] = block_1 + fake_quantize = _weight_fake_quantizer(minval_spec()) + return fake_quantize, fake_quantize(weight) + + +# --------------------------------------------------------------------------- +# NVFP4 numerics +# --------------------------------------------------------------------------- + + +def test_nvfp4_chooses_elements_against_the_quantized_block_scale(): + """Elements are chosen against the E4M3-quantized block scale. + + With a tensor scale of 1, a block maximum of 7 rounds the block scale 7/6 to the E4M3 value + 1.125, so 1.43 / 1.125 = 1.27 rounds up to the E2M1 value 1.5. Against the unrounded scale, + 1.43 / 1.167 = 1.23 would round down to 1.0 instead. + """ + weight = torch.zeros(2, 16) + weight[0, 0] = 2688.0 # 6 * 448, so the tensor scale is 1 + weight[1, 0] = 7.0 + weight[1, 1] = 1.43 + fake_quantize = _weight_fake_quantizer(nvfp4_spec()) + + out = fake_quantize(weight) + + calculator = fake_quantize.qparams_calculator + assert calculator.scale_fake_quantize.qparams_calculator.scale.flatten().tolist() == [1.0] + assert calculator.scale.flatten().tolist() == [448.0, 1.125] + assert out[0, 0].item() == 2688.0 + # 7 / 1.125 = 6.2 saturates at 6.0. + assert out[1, :3].tolist() == [6.75, 1.6875, 0.0] + + +def test_nvfp4_block_scale_underflow_decodes_to_zero(): + """A block far below the tensor maximum decodes to zeros rather than NaN. + + Its E4M3 scale code underflows to 0. + """ + weight = torch.randn(2, 32) + weight[0, :16] *= 1e6 + weight[1, 16:] *= 1e-6 + fake_quantize = _weight_fake_quantizer(nvfp4_spec()) + + out = fake_quantize(weight) + + assert fake_quantize.qparams_calculator.scale[1, 1].item() == 0.0 + assert torch.equal(out[1, 16:], torch.zeros(16)) + assert torch.isfinite(out).all() + + +def test_nvfp4_all_zero_weight_stays_zero(): + out = _weight_fake_quantizer(nvfp4_spec())(torch.zeros(4, 64)) + assert torch.equal(out, torch.zeros(4, 64)) + + +# --------------------------------------------------------------------------- +# Quantized minval numerics +# --------------------------------------------------------------------------- + + +def test_minval_elements_use_the_decoded_scale_and_minval(): + """Hand-computed values: elements are chosen against the decoded scale and minval. + + Block 0 is [-1.5, 0, 1.5, 3]: scale 4.5 / 15 = 0.3, minval -1.5. Block 1 is + [-0.28, 0.1, 0.3, 0.555]: scale 0.835 / 15 = 0.0557, minval -0.28. The scales share one + super-scale, 0.3 / 15 = 0.02, so block 1's scale becomes code 3 (2.78 rounded), 0.06. The + minvals get their own super-scale, 1.5 / 15 = 0.1, so block 1's minval becomes code 3 (2.8 + rounded), -0.3. Block 1 then decodes to -0.3 + 0.06 * [0, 7, 10, 14]. Without the nested + spec it would decode to [-0.28, 0.1097, 0.2767, 0.555]. + """ + fake_quantize = _weight_fake_quantizer(_worked_example_spec(QuantizationFormulation.MINVAL)) + weight = torch.tensor([[-1.5, 0.0, 1.5, 3.0, -0.28, 0.1, 0.3, 0.555]]) + + out = fake_quantize(weight) + + calculator = fake_quantize.qparams_calculator + _assert_values(calculator.scale_fake_quantize.qparams_calculator.scale, [0.02]) + _assert_values(calculator.minval_fake_quantize.qparams_calculator.scale, [0.1]) + _assert_values(calculator.scale, [0.3, 0.06]) + _assert_values(calculator.minval, [-1.5, -0.3]) + _assert_values(out, [-1.5, 0.0, 1.5, 3.0, -0.3, 0.12, 0.3, 0.54]) + + +def test_fp16_super_scales_are_fp16_values(): + """A float16 nested ``scale_dtype`` rounds the super-scales to FP16. + + The output then differs from FP32 super-scales; ``test_matches_reference`` checks its values. + """ + weight = torch.randn(8, 512) * 0.02 + fake_quantize = _weight_fake_quantizer(minval_fp16_super_scales_spec()) + + out = fake_quantize(weight) + + calculator = fake_quantize.qparams_calculator + for qparam_fq in (calculator.scale_fake_quantize, calculator.minval_fake_quantize): + super_scale = qparam_fq.qparams_calculator.scale + assert torch.equal(super_scale, super_scale.half().float()) + assert not torch.equal(out, _minval_reference(weight)) + + +def test_minval_qparams_decompose_into_codes_and_super_scales(): + """Scale and minval are uint8 codes times separate super-scales, one of each per 8 blocks. + + The nested specs see non-negative values, so their zero points are 0. + """ + weight = torch.randn(4, 512) * 0.02 + fake_quantize = _weight_fake_quantizer(minval_spec()) + fake_quantize(weight) + calculator = fake_quantize.qparams_calculator + + for name, decoded, qparam_fq in ( + ("scale", calculator.scale, calculator.scale_fake_quantize), + ("minval", -calculator.minval, calculator.minval_fake_quantize), + ): + super_scale = qparam_fq.qparams_calculator.scale + assert super_scale.shape == (4, 2), name + assert (qparam_fq.qparams_calculator.zero_point == 0).all(), name + codes = decoded.reshape(4, 2, 8) / super_scale.unsqueeze(-1) + torch.testing.assert_close(codes, codes.round()) + assert codes.round().min() >= 0 and codes.round().max() <= 255, name + assert calculator.scale_fake_quantize is not calculator.minval_fake_quantize + + +def test_minval_constant_negative_block_decodes_to_its_minval(): + """A constant negative block decodes to the nearest minval its super-scale can express.""" + fake_quantize, out = _minval_with_block_1(-0.65) + + min_super_scale = fake_quantize.qparams_calculator.minval_fake_quantize.qparams_calculator.scale + assert out[0, 32:64].unique().numel() == 1 + assert (out[0, 32] + 0.65).abs() <= min_super_scale.item() / 2 + + +def test_minval_scale_code_zero_decodes_to_minval(): + """A block with scale code 0 decodes to its minval rather than NaN. + + Its range, which always includes zero, is below 1/510 of the largest range in its group of + 8 blocks. + """ + fake_quantize, out = _minval_with_block_1(torch.linspace(0.0, 1e-4, 32)) + + calculator = fake_quantize.qparams_calculator + assert calculator.scale[0, 1].item() == 0.0 + assert torch.equal(out[0, 32:64], calculator.minval[0, 1].expand(32)) + + +def test_minval_all_positive_block_has_zero_minval(): + fake_quantize, out = _minval_with_block_1(torch.linspace(0.1, 0.9, 32)) + + assert fake_quantize.qparams_calculator.minval[0, 1].item() == 0.0 + assert (out[0, 32:64] >= 0).all() + + +# --------------------------------------------------------------------------- +# Shared behavior +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("spec_factory", "reference"), _REFERENCE_CASES) +def test_matches_reference(spec_factory, reference): + weight = torch.randn(8, 512) * 0.02 + out = _weight_fake_quantizer(spec_factory())(weight) + assert torch.equal(out, reference(weight)) + + +@pytest.mark.parametrize("spec_factory", _SPEC_FACTORIES) +def test_half_precision_weight_matches_fp32_qparams(spec_factory): + """The qparams are computed from an fp32 scale, so a bf16 weight is not rounded twice.""" + weight = (torch.randn(4, 512) * 0.02).bfloat16() + bf16_fake_quantize = _weight_fake_quantizer(spec_factory()) + fp32_fake_quantize = _weight_fake_quantizer(spec_factory()) + + bf16_fake_quantize(weight) + fp32_fake_quantize(weight.float()) + + torch.testing.assert_close( + bf16_fake_quantize.calculate_qparams(), + fp32_fake_quantize.calculate_qparams(), + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize("spec_factory", _SPEC_FACTORIES) +def test_gradient_flows_through_fake_quantize(spec_factory): + weight = (torch.randn(4, 512) * 0.02).requires_grad_() + _weight_fake_quantizer(spec_factory())(weight).sum().backward() + assert weight.grad is not None + assert torch.isfinite(weight.grad).all() + assert weight.grad.abs().sum() > 0 + + +def test_nested_block_axis_follows_the_weight_block_axis(): + """Nested per-block granularities take the block axis once weight axis defaults set it.""" + fake_quantize = _weight_fake_quantizer(minval_spec(axis=None)) + calculator = fake_quantize.qparams_calculator + + fake_quantize.granularity = fake_quantize.granularity.model_copy(update={"axis": 1}) + + assert calculator.scale_fake_quantize.granularity.axis == 1 + assert calculator.minval_fake_quantize.granularity.axis == 1 + + +def test_explicit_nested_axis_groups_along_that_axis(): + """A nested block axis other than the weight's groups the qparam grid along it.""" + spec = QuantizationSpec( + dtype=torch.int4, + granularity=PerBlockGranularity(axis=1, block_size=32), + qparam_qspec=QuantizationSpec( + dtype=torch.uint8, + qscheme=QuantizationScheme.ASYMMETRIC, + granularity=PerBlockGranularity(axis=0, block_size=2), + ), + ) + fake_quantize = _weight_fake_quantizer(spec) + + fake_quantize(torch.randn(8, 128)) + + # The 8 x 4 blocks share nested qparams in pairs of rows. + nested_calculator = fake_quantize.qparams_calculator.scale_fake_quantize.qparams_calculator + assert nested_calculator.scale.shape == (4, 4) + + +def test_ungroupable_blocks_disable_the_fake_quantizer(caplog): + """288 input channels form 9 blocks of 32, which do not divide into groups of 8.""" + weight = torch.randn(4, 288) + fake_quantize = _weight_fake_quantizer(minval_spec()) + + with caplog.at_level(logging.WARNING): + out = fake_quantize(weight) + + assert torch.equal(out, weight) + assert fake_quantize.is_disabled() + assert "qparam_qspec cannot group the per-block scale grid of shape (4, 9)" in caplog.text + + +def test_per_tensor_outer_spec_quantizes_its_scale_per_tensor(): + """An outer spec switched to per-tensor quantizes its single scale and minval per tensor. + + Graph mode forces some shared observers to per-tensor, which leaves nothing for a nested + per-block granularity to group. + """ + fake_quantize = _weight_fake_quantizer(minval_spec()) + fake_quantize.granularity = PerTensorGranularity() + weight = torch.randn(4, 512) * 0.02 + + out = fake_quantize(weight) + + calculator = fake_quantize.qparams_calculator + assert not fake_quantize.is_disabled() + for qparam_fq in (calculator.scale_fake_quantize, calculator.minval_fake_quantize): + assert qparam_fq.granularity == PerTensorGranularity() + assert not torch.equal(out, weight) + + +@pytest.mark.parametrize( + ("spec_factory", "quantizes_minval"), [(nvfp4_spec, False), (minval_spec, True)] +) +def test_qparam_fake_quantizers_per_formulation(spec_factory, quantizes_minval): + """The scale is always quantized; minval only under the MINVAL formulation.""" + partial = QuantizationComponentFactory.create_fake_quantizer_partial( + spec_factory(), CompressionTargetTensor.WEIGHT + ) + first, second = partial(), partial() + assert first.qparams_calculator.scale_fake_quantize is not None + assert (first.qparams_calculator.minval_fake_quantize is not None) == quantizes_minval + # Each constructed fake quantizer gets its own qparam fake quantizers. + assert ( + first.qparams_calculator.scale_fake_quantize + is not second.qparams_calculator.scale_fake_quantize + ) + + +# --------------------------------------------------------------------------- +# General qparam quantization +# --------------------------------------------------------------------------- + + +def test_int8_scale_codes_per_channel_match_reference(): + """int4 blocks whose scales are int8 codes, with one tensor scale per output channel.""" + spec = QuantizationSpec( + dtype=torch.int4, + granularity=PerBlockGranularity(axis=1, block_size=32), + qparam_qspec=QuantizationSpec(dtype=torch.int8, granularity=PerChannelGranularity(axis=0)), + ) + weight = torch.randn(8, 256) + + out = _weight_fake_quantizer(spec)(weight) + + blocks = weight.reshape(8, -1, 32) + ideal_scale = blocks.abs().amax(dim=-1) / 7.5 + channel_scale = ideal_scale.amax(dim=-1, keepdim=True) / 127.5 + scale = ((ideal_scale / channel_scale).round().clamp(-128, 127) * channel_scale).unsqueeze(-1) + expected = (blocks / nonzero_scale(scale)).round().clamp(-8, 7) * scale + assert torch.equal(out, expected.reshape(weight.shape)) + + +def test_zero_point_is_rederived_from_the_decoded_scale(): + """Hand-computed values: the zero point comes from the decoded scale, not the raw one. + + Block 0 is [-1.5, 0, 1.5, 3]: scale 0.3, zero point 5. Block 1 is [-0.3, 0, 0.3, 0.72]: + scale 1.02 / 15 = 0.068, which the shared super-scale 0.02 rounds to code 3, so 0.06. The zero + point of the raw scale, round(0.3 / 0.068) = 4, would clip -0.3 to -0.24. Re-derived from + 0.06 it is 5, so -0.3 stays exact and 0.72 clips to the top code, (15 - 5) * 0.06 = 0.6. + """ + fake_quantize = _weight_fake_quantizer(_worked_example_spec(QuantizationFormulation.ZP)) + weight = torch.tensor([[-1.5, 0.0, 1.5, 3.0, -0.3, 0.0, 0.3, 0.72]]) + + out = fake_quantize(weight) + + calculator = fake_quantize.qparams_calculator + _assert_values(calculator.scale_fake_quantize.qparams_calculator.scale, [0.02]) + _assert_values(calculator.scale, [0.3, 0.06]) + assert calculator.zero_point.flatten().tolist() == [5, 5] + _assert_values(out, [-1.5, 0.0, 1.5, 3.0, -0.3, 0.0, 0.3, 0.6]) + + +def test_qparam_qspec_nests_recursively(): + """A nested spec's own qparam_qspec quantizes the nested scales in turn.""" + fake_quantize = _weight_fake_quantizer(twice_quantized_spec()) + + fake_quantize(torch.randn(4, 256)) + + nested_calculator = fake_quantize.qparams_calculator.scale_fake_quantize.qparams_calculator + innermost_scale = nested_calculator.scale_fake_quantize.qparams_calculator.scale + # One nested scale per 4 of the 4 x 8 blocks, each an E4M3 code times one tensor scale. + assert nested_calculator.scale.shape == (4, 2) + codes = nested_calculator.scale / innermost_scale + torch.testing.assert_close(codes, codes.to(torch.float8_e4m3fn).float()) + + +def test_activation_qparam_fake_quantizer_targets_activations(): + """The nested fake quantizer shares the target, so "default" resolves to a moving average.""" + fake_quantize = QuantizationComponentFactory.create_fake_quantizer( + _nvfp4_activation_spec(), CompressionTargetTensor.ACTIVATION + ) + qparam_fq = fake_quantize.qparams_calculator.scale_fake_quantize + assert qparam_fq.quantization_target == CompressionTargetTensor.ACTIVATION + assert isinstance(qparam_fq.qparams_calculator, MovingAverageQParamsCalculator) + + +def test_nested_scale_follows_the_observer_while_block_scales_stay_dynamic(): + """The nested tensor scale freezes with the observer; dynamic block scales keep updating.""" + fake_quantize = QuantizationComponentFactory.create_fake_quantizer( + _nvfp4_activation_spec(tensor_scale_calculator="global_minmax"), + CompressionTargetTensor.ACTIVATION, + ) + tensor_scale = fake_quantize.qparams_calculator.scale_fake_quantize.qparams_calculator + inputs = torch.randn(2, 64) + fake_quantize(inputs) + calibrated = tensor_scale.scale.clone() + + fake_quantize.apply(disable_observer) + frozen_out = fake_quantize(inputs * 0.5) + assert torch.equal(tensor_scale.scale, calibrated) + assert torch.isfinite(frozen_out).all() + + fake_quantize.apply(enable_observer) + fake_quantize(inputs * 10) + assert tensor_scale.scale.item() > calibrated.item() + + +# --------------------------------------------------------------------------- +# End to end +# --------------------------------------------------------------------------- + + +@_EXECUTION_MODES +def test_weight_and_activation_qparam_qspec_end_to_end(execution_mode): + """NVFP4 weights and activations: calibrate, then run the prepared model in torch.""" + model = _mlp() + inputs = torch.randn(4, 256) + config = _config( + nvfp4_spec(axis=None), execution_mode, op_input_spec={"*": _nvfp4_activation_spec()} + ) + quantizer = Quantizer(model, config) + prepared_model = quantizer.prepare((inputs,)) + with quantizer.calibration_mode(): + prepared_model(inputs) + + out = prepared_model(inputs) + + assert torch.isfinite(out).all() + torch.testing.assert_close(out, model(inputs), atol=0.3, rtol=0.3) + assert torch.equal(quantizer.finalize(backend=ExportBackend.Torch)(inputs), out) + + +@_EXECUTION_MODES +@pytest.mark.parametrize(("spec_factory", "reference"), _REFERENCE_CASES) +def test_prepared_model_matches_reference(spec_factory, reference, execution_mode): + model = _mlp() + inputs = torch.randn(4, 256) + first, second = model[0], model[2] + hidden = F.relu(F.linear(inputs, reference(first.weight.detach()), first.bias)) + expected = F.linear(hidden, reference(second.weight.detach()), second.bias) + + quantizer = Quantizer(model, _config(spec_factory(axis=None), execution_mode)) + prepared_model = quantizer.prepare((inputs,)) + out = prepared_model(inputs) + + torch.testing.assert_close(out, expected) + finalized_model = quantizer.finalize(backend=ExportBackend.Torch) + assert torch.equal(finalized_model(inputs), out) + + +@_EXECUTION_MODES +@pytest.mark.parametrize( + ("weight_spec", "backend"), + [ + pytest.param(nvfp4_spec(axis=None), ExportBackend.CoreAI, id="nvfp4-coreai"), + pytest.param(minval_spec(axis=None), ExportBackend.CoreML, id="minval-coreml"), + pytest.param( + QuantizationSpec( + dtype=torch.int4, + granularity=PerBlockGranularity(axis=1, block_size=32), + qparam_qspec=QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.ASYMMETRIC, + granularity=PerChannelGranularity(axis=0), + ), + ), + ExportBackend.CoreAI, + id="scale_codes_with_an_offset-coreai", + ), + pytest.param( + twice_quantized_spec(), ExportBackend.CoreML, id="scale_quantized_twice-coreml" + ), + pytest.param( + QuantizationSpec( + dtype=torch.int8, + granularity=PerBlockGranularity(axis=1, block_size=32), + scale_dtype=torch.float16, + ), + ExportBackend.CoreAI, + id="float16_scale-coreai", + ), + pytest.param( + QuantizationSpec( + dtype=torch.int8, + granularity=PerBlockGranularity(axis=1, block_size=32), + scale_dtype=torch.float16, + ), + ExportBackend.CoreML, + id="float16_scale-coreml", + ), + ], +) +def test_export_is_not_supported_yet(weight_spec, backend, execution_mode): + """Any qparam_qspec or float16 scale_dtype blocks CoreAI and CoreML export for now.""" + inputs = torch.randn(4, 256) + quantizer = Quantizer(_mlp(), _config(weight_spec, execution_mode)) + quantizer.prepare((inputs,)) + + with pytest.raises(NotImplementedError, match="quantized qparams") as excinfo: + quantizer.finalize(backend=backend) + + # Only the weight fake quantizers are named, not the nested ones, which also quantize + # qparams when a spec is quantized twice. + assert "fake_quantize" not in str(excinfo.value).split("Affected FakeQuantize modules:")[1] + + +def test_qat_step_updates_weights(): + model = _mlp() + inputs = torch.randn(4, 256) + quantizer = Quantizer(model, _config(minval_spec(axis=None), ExecutionMode.EAGER)) + prepared_model = quantizer.prepare((inputs,)) + optimizer = torch.optim.SGD(prepared_model.parameters(), lr=0.1) + before = [p.detach().clone() for p in prepared_model.parameters()] + + with quantizer.training_mode(): + prepared_model(inputs).square().mean().backward() + optimizer.step() + + after = list(prepared_model.parameters()) + assert all(torch.isfinite(p).all() for p in after) + assert any(not torch.equal(b, a) for b, a in zip(before, after, strict=True)) + + +@_EXECUTION_MODES +def test_qat_schedule_freezes_nested_scales(execution_mode): + """A QAT schedule's disable_observer step also freezes the nested scales. + + The dynamic activation calculator quantizes its block scales on every forward, so the nested + tensor scale only stops updating if the schedule reaches its fake quantizer. + """ + inputs = torch.randn(4, 256) + config = _config( + nvfp4_spec(axis=None), + execution_mode, + op_input_spec={"*": _nvfp4_activation_spec(tensor_scale_calculator="global_minmax")}, + qat_schedule=QATSchedule(enable_observer=0, enable_fake_quant=0, disable_observer=1), + ) + quantizer = Quantizer(_mlp(), config) + prepared_model = quantizer.prepare((inputs,)) + nested_calculators = [ + mod.qparams_calculator.scale_fake_quantize.qparams_calculator + for mod in prepared_model.modules() + if isinstance(mod, FakeQuantizeImplBase) + and mod.quantization_target == CompressionTargetTensor.ACTIVATION + and mod.qparams_calculator.scale_fake_quantize is not None + ] + assert nested_calculators + + with quantizer.training_mode(): + prepared_model(inputs) + quantizer.step() + frozen = [calculator.scale.clone() for calculator in nested_calculators] + prepared_model(inputs * 10) + + for calculator, scale in zip(nested_calculators, frozen, strict=True): + assert torch.equal(calculator.scale, scale) diff --git a/tests/quantization/test_qparams_calculator.py b/tests/quantization/test_qparams_calculator.py index 2b5a9628..6659801c 100644 --- a/tests/quantization/test_qparams_calculator.py +++ b/tests/quantization/test_qparams_calculator.py @@ -1042,3 +1042,52 @@ 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]) +@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(8, 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 diff --git a/tests/quantization/test_qspec_reconcile.py b/tests/quantization/test_qspec_reconcile.py index 20b7cd1c..334f9763 100644 --- a/tests/quantization/test_qspec_reconcile.py +++ b/tests/quantization/test_qspec_reconcile.py @@ -22,13 +22,17 @@ ConcatPattern, _same_axis, ) +from coreai_opt.quantization._graph._provisional_qspec_generation import _FIELD_FROM_SPEC_ATTR from coreai_opt.quantization._graph._qspec_constraints import ( InheritFields, ShareFields, ShareObserverInstance, _reconcile_field, ) -from coreai_opt.quantization._graph._qspec_resolution import _build_concrete_spec +from coreai_opt.quantization._graph._qspec_resolution import ( + _SPEC_KWARG_FROM_FIELD, + _build_concrete_spec, +) from coreai_opt.quantization._graph._qspec_types import ( OP_INTRINSIC_PRIORITY, FieldName, @@ -40,10 +44,12 @@ SlotKind, ) from coreai_opt.quantization.spec import ( + PerBlockGranularity, PerChannelGranularity, PerTensorGranularity, QuantizationFormulation, QuantizationScheme, + QuantizationSpec, ) from coreai_opt.quantization.spec.fake_quantize import _DefaultFakeQuantizeImpl from coreai_opt.quantization.spec.qparams_calculator import ( @@ -52,6 +58,7 @@ _DefaultQParamsCalculator, ) from coreai_opt.quantization.spec.range_calculator import MinMaxRangeCalculator +from tests.test_utils.general import minval_qparam_qspec # --------------------------------------------------------------------------- # Test helpers. @@ -652,6 +659,7 @@ def _fields(**overrides): FieldName.QPARAM_CALCULATOR_CLS: _fv(MovingAverageQParamsCalculator), FieldName.RANGE_CALCULATOR_CLS: _fv(MinMaxRangeCalculator), FieldName.SCALE_DTYPE: _fv(None), + FieldName.QPARAM_QSPEC: _fv(None), FieldName.QUANTIZATION_TARGET: _fv(CompressionTargetTensor.ACTIVATION), } # Keyword keys arrive as strings; map them onto FieldName so they @@ -699,6 +707,26 @@ def test_target_selects_the_default_calculator(self) -> None: assert isinstance(self._calculator(weight), StaticQParamsCalculator) assert isinstance(self._calculator(activation), MovingAverageQParamsCalculator) + def test_reconciled_qparam_qspec_reaches_the_calculator(self) -> None: + """A reconciled ``qparam_qspec`` builds the nested scale and minval fake quantizers.""" + spec = _build_concrete_spec( + self._fields( + DTYPE=_fv(torch.uint4), + QSCHEME=_fv(QuantizationScheme.ASYMMETRIC), + QFORMULATION=_fv(QuantizationFormulation.MINVAL), + GRANULARITY=_fv(PerBlockGranularity(axis=1, block_size=32)), + QPARAM_CALCULATOR_CLS=_fv(StaticQParamsCalculator), + QPARAM_QSPEC=_fv(minval_qparam_qspec()), + QUANTIZATION_TARGET=_fv(CompressionTargetTensor.WEIGHT), + ) + ) + calculator = self._calculator(spec) + for qparam_fq in (calculator.scale_fake_quantize, calculator.minval_fake_quantize): + assert qparam_fq.dtype == torch.uint8 + # The nested block axis follows the weight's from construction on. + assert qparam_fq.granularity == PerBlockGranularity(axis=1, block_size=8) + assert calculator.scale_fake_quantize is not calculator.minval_fake_quantize + def test_empty_field_map_is_not_observed(self) -> None: assert _build_concrete_spec(ProvisionalQSpec(fields={})) is None @@ -710,6 +738,16 @@ def test_partial_field_map_raises(self) -> None: _build_concrete_spec(ProvisionalQSpec(fields={FieldName.DTYPE: _fv(torch.int8)})) +def test_field_maps_cover_every_quantization_spec_field() -> None: + """Every QuantizationSpec field appears in the graph-mode field maps. + + Graph mode rebuilds op-level specs from these maps, so a missing field is silently dropped. + """ + spec_fields = set(QuantizationSpec.model_fields) + assert set(_FIELD_FROM_SPEC_ATTR.values()) == spec_fields + assert set(_SPEC_KWARG_FROM_FIELD.values()) == spec_fields + + # --------------------------------------------------------------------------- # Explicitly-declined slots. # --------------------------------------------------------------------------- diff --git a/tests/quantization/test_quantization_config.py b/tests/quantization/test_quantization_config.py index 76be2371..09041727 100644 --- a/tests/quantization/test_quantization_config.py +++ b/tests/quantization/test_quantization_config.py @@ -29,6 +29,7 @@ from tests.quantization.test_quantization_spec import ( expanded_dtype_allowlist, # noqa: F401 ) +from tests.test_utils.general import minval_spec GLOBAL_CONFIG = "global_config" MODULE_TYPE_CONFIGS = "module_type_configs" @@ -357,6 +358,38 @@ def test_from_yaml_digit_string_keys_e2e(simple_linear_model, simple_linear_mode assert output.shape == simple_linear_model_input.shape +def test_from_yaml_qparam_qspec(): + """A nested qparam_qspec written in YAML builds the same config as in Python.""" + config_content = """\ +quantization_config: + global_config: + op_state_spec: + weight: + dtype: uint4 + qscheme: asymmetric + qformulation: minval + granularity: + type: per_block + block_size: 32 + qparam_qspec: + dtype: uint8 + qscheme: asymmetric + granularity: + type: per_block + block_size: 8 + op_input_spec: null + op_output_spec: null +""" + config = QuantizerConfig.from_yaml(io.StringIO(config_content)) + + weight_spec = minval_spec(axis=None) + assert config == QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_input_spec=None, op_output_spec=None, op_state_spec={"weight": weight_spec} + ) + ) + + def test_from_yaml_null_config(tmp_path): """Test from_yaml with op-level and module-level configs""" config_content = """ diff --git a/tests/quantization/test_quantization_spec.py b/tests/quantization/test_quantization_spec.py index a4d845e5..38adf594 100644 --- a/tests/quantization/test_quantization_spec.py +++ b/tests/quantization/test_quantization_spec.py @@ -31,6 +31,7 @@ _DefaultQParamsCalculator, ) from coreai_opt.quantization.spec.range_calculator import MinMaxRangeCalculator +from tests.test_utils.general import minval_qparam_qspec, minval_spec, nvfp4_spec @pytest.fixture @@ -876,6 +877,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): @@ -892,9 +896,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, @@ -911,6 +916,111 @@ def test_invalid_scale_dtype(dtype, scale_dtype): ) +def test_qparam_qspec_keeps_fp4_scale_dtype_unset(): + """FP4 defaults to e8m0 scales only when qparam_qspec does not quantize them.""" + spec = nvfp4_spec() + assert spec.scale_dtype is None + unnested = QuantizationSpec(dtype=torch.float4_e2m1fn_x2, granularity=spec.granularity) + assert unnested.scale_dtype == torch.float8_e8m0fnu + + +def test_qparam_qspec_from_dict_with_string_dtypes(): + """A nested spec loads from the dict form that YAML configs produce.""" + spec = QuantizationSpec( + **{ + "dtype": "uint4", + "qscheme": "asymmetric", + "qformulation": "minval", + "granularity": {"type": "per_block", "axis": 1, "block_size": 32}, + "qparam_qspec": { + "dtype": "uint8", + "qscheme": "asymmetric", + "granularity": {"type": "per_block", "block_size": 8}, + }, + } + ) + assert spec == minval_spec() + + +@pytest.mark.parametrize( + "overrides", + [ + pytest.param({"qparam_qspec": QuantizationSpec(dtype=torch.int8)}, id="signed_codes"), + pytest.param({"qparam_qspec": QuantizationSpec(dtype=torch.float8_e5m2)}, id="e5m2_codes"), + pytest.param( + {"qparam_qspec": minval_qparam_qspec(qscheme=QuantizationScheme.SYMMETRIC)}, + id="symmetric_unsigned_codes", + ), + pytest.param( + {"qparam_qspec": minval_qparam_qspec(qparam_calculator_cls="moving_average")}, + id="running_calculator", + ), + pytest.param( + {"qparam_qspec": minval_qparam_qspec(float_range=[None, 1.0])}, + id="pinned_float_range", + ), + pytest.param( + {"qparam_qspec": minval_qparam_qspec(granularity=PerChannelGranularity(axis=0))}, + id="per_channel_nested", + ), + pytest.param( + { + "qparam_qspec": minval_qparam_qspec( + granularity=PerBlockGranularity(block_size=(2, 8)) + ) + }, + id="multi_axis_nested", + ), + pytest.param( + { + "granularity": PerChannelGranularity(axis=0), + "qparam_qspec": minval_qparam_qspec(granularity=PerTensorGranularity()), + }, + id="per_channel_outer", + ), + pytest.param( + { + "qparam_qspec": minval_qparam_qspec( + qformulation=QuantizationFormulation.MINVAL, + qparam_qspec=minval_qparam_qspec(), + ) + }, + id="nested_twice", + ), + pytest.param({"qformulation": QuantizationFormulation.ZP}, id="asymmetric_zp_outer"), + ], +) +def test_general_qparam_qspec(overrides): + """Any quantization spec can quantize the qparams.""" + minval_spec(**overrides) # Validation runs on construction. + + +@pytest.mark.parametrize( + "overrides, match", + [ + pytest.param( + {"granularity": PerChannelGranularity(axis=0)}, + "follows this spec's block axis", + id="unresolvable_nested_block_axis", + ), + pytest.param( + {"qparam_qspec": minval_qparam_qspec(granularity=PerChannelGranularity())}, + "needs an explicit axis", + id="nested_per_channel_without_axis", + ), + ], +) +def test_invalid_qparam_qspec(overrides, match): + with pytest.raises(ValidationError, match=match): + minval_spec(**overrides) + + +def test_qparam_qspec_rejects_a_scale_dtype(): + """qparam_qspec defines how the scale is stored, so an explicit scale_dtype conflicts.""" + with pytest.raises(ValidationError, match="scale_dtype must be None when qparam_qspec is set"): + nvfp4_spec(scale_dtype=torch.float8_e8m0fnu) + + @pytest.mark.parametrize( "spec", [ @@ -947,6 +1057,8 @@ def test_invalid_scale_dtype(dtype, scale_dtype): qformulation=QuantizationFormulation.ZP, granularity=PerChannelGranularity(axis=0), ), + pytest.param(nvfp4_spec(), id="nvfp4"), + pytest.param(minval_spec(), id="minval_qparam_qspec"), ], ) def test_model_dump_round_trip(spec): diff --git a/tests/test_utils/general.py b/tests/test_utils/general.py index 6a42f42f..8acc4a3b 100644 --- a/tests/test_utils/general.py +++ b/tests/test_utils/general.py @@ -7,10 +7,19 @@ import importlib.util from decimal import ROUND_FLOOR, Decimal +from typing import Any import torch from torch.ops import coreai +from coreai_opt.quantization.spec import ( + PerBlockGranularity, + PerTensorGranularity, + QuantizationFormulation, + QuantizationScheme, + QuantizationSpec, +) + COREAI_AVAILABLE = importlib.util.find_spec("coreai") is not None @@ -161,3 +170,66 @@ def is_coreai_dequantize(target: object) -> bool: def get_quantize_dtype(node: torch.fx.Node) -> torch.dtype | None: """Return the quantized dtype carried in an FX node's args, else None.""" return next((a for a in node.args if isinstance(a, torch.dtype)), None) + + +def nvfp4_spec(axis: int | None = 1, **overrides: Any) -> QuantizationSpec: + """NVFP4: FP4 E2M1 elements in blocks of 16, FP8 E4M3 block scales under one tensor scale.""" + kwargs = { + "dtype": torch.float4_e2m1fn_x2, + "granularity": PerBlockGranularity(axis=axis, block_size=16), + "qparam_qspec": QuantizationSpec( + dtype=torch.float8_e4m3fn, granularity=PerTensorGranularity() + ), + } + return QuantizationSpec(**(kwargs | overrides)) + + +def minval_qparam_qspec(**overrides: Any) -> QuantizationSpec: + """uint8 asymmetric codes under one super-scale per 8 blocks.""" + kwargs = { + "dtype": torch.uint8, + "qscheme": QuantizationScheme.ASYMMETRIC, + "granularity": PerBlockGranularity(block_size=8), + } + return QuantizationSpec(**(kwargs | overrides)) + + +def minval_spec(axis: int | None = 1, **overrides: Any) -> QuantizationSpec: + """uint4 MINVAL elements in blocks of 32, with uint8 scale and minval codes. + + The two codes sit under separate super-scales, one of each per 8 blocks. + """ + kwargs = { + "dtype": torch.uint4, + "qscheme": QuantizationScheme.ASYMMETRIC, + "qformulation": QuantizationFormulation.MINVAL, + "granularity": PerBlockGranularity(axis=axis, block_size=32), + "qparam_qspec": minval_qparam_qspec(), + } + return QuantizationSpec(**(kwargs | overrides)) + + +def minval_fp16_super_scales_spec(axis: int | None = 1) -> QuantizationSpec: + """``minval_spec`` with FP16 super-scales, as GGUF Q4_K stores its ``d`` and ``dmin``.""" + return minval_spec(axis, qparam_qspec=minval_qparam_qspec(scale_dtype=torch.float16)) + + +def twice_quantized_spec(**nested_overrides: Any) -> QuantizationSpec: + """int4 blocks of 32 whose qparams are quantized twice. + + Each block's scale is a uint8 code under a super-scale shared by 4 blocks, and those + super-scales are FP8 E4M3 codes under one tensor scale. + """ + nested_kwargs = { + "dtype": torch.uint8, + "qscheme": QuantizationScheme.ASYMMETRIC, + "granularity": PerBlockGranularity(block_size=4), + "qparam_qspec": QuantizationSpec( + dtype=torch.float8_e4m3fn, granularity=PerTensorGranularity() + ), + } + return QuantizationSpec( + dtype=torch.int4, + granularity=PerBlockGranularity(axis=1, block_size=32), + qparam_qspec=QuantizationSpec(**(nested_kwargs | nested_overrides)), + )