diff --git a/modelopt/torch/kernels/quantization/ggml/q8_0.cu b/modelopt/torch/kernels/quantization/ggml/q8_0.cu index 5f9707bfcc7..a5bd9be02b1 100644 --- a/modelopt/torch/kernels/quantization/ggml/q8_0.cu +++ b/modelopt/torch/kernels/quantization/ggml/q8_0.cu @@ -53,7 +53,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, uint8_t *outpu payload[1] = static_cast(d_bits >> 8); } d = __shfl_sync(0xffffffff, d, 0); - const float rounded = d > 0.0f ? roundf(value / d) : 0.0f; + const float inverse = d > 0.0f ? 1.0f / d : 0.0f; + const float rounded = roundf(value * inverse); const int quant = static_cast(fminf(fmaxf(rounded, -kMaxQuant), kMaxQuant)); payload[kScaleBytes + lane] = static_cast(static_cast(quant)); } diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index 736b01ea7a9..5950fd7b087 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -27,7 +27,9 @@ from .iq2_xs import __all__ as _iq2_xs_all from .iq2_xxs import * from .iq2_xxs import __all__ as _iq2_xxs_all -from .registry import IQ_FORMAT_REGISTRY, IQFormat +from .q8_0 import * +from .q8_0 import __all__ as _q8_0_all +from .registry import GGML_FORMAT_REGISTRY, IQ_FORMAT_REGISTRY, GGMLFormat, IQFormat __all__ = [ # noqa: PLE0604 *_iq1_m_all, @@ -35,6 +37,9 @@ *_iq2_s_all, *_iq2_xs_all, *_iq2_xxs_all, + *_q8_0_all, + "GGML_FORMAT_REGISTRY", + "GGMLFormat", "IQ_FORMAT_REGISTRY", "IQFormat", ] diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py index 1f4a83eeff6..e8d9ecfd0cc 100644 --- a/modelopt/torch/quantization/ggml/backend.py +++ b/modelopt/torch/quantization/ggml/backend.py @@ -13,16 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TensorQuantizer backend dispatch for GGML-compatible weight-only IQ formats.""" +"""TensorQuantizer backend dispatch for GGML-compatible weight-only formats.""" import torch from ..nn.modules.tensor_quantizer import register_quant_backend -from .registry import IQ_FORMAT_REGISTRY +from .registry import GGML_FORMAT_REGISTRY def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: - """Dispatch an IQ quantizer to its format-specific implementation.""" + """Dispatch a GGML quantizer to its format-specific implementation.""" num_bits = getattr(quantizer, "num_bits", None) extra_args = getattr(quantizer, "backend_extra_args", None) or {} unknown_args = set(extra_args) - {"block_chunk_size", "decode_chunk_size"} @@ -30,11 +30,11 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}") # num_bits arrives untyped from the quantizer and is a tuple for scalar formats, # so narrow before the lookup rather than relying on the dict to reject it. - iq_format = IQ_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None - if iq_format is None: - supported = ", ".join(repr(name) for name in sorted(IQ_FORMAT_REGISTRY)) + ggml_format = GGML_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None + if ggml_format is None: + supported = ", ".join(repr(name) for name in sorted(GGML_FORMAT_REGISTRY)) raise ValueError(f"The ggml backend requires num_bits in ({supported})") - return iq_format.fake_quant(inputs, quantizer, **extra_args) + return ggml_format.fake_quant(inputs, quantizer, **extra_args) register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 6c14a29d0cf..96a490716e0 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -173,11 +173,11 @@ def fake_quantize_with_cache( @dataclass(frozen=True) -class IQFormat: - """Everything backend dispatch and export need to know about one IQ format. +class GGMLFormat: + """Everything backend dispatch and export need to know about one GGML format. Each format module declares one of these beside its encoder and decoder, and - :data:`~modelopt.torch.quantization.ggml.registry.IQ_FORMAT_REGISTRY` lists them. The + :data:`~modelopt.torch.quantization.ggml.registry.GGML_FORMAT_REGISTRY` lists them. The per-format pieces -- codebook, search, payload layout -- stay in the format's module; what lives here is the part every format does the same way. """ @@ -243,6 +243,11 @@ def pack(self, weight: torch.Tensor, quantizer=None) -> torch.Tensor: ) +# Compatibility alias for callers that imported the record type before the registry was +# generalized from the IQ family to every supported GGML block format. +IQFormat = GGMLFormat + + def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does. @@ -261,14 +266,16 @@ def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: return finite.float() -def validate_weight(weight: torch.Tensor, format_name: str) -> None: +def validate_weight( + weight: torch.Tensor, format_name: str, *, block_size: int = GGML_BLOCK_SIZE +) -> None: """Validate weight metadata accepted by the current GGML block encoders.""" if weight.numel() == 0: raise ValueError(f"{format_name} requires a non-empty weight") - if weight.dim() == 0 or weight.shape[-1] % GGML_BLOCK_SIZE: + if weight.dim() == 0 or weight.shape[-1] % block_size: raise ValueError( f"{format_name} requires the last weight dimension to be divisible by " - f"{GGML_BLOCK_SIZE}, got shape {tuple(weight.shape)}" + f"{block_size}, got shape {tuple(weight.shape)}" ) if not weight.is_floating_point(): raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") @@ -287,6 +294,7 @@ def validate_packed_weights( weight_shape: torch.Tensor, *, block_bytes: int, + block_size: int = GGML_BLOCK_SIZE, format_name: str, ) -> tuple[int, ...]: """Validate a packed payload and return its logical shape.""" @@ -303,9 +311,9 @@ def validate_packed_weights( if weight_shape.dim() != 1 or weight_shape.dtype not in integral_dtypes: raise ValueError("weight_shape must be a one-dimensional integral tensor") shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist()) - if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % GGML_BLOCK_SIZE: + if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % block_size: raise ValueError(f"invalid {format_name} logical weight shape: {shape}") - expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes + expected_payload_values = math.prod(shape) // block_size * block_bytes if packed_weights.numel() != expected_payload_values: raise ValueError("packed_weights size does not match weight_shape") return shape diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 90422b7a6ca..83aff654b8a 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -274,5 +274,5 @@ def dequantize_iq1_s( # Kept for callers of the per-format entry point. The record captured quantize_iq1_s and # dequantize_iq1_s when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq1_s_fake_quant = IQ1_S_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index a3cc22d2b54..9045d95de91 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -272,5 +272,5 @@ def dequantize_iq2_xs( # Kept for callers of the per-format entry point. The record captured quantize_iq2_xs and # dequantize_iq2_xs when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq2_xs_fake_quant = IQ2_XS_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py index ee35d44e3d4..3ac4d8a079b 100644 --- a/modelopt/torch/quantization/ggml/iq2_xxs.py +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -282,5 +282,5 @@ def dequantize_iq2_xxs( # Kept for callers of the per-format entry point. The record captured quantize_iq2_xxs and # dequantize_iq2_xxs when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq2_xxs_fake_quant = IQ2_XXS_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/q8_0.py b/modelopt/torch/quantization/ggml/q8_0.py new file mode 100644 index 00000000000..242a823835d --- /dev/null +++ b/modelopt/torch/quantization/ggml/q8_0.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Q8_0 fake quantization and GGML-compatible block packing. + +Every 32 logical values become one 34-byte ``block_q8_0`` payload: a +little-endian FP16 block scale followed by 32 signed int8 quants. +""" + +import torch + +from ..extensions import get_cuda_ext_ggml +from .common import ( + GGMLFormat, + narrow_to_float32, + validate_block_chunk_size, + validate_packed_weights, + validate_weight, +) + +__all__ = [ + "Q8_0_BLOCK_BYTES", + "Q8_0_BLOCK_SIZE", + "Q8_0_EFFECTIVE_BITS", + "Q8_0_FORMAT", + "dequantize_q8_0", + "q8_0_fake_quant", + "quantize_q8_0", +] + +Q8_0_BLOCK_SIZE = 32 +Q8_0_BLOCK_BYTES = 34 +Q8_0_EFFECTIVE_BITS = Q8_0_BLOCK_BYTES * 8 / Q8_0_BLOCK_SIZE +_Q8_0_MAX_QUANT = 127 +_DEFAULT_BLOCK_CHUNK_SIZE = 4096 +_DEFAULT_DECODE_CHUNK_SIZE = 16384 + + +def _encode_blocks(blocks: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 32-value blocks.""" + x = narrow_to_float32(blocks) + block_count = x.shape[0] + amax = x.abs().amax(dim=1) + # Keep the serialized fp16 scale finite. This extends the existing GGML backend policy for + # finite float64 values outside the float32/fp16 scale range. + d_float = (amax / _Q8_0_MAX_QUANT).clamp(max=65504.0) + d = d_float.to(torch.float16) + # Select quants with the unrounded float32 scale, not the serialized fp16 scale. + inverse = torch.where(d_float > 0, d_float.reciprocal(), torch.zeros_like(d_float)) + normalized = x * inverse.unsqueeze(1) + # C roundf, used by the canonical Q8_0 encoder, rounds half-way cases away from zero. + magnitude = normalized.abs() + rounded = normalized.sign() * torch.where( + magnitude.frac() == 0.5, + magnitude.floor() + 1, + magnitude.round(), + ) + quants = rounded.clamp(-_Q8_0_MAX_QUANT, _Q8_0_MAX_QUANT).to(torch.int8) + + packed = torch.empty((block_count, Q8_0_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:] = quants.view(torch.uint8) + return packed + + +@torch.no_grad() +def quantize_q8_0( + weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible Q8_0 blocks.""" + validate_weight(weight, "Q8_0", block_size=Q8_0_BLOCK_SIZE) + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, Q8_0_BLOCK_SIZE) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + packed = extension.q8_0_pack(blocks) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // Q8_0_BLOCK_SIZE, + Q8_0_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size]) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // Q8_0_BLOCK_SIZE, + Q8_0_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_q8_0( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, + block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """Decode GGML-compatible Q8_0 payload bytes.""" + shape = validate_packed_weights( + packed_weights, + weight_shape, + block_bytes=Q8_0_BLOCK_BYTES, + block_size=Q8_0_BLOCK_SIZE, + format_name="Q8_0", + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, Q8_0_BLOCK_BYTES) + decoded = torch.empty((blocks.shape[0], Q8_0_BLOCK_SIZE), dtype=dtype, device=blocks.device) + for start in range(0, blocks.shape[0], block_chunk_size): + stop = min(start + block_chunk_size, blocks.shape[0]) + block_chunk = blocks[start:stop] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + quants = block_chunk[:, 2:].contiguous().view(torch.int8).float() + decoded[start:stop] = (d.unsqueeze(1) * quants).to(dtype) + return decoded.reshape(shape) + + +Q8_0_FORMAT = GGMLFormat( + name="q8_0", + block_size=Q8_0_BLOCK_SIZE, + block_bytes=Q8_0_BLOCK_BYTES, + quantize=quantize_q8_0, + dequantize=dequantize_q8_0, + block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, +) + +# Keep the public per-format entry point on the same record used by backend dispatch. +q8_0_fake_quant = Q8_0_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/registry.py b/modelopt/torch/quantization/ggml/registry.py index 8183b3a93ea..c95e9569fe7 100644 --- a/modelopt/torch/quantization/ggml/registry.py +++ b/modelopt/torch/quantization/ggml/registry.py @@ -13,24 +13,38 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The GGML IQ formats, listed once for backend dispatch and export.""" +"""GGML block formats, listed once for backend dispatch and export.""" -from .common import IQFormat +from .common import GGMLFormat, IQFormat from .iq1_m import IQ1_M_FORMAT from .iq1_s import IQ1_S_FORMAT from .iq2_s import IQ2_S_FORMAT from .iq2_xs import IQ2_XS_FORMAT from .iq2_xxs import IQ2_XXS_FORMAT +from .q8_0 import Q8_0_FORMAT -__all__ = ["IQ_FORMAT_REGISTRY", "IQFormat"] +__all__ = ["GGML_FORMAT_REGISTRY", "IQ_FORMAT_REGISTRY", "GGMLFormat", "IQFormat"] -# Every IQ format, keyed by the name a quantizer's num_bits carries, in increasing bits per -# weight. Backend dispatch and both exporters read this mapping and export's IQ_FORMATS is derived -# from it, so adding a format is one entry here rather than a row in several parallel tables. +# Every GGML block format, keyed by the name a quantizer's num_bits carries, in increasing bits +# per weight. Backend dispatch reads this mapping, and export will derive its supported formats +# from it rather than maintaining parallel metadata and packer tables. # # It is an explicit list rather than formats registering themselves on import, so its contents # never depend on which modules happen to have been imported first. -IQ_FORMAT_REGISTRY: dict[str, IQFormat] = { +GGML_FORMAT_REGISTRY: dict[str, GGMLFormat] = { fmt.name: fmt - for fmt in (IQ1_S_FORMAT, IQ1_M_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, IQ2_S_FORMAT) + for fmt in ( + IQ1_S_FORMAT, + IQ1_M_FORMAT, + IQ2_XXS_FORMAT, + IQ2_XS_FORMAT, + IQ2_S_FORMAT, + Q8_0_FORMAT, + ) +} + +# Compatibility view for the already-merged IQ export path. Export migrates to the general +# registry in the Q8_0 export PR; IQ-specific conformance tests continue to use this subset. +IQ_FORMAT_REGISTRY: dict[str, GGMLFormat] = { + name: fmt for name, fmt in GGML_FORMAT_REGISTRY.items() if name.startswith("iq") } diff --git a/tests/gpu/torch/quantization/test_q8_0_cuda.py b/tests/gpu/torch/quantization/test_q8_0_cuda.py new file mode 100644 index 00000000000..1c8f63fcbae --- /dev/null +++ b/tests/gpu/torch/quantization/test_q8_0_cuda.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +import modelopt.torch.quantization.ggml.q8_0 as q8_0_module +from modelopt.torch.quantization.extensions import get_cuda_ext_ggml +from modelopt.torch.quantization.ggml.q8_0 import dequantize_q8_0, quantize_q8_0 + + +def _extension(): + extension = get_cuda_ext_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +def test_q8_0_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 64), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _extension().q8_0_pack(weight).reshape(8, 2, 34) + packed_again = _extension().q8_0_pack(weight).reshape(8, 2, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, shape = quantize_q8_0(weight) + reconstructed = dequantize_q8_0(packed, shape) + + assert packed.shape == (8, 2, 34) + assert torch.equal(packed, packed_again) + assert torch.equal(packed, reference) + assert shape.device.type == "cpu" + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 1e-4 + + +def test_q8_0_cuda_zero_and_nonfinite_encoding_match_reference(monkeypatch): + weight = torch.zeros((2, 32), device="cuda", dtype=torch.bfloat16) + weight[1, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") + + packed = _extension().q8_0_pack(weight).reshape(2, 1, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, shape = quantize_q8_0(weight) + + assert not packed.any() + assert torch.equal(packed, reference) + assert torch.equal(dequantize_q8_0(packed, shape), torch.zeros_like(weight)) + + +def test_q8_0_cuda_falls_back_to_pytorch_encoder(monkeypatch): + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((2, 32), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight) + reconstructed = dequantize_q8_0(packed, shape) + + assert packed.shape == (2, 1, 34) + assert torch.isfinite(reconstructed).all() + + +def test_q8_0_cuda_float64_matches_pytorch_encoder(): + weight = torch.randn(4, 32, dtype=torch.float64, generator=torch.Generator().manual_seed(7)) + + reference, _ = quantize_q8_0(weight) + packed, _ = quantize_q8_0(weight.cuda()) + + assert torch.equal(reference, packed.cpu()) + + +@pytest.mark.parametrize( + ("values", "expected"), + [ + ([5207.0, 20.5, -20.5], [127, 0, 0]), + ([1.0, 3.4999 / 127, -3.4999 / 127], [127, 3, -3]), + ], + ids=["reciprocal-rounding", "unrounded-scale"], +) +def test_q8_0_cuda_reciprocal_rounding_boundary_matches_pytorch_encoder( + monkeypatch, values, expected +): + weight = torch.zeros((1, 32), device="cuda", dtype=torch.float32) + weight[0, :3] = torch.tensor(values, device="cuda") + + packed = _extension().q8_0_pack(weight).reshape(1, 1, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, _ = quantize_q8_0(weight) + + assert torch.equal(packed, reference) + assert packed[0, 0, 2:5].view(torch.int8).tolist() == expected diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index cffea63cb74..29c122c7107 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -25,15 +25,16 @@ import modelopt.torch.quantization.ggml as ggml import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +import modelopt.torch.quantization.ggml.q8_0 as q8_0_module from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.backend import ggml_fake_quant from modelopt.torch.quantization.ggml.common import narrow_to_float32 from modelopt.torch.quantization.nn import TensorQuantizer # Every registered format, paired with its own module. The module -- not the registry -- supplies # what a test expects, so a mis-wired registry entry cannot make both sides of an assertion agree. -FORMAT_NAMES = sorted(IQ_FORMAT_REGISTRY) +FORMAT_NAMES = sorted(GGML_FORMAT_REGISTRY) FORMAT_MODULES = [ (name, importlib.import_module(f"modelopt.torch.quantization.ggml.{name}")) for name in FORMAT_NAMES @@ -43,11 +44,11 @@ def _substitute(monkeypatch, name, **changes): """Replace parts of one registered format for the duration of a test. - Dispatch reads IQ_FORMAT_REGISTRY, so that is where a substitute has to go: patching the + Dispatch reads GGML_FORMAT_REGISTRY, so that is where a substitute has to go: patching the format module's function would not reach it. """ monkeypatch.setitem( - IQ_FORMAT_REGISTRY, name, dataclasses.replace(IQ_FORMAT_REGISTRY[name], **changes) + GGML_FORMAT_REGISTRY, name, dataclasses.replace(GGML_FORMAT_REGISTRY[name], **changes) ) @@ -89,6 +90,7 @@ def test_ggml_backend_rejects_unknown_format(): def test_ggml_codecs_are_exported_from_quantization_package(): assert mtq.quantize_iq1_s is iq1_s_module.quantize_iq1_s assert mtq.quantize_iq2_xs is iq2_xs_module.quantize_iq2_xs + assert mtq.quantize_q8_0 is q8_0_module.quantize_q8_0 @pytest.mark.parametrize( @@ -108,7 +110,7 @@ def fake_quant(inputs, _quantizer, **kwargs): return inputs # The dispatcher resolves through the registry, so that is the seam to patch. - monkeypatch.setitem(IQ_FORMAT_REGISTRY, "iq1_s", SimpleNamespace(fake_quant=fake_quant)) + monkeypatch.setitem(GGML_FORMAT_REGISTRY, "iq1_s", SimpleNamespace(fake_quant=fake_quant)) inputs = torch.ones(1, 256) quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args=extra_args) @@ -125,7 +127,8 @@ def test_ggml_backend_rejects_unknown_extra_arg(): @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_ggml_backend_caches_packed_weight_and_invalidates_on_change(monkeypatch, num_bits, module): - weight = torch.randn(1, 256) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size + weight = torch.randn(1, block_size) quantizer = SimpleNamespace(num_bits=num_bits, _quantizer_cache=None) original_quantize = getattr(module, f"quantize_{num_bits}") call_count = 0 @@ -136,7 +139,7 @@ def counted_quantize(*args, **kwargs): return original_quantize(*args, **kwargs) _substitute(monkeypatch, num_bits, quantize=counted_quantize) - fake_quant = IQ_FORMAT_REGISTRY[num_bits].fake_quant + fake_quant = GGML_FORMAT_REGISTRY[num_bits].fake_quant fake_quant(weight, quantizer, block_chunk_size=1) fake_quant(weight, quantizer, block_chunk_size=1) @@ -183,8 +186,9 @@ def counting(weight, **kwargs): return original(weight, **kwargs) _substitute(monkeypatch, num_bits, quantize=counting) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: block_size}, backend="ggml") ) weight = torch.randn(4, 256) @@ -192,7 +196,9 @@ def counting(weight, **kwargs): for _ in range(5): quantizer(weight) - assert calls == [(4, 256)], f"expected one pack, got {len(calls)}" + assert calls == [(weight.numel() // block_size, block_size)], ( + f"expected one pack, got {len(calls)}" + ) @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) @@ -212,8 +218,9 @@ def recording(packed_weights, weight_shape, **kwargs): return original(packed_weights, weight_shape, **kwargs) _substitute(monkeypatch, num_bits, dequantize=recording) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: block_size}, backend="ggml") ) quantizer(torch.randn(4, 256)) @@ -226,13 +233,23 @@ def test_registry_lists_every_exported_encoder(): encoders = { name.removeprefix("quantize_") for name in ggml.__all__ if name.startswith("quantize_") } - assert encoders == set(IQ_FORMAT_REGISTRY) + assert encoders == set(GGML_FORMAT_REGISTRY) + + +def test_iq_registry_remains_a_compatible_subset_of_ggml_registry(): + assert ggml.IQFormat is ggml.GGMLFormat + assert set(ggml.IQ_FORMAT_REGISTRY) == { + name for name in GGML_FORMAT_REGISTRY if name.startswith("iq") + } + assert "q8_0" not in ggml.IQ_FORMAT_REGISTRY + for name, record in ggml.IQ_FORMAT_REGISTRY.items(): + assert record is GGML_FORMAT_REGISTRY[name] @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_registry_record_is_wired_to_its_own_codec(num_bits, module): """Each record points at its own format's encoder, decoder, geometry and chunk defaults.""" - record = IQ_FORMAT_REGISTRY[num_bits] + record = GGML_FORMAT_REGISTRY[num_bits] upper = num_bits.upper() assert record.name == num_bits @@ -248,7 +265,7 @@ def test_registry_record_is_wired_to_its_own_codec(num_bits, module): @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_public_fake_quant_is_the_registered_record(num_bits, module): """The per-format entry point and backend dispatch run the same code path.""" - assert getattr(module, f"{num_bits}_fake_quant").__self__ is IQ_FORMAT_REGISTRY[num_bits] + assert getattr(module, f"{num_bits}_fake_quant").__self__ is GGML_FORMAT_REGISTRY[num_bits] @pytest.mark.parametrize("num_bits", FORMAT_NAMES) @@ -260,6 +277,9 @@ def test_format_fake_quant_rejects_another_formats_quantizer(num_bits): """ other = next(name for name in FORMAT_NAMES if name != num_bits) expected = f"The ggml {num_bits.upper()} backend requires num_bits={num_bits!r}" + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size with pytest.raises(ValueError, match=re.escape(expected)): - IQ_FORMAT_REGISTRY[num_bits].fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits=other)) + GGML_FORMAT_REGISTRY[num_bits].fake_quant( + torch.ones(1, block_size), SimpleNamespace(num_bits=other) + ) diff --git a/tests/unit/torch/quantization/test_q8_0.py b/tests/unit/torch/quantization/test_q8_0.py new file mode 100644 index 00000000000..53f9f1c72e7 --- /dev/null +++ b/tests/unit/torch/quantization/test_q8_0.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.ggml.q8_0 import ( + Q8_0_BLOCK_BYTES, + dequantize_q8_0, + q8_0_fake_quant, + quantize_q8_0, +) + + +def test_q8_0_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 32), dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight) + + assert packed.shape == (2, 1, Q8_0_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert torch.equal(dequantize_q8_0(packed, shape), weight) + + +def test_q8_0_uses_roundf_ties_away_from_zero(): + weight = torch.zeros((1, 32), dtype=torch.float32) + below_half = torch.nextafter(torch.tensor(0.5), torch.tensor(0.0)) + weight[0, :5] = torch.tensor([127.0, 0.5, -0.5, 1.5, -1.5]) + weight[0, 5:7] = torch.stack((below_half, -below_half)) + + packed, shape = quantize_q8_0(weight) + + assert packed[0, 0, :2].contiguous().view(torch.float16).item() == 1.0 + assert packed[0, 0, 2:9].contiguous().view(torch.int8).tolist() == [ + 127, + 1, + -1, + 2, + -2, + 0, + 0, + ] + expected = weight.clone() + expected[0, 1:5] = torch.tensor([1.0, -1.0, 2.0, -2.0]) + expected[0, 5:7] = 0 + assert torch.equal(dequantize_q8_0(packed, shape, dtype=torch.float32), expected) + + +def test_q8_0_selects_quants_before_rounding_scale_to_fp16(): + weight = torch.zeros((1, 32), dtype=torch.float32) + weight[0, :3] = torch.tensor([1.0, 3.4999 / 127, -3.4999 / 127]) + + packed, _ = quantize_q8_0(weight) + + scale = packed[0, 0, :2].contiguous().view(torch.float16).item() + assert scale == torch.tensor(1.0 / 127, dtype=torch.float16).item() + # The unrounded scale gives +/-3.4999; using the stored scale would cross +/-3.5. + assert packed[0, 0, 2:5].contiguous().view(torch.int8).tolist() == [127, 3, -3] + + +def test_q8_0_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 64), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight, block_chunk_size=1) + reconstructed = dequantize_q8_0(packed, shape, dtype=torch.float32, block_chunk_size=1) + + assert packed.shape == (2, 2, 34) + normalized_mse = ( + reconstructed - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 1e-4 + assert torch.all(packed.reshape(-1, 34)[:, :2].contiguous().view(torch.float16) > 0) + + +def test_q8_0_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="divisible by 32"): + quantize_q8_0(torch.ones(2, 33)) + + +def test_q8_0_treats_nonfinite_values_as_zero(): + weight = torch.randn(1, 32) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf]) + + packed, _ = quantize_q8_0(weight) + expected, _ = quantize_q8_0(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0)) + + assert torch.equal(packed, expected) + + +@pytest.mark.parametrize("shape", [[1, 31], [0, 32]]) +def test_q8_0_rejects_invalid_shape_metadata(shape): + packed = torch.zeros((1, 1, 34), dtype=torch.uint8) + + with pytest.raises(ValueError, match="logical weight shape"): + dequantize_q8_0(packed, torch.tensor(shape)) + + +def test_q8_0_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "q8_0" + + weight = torch.randn(1, 32, requires_grad=True) + q8_0_fake_quant(weight, Quantizer()).sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight))