mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
[OMNIML-5899] Add IQ quantization codecs and backend (#2446)
## Summary - add IQ1_S and IQ2_XS reference codecs and a weight-only fake-quant backend - register and export both formats from the quantization package - cache compact packed weights across unchanged forwards and invalidate on tensor or config changes - use one Python-side IQ2_XS FP16 scale predictor for both reference and CUDA packing - validate packed payload metadata, normalize CUDA cache keys, and define a shared non-finite policy ## PR split This work is split into four focused PRs. Each PR targets `main` and owns a disjoint file set: 1. **Kernel** — [#2448: Add CUDA kernels for IQ packing](https://github.com/NVIDIA/Model-Optimizer/pull/2448) 2. **Quantization** — [#2446: Add IQ quantization codecs and backend](https://github.com/NVIDIA/Model-Optimizer/pull/2446) 3. **Export** — [#2447: Export IQ checkpoints from HF and Megatron](https://github.com/NVIDIA/Model-Optimizer/pull/2447) 4. **Recipes** — [#2449: Add IQ post-training quantization recipes](https://github.com/NVIDIA/Model-Optimizer/pull/2449) The required merge order is #2448, #2446, #2447, then #2449. ## Scope This PR owns the Python codecs, backend dispatch, package registration, license attribution, CPU codec/backend tests, and CUDA numerical/reference-path tests. The native CUDA layer and direct extension tests remain in #2448; export and recipes remain in their own PRs. ## Why the codecs are separate from `qtensor` The new `ggml/` package contains stateless reference codecs and fake-quant backend functions. They transform ordinary tensors into packed format payloads and reconstruct tensors for fake quantization; they do not define persistent runtime quantized-tensor objects. `BaseQuantizedTensor` subclasses under `qtensor/` own runtime tensor objects and execution dispatch. Keeping the codecs separate avoids claiming a runtime tensor contract that these formats do not yet provide. A `qtensor` type can be added later if a runtime execution path requires one. ## Compatibility boundary The Python encoders intentionally use fixed-scale, unweighted searches. They are not intended to reproduce another encoder's bytes for every input when that encoder performs iterative scale refinement or importance weighting. Compatibility is defined by the canonical codebooks, 50/74-byte payload layouts, and pinned dequantization formulas. IQ2_XS computes the FP16 superblock scale once in the Python predictor and passes it to the CUDA packer. This removes a duplicate floating-point reduction and makes native/reference byte parity use the same scale. Non-finite input elements are treated as zero during packing in both implementations. The unit tests construct nonzero payload fields independently and validate metadata, signs, local scales, and global scales. The CUDA tests compare native packed bytes with this Python reference encoder. ## Test coverage - [IQ1_S CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq1_s.py) - [IQ2_XS CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq2_xs.py) - [registered backend and cache tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_ggml_backend.py) - [IQ1_S CUDA byte-parity, numerical, non-finite, zero-payload, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq1_s_cuda.py) - [IQ2_XS CUDA byte-parity, numerical, non-finite, zero-payload, underflow, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq2_xs_cuda.py) ## Licensing The embedded codebook data cites the pinned upstream MIT source, carries its license notice, and uses the repository's third-party license mechanism. Human OSRB/code-owner confirmation is still required; this PR does not claim that approval. ## Validation - focused lint, format, and type checks pass for all changed Python files - 36 focused CPU codec and backend tests pass locally - all 20 direct-extension and CUDA integration test cases collect locally; runtime CUDA execution remains delegated to GPU CI - restricted-term scan passes <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added GGML quantization support for IQ1_S and IQ2_XS formats. - Added quantization, dequantization, and fake-quantization workflows with pass-through gradients. - Added CPU fallback when CUDA acceleration is unavailable. - Added validation for packed weights, tensor shapes, formats, and backend options. - Added configurable chunk processing and caching for repeated quantization. - **Tests** - Added comprehensive CPU and CUDA coverage for accuracy, validation, caching, fallback behavior, and edge cases. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
hychiang
co-authored by
Chenjie Luo
Claude Opus 5
Chenjie Luo
Claude Opus 5
parent
9895d6f129
commit
9e3d555aa1
@@ -103,6 +103,7 @@ repos:
|
|||||||
exclude: >
|
exclude: >
|
||||||
(?x)^(
|
(?x)^(
|
||||||
modelopt/torch/quantization/utils/calib_utils.py|
|
modelopt/torch/quantization/utils/calib_utils.py|
|
||||||
|
modelopt/torch/quantization/ggml/codebooks.py|
|
||||||
modelopt/onnx/quantization/operators.py|
|
modelopt/onnx/quantization/operators.py|
|
||||||
modelopt/onnx/quantization/ort_patching.py|
|
modelopt/onnx/quantization/ort_patching.py|
|
||||||
modelopt/torch/_deploy/utils/onnx_utils.py|
|
modelopt/torch/_deploy/utils/onnx_utils.py|
|
||||||
|
|||||||
@@ -250,6 +250,7 @@ the following copyright holders, licensed under the MIT License:
|
|||||||
Copyright (c) 2023 DeepSeek
|
Copyright (c) 2023 DeepSeek
|
||||||
Copyright (c) 2025 sgl-project
|
Copyright (c) 2025 sgl-project
|
||||||
Copyright (c) 2026 The DeepSpec Authors
|
Copyright (c) 2026 The DeepSpec Authors
|
||||||
|
Copyright (c) 2023-2026 The ggml authors
|
||||||
|
|
||||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
of this software and associated documentation files (the "Software"), to deal
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
|||||||
@@ -15,6 +15,8 @@
|
|||||||
|
|
||||||
"""Quantization package."""
|
"""Quantization package."""
|
||||||
|
|
||||||
|
from importlib import import_module as _import_module
|
||||||
|
|
||||||
# Initialize mode and plugins
|
# Initialize mode and plugins
|
||||||
from . import mode, plugins, utils
|
from . import mode, plugins, utils
|
||||||
|
|
||||||
@@ -25,3 +27,9 @@ from .conversion import *
|
|||||||
from .model_quant import *
|
from .model_quant import *
|
||||||
from .nn.modules.quant_module import QuantModuleRegistry
|
from .nn.modules.quant_module import QuantModuleRegistry
|
||||||
from .utils import update_quant_cfg_with_kv_cache_quant
|
from .utils import update_quant_cfg_with_kv_cache_quant
|
||||||
|
|
||||||
|
# Imported last to register the backend without cycling through quantization.qtensor.
|
||||||
|
# A dynamic import prevents isort from hoisting it into the import block above.
|
||||||
|
ggml = _import_module(".ggml", __name__)
|
||||||
|
globals().update({name: getattr(ggml, name) for name in ggml.__all__})
|
||||||
|
del _import_module
|
||||||
|
|||||||
@@ -0,0 +1,25 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""GGML-compatible block quantization formats."""
|
||||||
|
|
||||||
|
# Importing the backend installs its TensorQuantizer dispatch entry.
|
||||||
|
from . import backend as _backend
|
||||||
|
from .iq1_s import *
|
||||||
|
from .iq1_s import __all__ as _iq1_s_all
|
||||||
|
from .iq2_xs import *
|
||||||
|
from .iq2_xs import __all__ as _iq2_xs_all
|
||||||
|
|
||||||
|
__all__ = [*_iq1_s_all, *_iq2_xs_all] # noqa: PLE0604
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""TensorQuantizer backend dispatch for GGML-compatible weight-only IQ formats."""
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from ..nn.modules.tensor_quantizer import register_quant_backend
|
||||||
|
from .iq1_s import iq1_s_fake_quant
|
||||||
|
from .iq2_xs import iq2_xs_fake_quant
|
||||||
|
|
||||||
|
|
||||||
|
def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor:
|
||||||
|
"""Dispatch an IQ 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"}
|
||||||
|
if unknown_args:
|
||||||
|
raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}")
|
||||||
|
if num_bits == "iq1_s":
|
||||||
|
return iq1_s_fake_quant(inputs, quantizer, **extra_args)
|
||||||
|
if num_bits == "iq2_xs":
|
||||||
|
return iq2_xs_fake_quant(inputs, quantizer, **extra_args)
|
||||||
|
raise ValueError("The ggml backend requires num_bits='iq1_s' or 'iq2_xs'")
|
||||||
|
|
||||||
|
|
||||||
|
register_quant_backend("ggml", ggml_fake_quant)
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
# This file includes the IQ1_S and IQ2_XS codebooks adapted from:
|
||||||
|
# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
|
||||||
|
#
|
||||||
|
# MIT License
|
||||||
|
#
|
||||||
|
# Copyright (c) 2023-2026 The ggml authors
|
||||||
|
#
|
||||||
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
# of this software and associated documentation files (the "Software"), to deal
|
||||||
|
# in the Software without restriction, including without limitation the rights
|
||||||
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
# copies of the Software, and to permit persons to whom the Software is
|
||||||
|
# furnished to do so, subject to the following conditions:
|
||||||
|
#
|
||||||
|
# The above copyright notice and this permission notice shall be included in all
|
||||||
|
# copies or substantial portions of the Software.
|
||||||
|
#
|
||||||
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
# SOFTWARE.
|
||||||
|
#
|
||||||
|
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||||
|
# SPDX-License-Identifier: Apache-2.0 AND MIT
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
"""Canonical IQ codebooks, carried verbatim from the GGML tables.
|
||||||
|
|
||||||
|
This module holds only data and its decoding. The tables are the one piece of GGML material
|
||||||
|
reproduced in this package, so keeping them here bounds the MIT-licensed surface to a single
|
||||||
|
file and leaves the codecs beside it as reviewable logic.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import zlib
|
||||||
|
from functools import cache
|
||||||
|
|
||||||
|
# zlib-compressed little-endian bytes of the canonical uint64_t table. The
|
||||||
|
# decoded int8 values are -1, 0, and 1.
|
||||||
|
_IQ1_S_GRID_ZLIB_B64 = (
|
||||||
|
"eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv"
|
||||||
|
"c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY"
|
||||||
|
"4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX"
|
||||||
|
"bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8"
|
||||||
|
"Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W"
|
||||||
|
"YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH"
|
||||||
|
"9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL"
|
||||||
|
"2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E"
|
||||||
|
"UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg"
|
||||||
|
"MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0"
|
||||||
|
"jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+"
|
||||||
|
"LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs"
|
||||||
|
"fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2"
|
||||||
|
"a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g"
|
||||||
|
"n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq"
|
||||||
|
"ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ"
|
||||||
|
"q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/"
|
||||||
|
"1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2"
|
||||||
|
"kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM"
|
||||||
|
"T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8"
|
||||||
|
"0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT"
|
||||||
|
"oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy"
|
||||||
|
"i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44"
|
||||||
|
"rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg"
|
||||||
|
"O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO"
|
||||||
|
"k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI"
|
||||||
|
"Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8"
|
||||||
|
"73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn"
|
||||||
|
"gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6"
|
||||||
|
"hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+"
|
||||||
|
"EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA"
|
||||||
|
"0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH"
|
||||||
|
"OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD"
|
||||||
|
"FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn"
|
||||||
|
"h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6"
|
||||||
|
"DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB"
|
||||||
|
"MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI"
|
||||||
|
"4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh"
|
||||||
|
"eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ"
|
||||||
|
"zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep"
|
||||||
|
"V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5"
|
||||||
|
"/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C"
|
||||||
|
"qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T"
|
||||||
|
"rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD"
|
||||||
|
"hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr"
|
||||||
|
"x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO"
|
||||||
|
"umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0"
|
||||||
|
"Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA"
|
||||||
|
"MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ"
|
||||||
|
"jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8="
|
||||||
|
)
|
||||||
|
|
||||||
|
# Compact byte representation of the canonical [512, 8] grid. Values are only
|
||||||
|
# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids
|
||||||
|
# adding a pickle-backed torch.save artifact to the wheel.
|
||||||
|
_IQ2_XS_GRID_B64 = (
|
||||||
|
"CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr"
|
||||||
|
"CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI"
|
||||||
|
"CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI"
|
||||||
|
"GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI"
|
||||||
|
"GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI"
|
||||||
|
"CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr"
|
||||||
|
"CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr"
|
||||||
|
"CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI"
|
||||||
|
"CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI"
|
||||||
|
"KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ"
|
||||||
|
"CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI"
|
||||||
|
"CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI"
|
||||||
|
"CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI"
|
||||||
|
"CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ"
|
||||||
|
"GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI"
|
||||||
|
"CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI"
|
||||||
|
"KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ"
|
||||||
|
"CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI"
|
||||||
|
"GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI"
|
||||||
|
"GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI"
|
||||||
|
"GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ"
|
||||||
|
"CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI"
|
||||||
|
"CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI"
|
||||||
|
"CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI"
|
||||||
|
"KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI"
|
||||||
|
"GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr"
|
||||||
|
"KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr"
|
||||||
|
"KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI"
|
||||||
|
"CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI"
|
||||||
|
"GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ"
|
||||||
|
"GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr"
|
||||||
|
"CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ"
|
||||||
|
"CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI"
|
||||||
|
"CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr"
|
||||||
|
"CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI"
|
||||||
|
"CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI"
|
||||||
|
"CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ"
|
||||||
|
"GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI"
|
||||||
|
"CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr"
|
||||||
|
"GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ"
|
||||||
|
"GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI"
|
||||||
|
"CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ"
|
||||||
|
"KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI"
|
||||||
|
"KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr"
|
||||||
|
"CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI"
|
||||||
|
"KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI"
|
||||||
|
"GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr"
|
||||||
|
"GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr"
|
||||||
|
"CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI"
|
||||||
|
"KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ"
|
||||||
|
"GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI"
|
||||||
|
"CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ"
|
||||||
|
"GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr"
|
||||||
|
"KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ"
|
||||||
|
"CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw=="
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@cache
|
||||||
|
def iq1_s_grid_bytes() -> bytes:
|
||||||
|
"""Decoded little-endian int8 bytes of the [2048, 8] IQ1_S ternary table."""
|
||||||
|
return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64))
|
||||||
|
|
||||||
|
|
||||||
|
@cache
|
||||||
|
def iq2_xs_grid_bytes() -> bytes:
|
||||||
|
"""Decoded bytes of the [512, 8] IQ2_XS magnitude table."""
|
||||||
|
return base64.b64decode(_IQ2_XS_GRID_B64)
|
||||||
@@ -0,0 +1,163 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""Shared validation for GGML-compatible block quantizers."""
|
||||||
|
|
||||||
|
import math
|
||||||
|
import weakref
|
||||||
|
from collections.abc import Callable
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
GGML_BLOCK_SIZE = 256
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _PackedWeightCache:
|
||||||
|
input_ref: weakref.ReferenceType
|
||||||
|
input_key: tuple[object, ...]
|
||||||
|
format_name: str
|
||||||
|
block_chunk_size: int
|
||||||
|
packed_weights: torch.Tensor
|
||||||
|
weight_shape: torch.Tensor
|
||||||
|
|
||||||
|
|
||||||
|
def _input_cache_key(inputs: torch.Tensor) -> tuple[object, ...] | None:
|
||||||
|
try:
|
||||||
|
version = inputs._version
|
||||||
|
except RuntimeError:
|
||||||
|
# Inference tensors can omit version counters, so changes cannot be detected safely.
|
||||||
|
return None
|
||||||
|
return (
|
||||||
|
inputs.data_ptr(),
|
||||||
|
tuple(inputs.shape),
|
||||||
|
tuple(inputs.stride()),
|
||||||
|
inputs.dtype,
|
||||||
|
inputs.device,
|
||||||
|
version,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def fake_quantize_with_cache(
|
||||||
|
inputs: torch.Tensor,
|
||||||
|
quantizer,
|
||||||
|
*,
|
||||||
|
format_name: str,
|
||||||
|
block_chunk_size: int,
|
||||||
|
quantize: Callable[..., tuple[torch.Tensor, torch.Tensor]],
|
||||||
|
dequantize: Callable[..., torch.Tensor],
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Fake-quantize a weight while caching its compact packed representation."""
|
||||||
|
input_key = _input_cache_key(inputs)
|
||||||
|
cache = getattr(quantizer, "_quantizer_cache", None)
|
||||||
|
if (
|
||||||
|
isinstance(cache, _PackedWeightCache)
|
||||||
|
and input_key is not None
|
||||||
|
and cache.input_ref() is inputs
|
||||||
|
and cache.input_key == input_key
|
||||||
|
and cache.format_name == format_name
|
||||||
|
and cache.block_chunk_size == block_chunk_size
|
||||||
|
):
|
||||||
|
packed_weights, weight_shape = cache.packed_weights, cache.weight_shape
|
||||||
|
else:
|
||||||
|
packed_weights, weight_shape = quantize(inputs, block_chunk_size=block_chunk_size)
|
||||||
|
if input_key is not None:
|
||||||
|
quantizer._quantizer_cache = _PackedWeightCache(
|
||||||
|
input_ref=weakref.ref(inputs),
|
||||||
|
input_key=input_key,
|
||||||
|
format_name=format_name,
|
||||||
|
block_chunk_size=block_chunk_size,
|
||||||
|
packed_weights=packed_weights,
|
||||||
|
weight_shape=weight_shape,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
quantizer._quantizer_cache = None
|
||||||
|
|
||||||
|
reconstructed = dequantize(
|
||||||
|
packed_weights,
|
||||||
|
weight_shape,
|
||||||
|
dtype=inputs.dtype,
|
||||||
|
block_chunk_size=block_chunk_size,
|
||||||
|
)
|
||||||
|
return inputs + (reconstructed - inputs).detach()
|
||||||
|
|
||||||
|
|
||||||
|
def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does.
|
||||||
|
|
||||||
|
Non-finite elements become zero, and finite elements outside the float32 range saturate
|
||||||
|
instead of overflowing to infinity and then being zeroed. Sanitizing at the source precision
|
||||||
|
is what keeps the reference encoders byte-identical to the extension for float64 weights;
|
||||||
|
converting first would turn a finite 1e100 into zero on this path and into the float32
|
||||||
|
maximum on the CUDA one.
|
||||||
|
"""
|
||||||
|
finite = torch.nan_to_num(blocks, nan=0.0, posinf=0.0, neginf=0.0)
|
||||||
|
if finite.dtype == torch.float64:
|
||||||
|
# Only float64 can hold a finite value the narrowing would overflow. The float32 bounds
|
||||||
|
# do not fit in the narrower dtypes, so clamping them would raise rather than no-op.
|
||||||
|
info = torch.finfo(torch.float32)
|
||||||
|
finite = finite.clamp(info.min, info.max)
|
||||||
|
return finite.float()
|
||||||
|
|
||||||
|
|
||||||
|
def validate_weight(weight: torch.Tensor, format_name: str) -> 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:
|
||||||
|
raise ValueError(
|
||||||
|
f"{format_name} requires the last weight dimension to be divisible by "
|
||||||
|
f"{GGML_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}")
|
||||||
|
|
||||||
|
|
||||||
|
def validate_block_chunk_size(block_chunk_size: int) -> None:
|
||||||
|
"""Validate the common encoder and decoder block-chunk limit."""
|
||||||
|
if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int):
|
||||||
|
raise TypeError("block_chunk_size must be an integer")
|
||||||
|
if block_chunk_size <= 0:
|
||||||
|
raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}")
|
||||||
|
|
||||||
|
|
||||||
|
def validate_packed_weights(
|
||||||
|
packed_weights: torch.Tensor,
|
||||||
|
weight_shape: torch.Tensor,
|
||||||
|
*,
|
||||||
|
block_bytes: int,
|
||||||
|
format_name: str,
|
||||||
|
) -> tuple[int, ...]:
|
||||||
|
"""Validate a packed payload and return its logical shape."""
|
||||||
|
if (
|
||||||
|
packed_weights.dim() == 0
|
||||||
|
or packed_weights.dtype != torch.uint8
|
||||||
|
or packed_weights.shape[-1] != block_bytes
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
f"packed_weights must be uint8 with last dimension {block_bytes}, "
|
||||||
|
f"got {packed_weights.dtype} {tuple(packed_weights.shape)}"
|
||||||
|
)
|
||||||
|
integral_dtypes = {torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64}
|
||||||
|
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:
|
||||||
|
raise ValueError(f"invalid {format_name} logical weight shape: {shape}")
|
||||||
|
expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes
|
||||||
|
if packed_weights.numel() != expected_payload_values:
|
||||||
|
raise ValueError("packed_weights size does not match weight_shape")
|
||||||
|
return shape
|
||||||
@@ -0,0 +1,248 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""IQ1_S fake quantization and GGML-compatible block packing.
|
||||||
|
|
||||||
|
The encoder performs a single-pass squared-error grid search at a fixed,
|
||||||
|
empirically anchored super-block scale. It does not iteratively refine the
|
||||||
|
scale or apply importance weights. Every 256 logical values become one 50-byte
|
||||||
|
``block_iq1_s`` payload:
|
||||||
|
|
||||||
|
* bytes 0..1: little-endian FP16 super-block scale ``d``
|
||||||
|
* bytes 2..33: low eight bits of 32 codebook indices
|
||||||
|
* bytes 34..49: eight little-endian uint16 metadata words
|
||||||
|
|
||||||
|
Each metadata word describes four consecutive eight-value vectors. Bits 0..11
|
||||||
|
hold the three high index bits, bits 12..14 select one of eight local scales,
|
||||||
|
and bit 15 selects the shared -0.125 rather than +0.125 delta. The canonical
|
||||||
|
2048 x 8 ternary grid lives in :mod:`.codebooks`, carried from llama.cpp
|
||||||
|
``ggml-common.h`` revision
|
||||||
|
9b05354ec6fb58b4e665e9a39ebc40285c015638.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from ..extensions import get_cuda_ext_ggml
|
||||||
|
from .codebooks import iq1_s_grid_bytes
|
||||||
|
from .common import (
|
||||||
|
GGML_BLOCK_SIZE,
|
||||||
|
fake_quantize_with_cache,
|
||||||
|
narrow_to_float32,
|
||||||
|
validate_block_chunk_size,
|
||||||
|
validate_packed_weights,
|
||||||
|
validate_weight,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"IQ1_S_BLOCK_BYTES",
|
||||||
|
"IQ1_S_BLOCK_SIZE",
|
||||||
|
"IQ1_S_EFFECTIVE_BITS",
|
||||||
|
"dequantize_iq1_s",
|
||||||
|
"iq1_s_fake_quant",
|
||||||
|
"iq1_s_grid",
|
||||||
|
"quantize_iq1_s",
|
||||||
|
]
|
||||||
|
|
||||||
|
IQ1_S_BLOCK_SIZE = GGML_BLOCK_SIZE
|
||||||
|
IQ1_S_BLOCK_BYTES = 50
|
||||||
|
IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE
|
||||||
|
_IQ1_S_DELTA = 0.125
|
||||||
|
_IQ1_S_NATIVE_MAX = 16.875
|
||||||
|
_IQ1_S_SCALE_ANCHOR = 0.61
|
||||||
|
# At 1024 blocks, each largest IQ1_S search temporary is about 16 MiB in FP32.
|
||||||
|
_DEFAULT_BLOCK_CHUNK_SIZE = 1024
|
||||||
|
|
||||||
|
|
||||||
|
_GRID_CACHE: dict[torch.device, torch.Tensor] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor:
|
||||||
|
"""Return the canonical IQ1_S ternary grid as float32."""
|
||||||
|
resolved_device = torch.device(device or "cpu")
|
||||||
|
if resolved_device.type == "cuda" and resolved_device.index is None:
|
||||||
|
resolved_device = torch.device("cuda", torch.cuda.current_device())
|
||||||
|
if resolved_device not in _GRID_CACHE:
|
||||||
|
raw = torch.tensor(list(iq1_s_grid_bytes()), dtype=torch.uint8).view(torch.int8)
|
||||||
|
_GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to(
|
||||||
|
device=resolved_device, dtype=torch.float32
|
||||||
|
)
|
||||||
|
return _GRID_CACHE[resolved_device]
|
||||||
|
|
||||||
|
|
||||||
|
def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Encode a moderate-size batch of flattened 256-value blocks."""
|
||||||
|
x = narrow_to_float32(blocks)
|
||||||
|
block_count = x.shape[0]
|
||||||
|
vectors = x.reshape(block_count, 32, 8)
|
||||||
|
xnorm = vectors.square().sum(dim=-1)
|
||||||
|
xsum = vectors.sum(dim=-1)
|
||||||
|
|
||||||
|
amax = x.abs().amax(dim=1)
|
||||||
|
# The fixed-scale search favors a compressed super-block scale. This
|
||||||
|
# empirical anchor initializes d below the full-range value.
|
||||||
|
d = ((amax / _IQ1_S_NATIVE_MAX) * _IQ1_S_SCALE_ANCHOR).clamp(max=65504.0).to(torch.float16)
|
||||||
|
d_float = d.float()
|
||||||
|
|
||||||
|
best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device)
|
||||||
|
best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device)
|
||||||
|
grid_norm = grid.square().sum(dim=-1)
|
||||||
|
grid_sum = grid.sum(dim=-1)
|
||||||
|
|
||||||
|
# Tile the 2048-entry codebook to bound temporary memory. A strict update
|
||||||
|
# retains the lowest codebook index when two candidates have equal error.
|
||||||
|
for entry_start in range(0, 2048, 128):
|
||||||
|
grid_tile = grid[entry_start : entry_start + 128]
|
||||||
|
dot = torch.matmul(vectors, grid_tile.T)
|
||||||
|
tile_norm = grid_norm[entry_start : entry_start + 128].reshape(1, 1, -1)
|
||||||
|
tile_sum = grid_sum[entry_start : entry_start + 128].reshape(1, 1, -1)
|
||||||
|
|
||||||
|
for shift in range(2):
|
||||||
|
delta = -_IQ1_S_DELTA if shift else _IQ1_S_DELTA
|
||||||
|
shifted_dot = dot + delta * xsum.unsqueeze(-1)
|
||||||
|
shifted_norm = tile_norm + 2 * delta * tile_sum + 8 * delta * delta
|
||||||
|
for local in range(8):
|
||||||
|
choice = shift * 8 + local
|
||||||
|
scale = d_float.reshape(-1, 1, 1) * (2 * local + 1)
|
||||||
|
error = (
|
||||||
|
xnorm.unsqueeze(-1) - 2 * scale * shifted_dot + scale.square() * shifted_norm
|
||||||
|
).clamp_min_(0)
|
||||||
|
tile_error, tile_index = error.min(dim=-1)
|
||||||
|
replace = tile_error < best_error[:, :, choice]
|
||||||
|
best_error[:, :, choice] = torch.where(
|
||||||
|
replace, tile_error, best_error[:, :, choice]
|
||||||
|
)
|
||||||
|
best_entry[:, :, choice] = torch.where(
|
||||||
|
replace, tile_index + entry_start, best_entry[:, :, choice]
|
||||||
|
)
|
||||||
|
|
||||||
|
group_error = best_error.reshape(block_count, 8, 4, 16).sum(dim=2)
|
||||||
|
selected_choice = group_error.argmin(dim=-1)
|
||||||
|
vector_choice = selected_choice.repeat_interleave(4, dim=1)
|
||||||
|
selected_entry = best_entry.gather(2, vector_choice.unsqueeze(-1)).squeeze(-1)
|
||||||
|
selected_local = selected_choice & 0x7
|
||||||
|
selected_shift = selected_choice >> 3
|
||||||
|
|
||||||
|
high = (selected_entry >> 8).reshape(block_count, 8, 4)
|
||||||
|
qh = (
|
||||||
|
high[:, :, 0]
|
||||||
|
| (high[:, :, 1] << 3)
|
||||||
|
| (high[:, :, 2] << 6)
|
||||||
|
| (high[:, :, 3] << 9)
|
||||||
|
| (selected_local << 12)
|
||||||
|
| (selected_shift << 15)
|
||||||
|
)
|
||||||
|
|
||||||
|
packed = torch.empty((block_count, IQ1_S_BLOCK_BYTES), dtype=torch.uint8, device=x.device)
|
||||||
|
packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2)
|
||||||
|
packed[:, 2:34] = (selected_entry & 0xFF).to(torch.uint8)
|
||||||
|
packed[:, 34:50:2] = (qh & 0xFF).to(torch.uint8)
|
||||||
|
packed[:, 35:50:2] = (qh >> 8).to(torch.uint8)
|
||||||
|
return torch.where((d_float == 0).unsqueeze(1), 0, packed)
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def quantize_iq1_s(
|
||||||
|
weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
"""Pack a floating-point weight into GGML-compatible IQ1_S blocks.
|
||||||
|
|
||||||
|
Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 50]``
|
||||||
|
and ``[weight.ndim]``. The packed payload remains on the weight's device;
|
||||||
|
the logical-shape metadata is kept on CPU. Non-finite input elements are
|
||||||
|
treated as zero during packing.
|
||||||
|
"""
|
||||||
|
validate_weight(weight, "IQ1_S")
|
||||||
|
validate_block_chunk_size(block_chunk_size)
|
||||||
|
|
||||||
|
logical_shape = torch.tensor(weight.shape, dtype=torch.int64)
|
||||||
|
blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE)
|
||||||
|
grid = iq1_s_grid(weight.device)
|
||||||
|
if weight.is_cuda:
|
||||||
|
extension = get_cuda_ext_ggml()
|
||||||
|
if extension is not None:
|
||||||
|
packed = extension.iq1_s_pack(blocks, grid)
|
||||||
|
packed_shape = (
|
||||||
|
*weight.shape[:-1],
|
||||||
|
weight.shape[-1] // IQ1_S_BLOCK_SIZE,
|
||||||
|
IQ1_S_BLOCK_BYTES,
|
||||||
|
)
|
||||||
|
return packed.reshape(packed_shape), logical_shape
|
||||||
|
|
||||||
|
chunks = [
|
||||||
|
_encode_blocks(blocks[start : start + block_chunk_size], grid)
|
||||||
|
for start in range(0, blocks.shape[0], block_chunk_size)
|
||||||
|
]
|
||||||
|
packed_shape = (
|
||||||
|
*weight.shape[:-1],
|
||||||
|
weight.shape[-1] // IQ1_S_BLOCK_SIZE,
|
||||||
|
IQ1_S_BLOCK_BYTES,
|
||||||
|
)
|
||||||
|
return torch.cat(chunks).reshape(packed_shape), logical_shape
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def dequantize_iq1_s(
|
||||||
|
packed_weights: torch.Tensor,
|
||||||
|
weight_shape: torch.Tensor,
|
||||||
|
*,
|
||||||
|
dtype: torch.dtype = torch.bfloat16,
|
||||||
|
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Decode GGML-compatible IQ1_S payload bytes."""
|
||||||
|
shape = validate_packed_weights(
|
||||||
|
packed_weights, weight_shape, block_bytes=IQ1_S_BLOCK_BYTES, format_name="IQ1_S"
|
||||||
|
)
|
||||||
|
validate_block_chunk_size(block_chunk_size)
|
||||||
|
|
||||||
|
blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES)
|
||||||
|
shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device)
|
||||||
|
grid = iq1_s_grid(blocks.device)
|
||||||
|
decoded = torch.empty((blocks.shape[0], IQ1_S_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()
|
||||||
|
low = block_chunk[:, 2:34].to(torch.int64).reshape(-1, 8, 4)
|
||||||
|
qh = block_chunk[:, 34:50:2].to(torch.int64) | (
|
||||||
|
block_chunk[:, 35:50:2].to(torch.int64) << 8
|
||||||
|
)
|
||||||
|
high = (qh.unsqueeze(-1) >> shifts) & 0x7
|
||||||
|
entries = low | (high << 8)
|
||||||
|
local = (qh >> 12) & 0x7
|
||||||
|
delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA)
|
||||||
|
values = grid[entries] + delta.unsqueeze(-1).unsqueeze(-1)
|
||||||
|
scales = d.unsqueeze(-1) * (2 * local + 1).float()
|
||||||
|
chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1)
|
||||||
|
decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_S_BLOCK_SIZE)
|
||||||
|
return decoded.reshape(shape)
|
||||||
|
|
||||||
|
|
||||||
|
def iq1_s_fake_quant(
|
||||||
|
inputs: torch.Tensor,
|
||||||
|
quantizer,
|
||||||
|
*,
|
||||||
|
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""IQ1_S weight backend for TensorQuantizer, with pass-through backward."""
|
||||||
|
if getattr(quantizer, "num_bits", None) != "iq1_s":
|
||||||
|
raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'")
|
||||||
|
return fake_quantize_with_cache(
|
||||||
|
inputs,
|
||||||
|
quantizer,
|
||||||
|
format_name="iq1_s",
|
||||||
|
block_chunk_size=block_chunk_size,
|
||||||
|
quantize=quantize_iq1_s,
|
||||||
|
dequantize=dequantize_iq1_s,
|
||||||
|
)
|
||||||
@@ -0,0 +1,262 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
"""IQ2_XS fake quantization and GGML-compatible block packing.
|
||||||
|
|
||||||
|
The encoder performs a single-pass squared-error grid search at a fixed,
|
||||||
|
empirically anchored super-block scale. It does not iteratively refine the
|
||||||
|
scale or apply importance weights. Every 256 logical values become one 74-byte
|
||||||
|
block_iq2_xs payload:
|
||||||
|
|
||||||
|
* bytes 0..1: little-endian FP16 super-block scale d
|
||||||
|
* bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign)
|
||||||
|
* bytes 66..73: 16 four-bit local scales, two per byte
|
||||||
|
|
||||||
|
The canonical 512 x 8 magnitude grid lives in :mod:`.codebooks`, carried from
|
||||||
|
llama.cpp ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638.
|
||||||
|
The matching dequantization formula is in ggml-quants.c at the same revision:
|
||||||
|
https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-quants.c#L2516-L2538
|
||||||
|
"""
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from ..extensions import get_cuda_ext_ggml
|
||||||
|
from .codebooks import iq2_xs_grid_bytes
|
||||||
|
from .common import (
|
||||||
|
GGML_BLOCK_SIZE,
|
||||||
|
fake_quantize_with_cache,
|
||||||
|
narrow_to_float32,
|
||||||
|
validate_block_chunk_size,
|
||||||
|
validate_packed_weights,
|
||||||
|
validate_weight,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"IQ2_XS_BLOCK_BYTES",
|
||||||
|
"IQ2_XS_BLOCK_SIZE",
|
||||||
|
"IQ2_XS_EFFECTIVE_BITS",
|
||||||
|
"dequantize_iq2_xs",
|
||||||
|
"iq2_xs_fake_quant",
|
||||||
|
"iq2_xs_grid",
|
||||||
|
"quantize_iq2_xs",
|
||||||
|
]
|
||||||
|
|
||||||
|
IQ2_XS_BLOCK_SIZE = GGML_BLOCK_SIZE
|
||||||
|
IQ2_XS_BLOCK_BYTES = 74
|
||||||
|
IQ2_XS_EFFECTIVE_BITS = IQ2_XS_BLOCK_BYTES * 8 / IQ2_XS_BLOCK_SIZE
|
||||||
|
_IQ2_XS_NATIVE_MAX = 43 * 31 / 8
|
||||||
|
_IQ2_XS_SCALE_ANCHOR_MIN = 0.65
|
||||||
|
_IQ2_XS_SCALE_ANCHOR_MAX = 0.92
|
||||||
|
_IQ2_XS_PEAK_TO_RMS_TAPER = 0.035
|
||||||
|
# At 256 blocks, the largest IQ2_XS search temporary is about 16 MiB in FP32.
|
||||||
|
_DEFAULT_BLOCK_CHUNK_SIZE = 256
|
||||||
|
_SCALE_BLOCK_CHUNK_SIZE = 4096
|
||||||
|
|
||||||
|
|
||||||
|
_GRID_CACHE: dict[torch.device, torch.Tensor] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor:
|
||||||
|
"""Return the canonical IQ2_XS magnitude grid as float32."""
|
||||||
|
resolved_device = torch.device(device or "cpu")
|
||||||
|
if resolved_device.type == "cuda" and resolved_device.index is None:
|
||||||
|
resolved_device = torch.device("cuda", torch.cuda.current_device())
|
||||||
|
if resolved_device not in _GRID_CACHE:
|
||||||
|
values = torch.tensor(list(iq2_xs_grid_bytes()), dtype=torch.float32)
|
||||||
|
_GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device)
|
||||||
|
return _GRID_CACHE[resolved_device]
|
||||||
|
|
||||||
|
|
||||||
|
def _predict_iq2_xs_scales(blocks: torch.Tensor) -> torch.Tensor:
|
||||||
|
"""Predict one FP16 super-block scale for each flattened block."""
|
||||||
|
x = narrow_to_float32(blocks)
|
||||||
|
amax = x.abs().amax(dim=1)
|
||||||
|
rms = x.square().mean(dim=1).sqrt()
|
||||||
|
peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms))
|
||||||
|
# The fixed-scale search favors a compressed super-block scale. This
|
||||||
|
# empirical predictor tapers the anchor for outlier-heavy blocks.
|
||||||
|
anchor_ratio = (1.0 - _IQ2_XS_PEAK_TO_RMS_TAPER * peak_to_rms).clamp(
|
||||||
|
_IQ2_XS_SCALE_ANCHOR_MIN, _IQ2_XS_SCALE_ANCHOR_MAX
|
||||||
|
)
|
||||||
|
return ((amax / _IQ2_XS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16)
|
||||||
|
|
||||||
|
|
||||||
|
def _encode_blocks(
|
||||||
|
blocks: torch.Tensor, grid: torch.Tensor, scales: torch.Tensor | None = None
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Encode a moderate-size batch of flattened 256-value blocks."""
|
||||||
|
x = narrow_to_float32(blocks)
|
||||||
|
block_count = x.shape[0]
|
||||||
|
vectors = x.reshape(block_count, 32, 8)
|
||||||
|
magnitudes = vectors.abs()
|
||||||
|
negative = vectors < 0
|
||||||
|
odd_parity = negative.sum(dim=-1).remainder(2).bool()
|
||||||
|
|
||||||
|
d = _predict_iq2_xs_scales(x) if scales is None else scales
|
||||||
|
d_float = d.float()
|
||||||
|
|
||||||
|
xnorm = vectors.square().sum(dim=-1)
|
||||||
|
qnorm = grid.square().sum(dim=-1)
|
||||||
|
best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device)
|
||||||
|
best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device)
|
||||||
|
# Search the codebook in tiles to cap temporary memory. Strict comparison
|
||||||
|
# preserves the lowest grid index on equal error, matching the CUDA key.
|
||||||
|
for entry_start in range(0, 512, 64):
|
||||||
|
grid_tile = grid[entry_start : entry_start + 64]
|
||||||
|
products = magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8)
|
||||||
|
dot = products.sum(dim=-1)
|
||||||
|
dot = torch.where(odd_parity.unsqueeze(-1), dot - 2.0 * products.amin(dim=-1), dot)
|
||||||
|
tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1)
|
||||||
|
|
||||||
|
for local in range(16):
|
||||||
|
scale = d_float.reshape(-1, 1, 1) * ((2 * local + 1) / 8.0)
|
||||||
|
error = (
|
||||||
|
xnorm.unsqueeze(-1) - 2.0 * scale * dot + scale.square() * tile_qnorm
|
||||||
|
).clamp_min_(0)
|
||||||
|
tile_error, tile_index = error.min(dim=-1)
|
||||||
|
replace = tile_error < best_error[:, :, local]
|
||||||
|
best_error[:, :, local] = torch.where(replace, tile_error, best_error[:, :, local])
|
||||||
|
best_entry[:, :, local] = torch.where(
|
||||||
|
replace, tile_index + entry_start, best_entry[:, :, local]
|
||||||
|
)
|
||||||
|
|
||||||
|
group_error = best_error.reshape(block_count, 16, 2, 16).sum(dim=2)
|
||||||
|
selected_local = group_error.argmin(dim=-1)
|
||||||
|
vector_local = selected_local.repeat_interleave(2, dim=1)
|
||||||
|
selected_entry = best_entry.gather(2, vector_local.unsqueeze(-1)).squeeze(-1)
|
||||||
|
|
||||||
|
selected_grid = grid[selected_entry]
|
||||||
|
weakest_index = (magnitudes * selected_grid).argmin(dim=-1)
|
||||||
|
flip = torch.nn.functional.one_hot(weakest_index, num_classes=8).bool()
|
||||||
|
encoded_negative = negative ^ (flip & odd_parity.unsqueeze(-1))
|
||||||
|
sign_bits = torch.arange(8, dtype=torch.int64, device=x.device)
|
||||||
|
sign_mask = (encoded_negative.to(torch.int64) << sign_bits).sum(dim=-1)
|
||||||
|
|
||||||
|
codes = selected_entry | ((sign_mask & 0x7F) << 9)
|
||||||
|
packed = torch.empty((block_count, IQ2_XS_BLOCK_BYTES), dtype=torch.uint8, device=x.device)
|
||||||
|
packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2)
|
||||||
|
packed[:, 2:66:2] = (codes & 0xFF).to(torch.uint8)
|
||||||
|
packed[:, 3:66:2] = (codes >> 8).to(torch.uint8)
|
||||||
|
packed[:, 66:] = (selected_local[:, 0::2] | (selected_local[:, 1::2] << 4)).to(torch.uint8)
|
||||||
|
return torch.where((d_float == 0).unsqueeze(1), 0, packed)
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def quantize_iq2_xs(
|
||||||
|
weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
"""Pack a floating-point weight into GGML-compatible IQ2_XS blocks.
|
||||||
|
|
||||||
|
Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 74]``
|
||||||
|
and ``[weight.ndim]``. The packed payload remains on the weight's device;
|
||||||
|
the logical-shape metadata is kept on CPU. Non-finite input elements are
|
||||||
|
treated as zero during packing.
|
||||||
|
"""
|
||||||
|
validate_weight(weight, "IQ2_XS")
|
||||||
|
validate_block_chunk_size(block_chunk_size)
|
||||||
|
|
||||||
|
logical_shape = torch.tensor(weight.shape, dtype=torch.int64)
|
||||||
|
blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE)
|
||||||
|
grid = iq2_xs_grid(weight.device)
|
||||||
|
if weight.is_cuda:
|
||||||
|
extension = get_cuda_ext_ggml()
|
||||||
|
if extension is not None:
|
||||||
|
scale_chunks = [
|
||||||
|
_predict_iq2_xs_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE])
|
||||||
|
for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE)
|
||||||
|
]
|
||||||
|
packed = extension.iq2_xs_pack(blocks, grid, torch.cat(scale_chunks))
|
||||||
|
packed_shape = (
|
||||||
|
*weight.shape[:-1],
|
||||||
|
weight.shape[-1] // IQ2_XS_BLOCK_SIZE,
|
||||||
|
IQ2_XS_BLOCK_BYTES,
|
||||||
|
)
|
||||||
|
return packed.reshape(packed_shape), logical_shape
|
||||||
|
|
||||||
|
chunks = []
|
||||||
|
for start in range(0, blocks.shape[0], block_chunk_size):
|
||||||
|
block_chunk = blocks[start : start + block_chunk_size]
|
||||||
|
scales = _predict_iq2_xs_scales(block_chunk)
|
||||||
|
chunks.append(_encode_blocks(block_chunk, grid, scales))
|
||||||
|
packed_shape = (
|
||||||
|
*weight.shape[:-1],
|
||||||
|
weight.shape[-1] // IQ2_XS_BLOCK_SIZE,
|
||||||
|
IQ2_XS_BLOCK_BYTES,
|
||||||
|
)
|
||||||
|
return torch.cat(chunks).reshape(packed_shape), logical_shape
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def dequantize_iq2_xs(
|
||||||
|
packed_weights: torch.Tensor,
|
||||||
|
weight_shape: torch.Tensor,
|
||||||
|
*,
|
||||||
|
dtype: torch.dtype = torch.bfloat16,
|
||||||
|
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Decode GGML-compatible IQ2_XS payload bytes."""
|
||||||
|
shape = validate_packed_weights(
|
||||||
|
packed_weights, weight_shape, block_bytes=IQ2_XS_BLOCK_BYTES, format_name="IQ2_XS"
|
||||||
|
)
|
||||||
|
validate_block_chunk_size(block_chunk_size)
|
||||||
|
|
||||||
|
blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES)
|
||||||
|
bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device)
|
||||||
|
grid = iq2_xs_grid(blocks.device)
|
||||||
|
decoded = torch.empty((blocks.shape[0], IQ2_XS_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()
|
||||||
|
codes = block_chunk[:, 2:66:2].to(torch.int64) | (
|
||||||
|
block_chunk[:, 3:66:2].to(torch.int64) << 8
|
||||||
|
)
|
||||||
|
entries = codes & 0x1FF
|
||||||
|
sign_index = codes >> 9
|
||||||
|
parity = torch.zeros_like(sign_index)
|
||||||
|
for bit in range(7):
|
||||||
|
parity ^= (sign_index >> bit) & 1
|
||||||
|
sign_mask = sign_index | (parity << 7)
|
||||||
|
signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float()
|
||||||
|
|
||||||
|
scale_bytes = block_chunk[:, 66:].to(torch.int64)
|
||||||
|
local = torch.empty((block_chunk.shape[0], 16), dtype=torch.int64, device=blocks.device)
|
||||||
|
local[:, 0::2] = scale_bytes & 0x0F
|
||||||
|
local[:, 1::2] = scale_bytes >> 4
|
||||||
|
# Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8.
|
||||||
|
scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0
|
||||||
|
values = grid[entries] * signs
|
||||||
|
chunk_decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1)
|
||||||
|
decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XS_BLOCK_SIZE)
|
||||||
|
return decoded.reshape(shape)
|
||||||
|
|
||||||
|
|
||||||
|
def iq2_xs_fake_quant(
|
||||||
|
inputs: torch.Tensor,
|
||||||
|
quantizer,
|
||||||
|
*,
|
||||||
|
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""IQ2_XS weight backend for TensorQuantizer, with pass-through backward."""
|
||||||
|
if getattr(quantizer, "num_bits", None) != "iq2_xs":
|
||||||
|
raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'")
|
||||||
|
return fake_quantize_with_cache(
|
||||||
|
inputs,
|
||||||
|
quantizer,
|
||||||
|
format_name="iq2_xs",
|
||||||
|
block_chunk_size=block_chunk_size,
|
||||||
|
quantize=quantize_iq2_xs,
|
||||||
|
dequantize=dequantize_iq2_xs,
|
||||||
|
)
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
# 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 torch
|
||||||
|
|
||||||
|
import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
|
||||||
|
from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
|
||||||
|
from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, iq1_s_grid, quantize_iq1_s
|
||||||
|
|
||||||
|
|
||||||
|
def _extension():
|
||||||
|
extension = get_cuda_ext_ggml(raise_if_failed=True)
|
||||||
|
assert extension is not None
|
||||||
|
return extension
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch):
|
||||||
|
generator = torch.Generator(device="cuda").manual_seed(1234)
|
||||||
|
weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50)
|
||||||
|
packed_again = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50)
|
||||||
|
monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
reference, shape = quantize_iq1_s(weight)
|
||||||
|
reconstructed = dequantize_iq1_s(packed, shape)
|
||||||
|
|
||||||
|
assert packed.shape == (8, 1, 50)
|
||||||
|
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 < 0.25
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout():
|
||||||
|
weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16)
|
||||||
|
packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50)
|
||||||
|
shape = torch.tensor(weight.shape, device="cuda")
|
||||||
|
|
||||||
|
assert not packed.any()
|
||||||
|
assert torch.equal(dequantize_iq1_s(packed, shape), weight)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch):
|
||||||
|
weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16)
|
||||||
|
weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda")
|
||||||
|
|
||||||
|
packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50)
|
||||||
|
monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
reference, _ = quantize_iq1_s(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, reference)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch):
|
||||||
|
monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
generator = torch.Generator(device="cuda").manual_seed(1234)
|
||||||
|
weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq1_s(weight)
|
||||||
|
reconstructed = dequantize_iq1_s(packed, shape)
|
||||||
|
normalized_mse = (
|
||||||
|
reconstructed.float() - weight.float()
|
||||||
|
).square().mean() / weight.float().square().mean()
|
||||||
|
|
||||||
|
assert packed.shape == (2, 1, 50)
|
||||||
|
assert normalized_mse < 0.25
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_float64_matches_pytorch_encoder():
|
||||||
|
"""float64 weights inside the float32 range must pack identically on both paths."""
|
||||||
|
weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7))
|
||||||
|
|
||||||
|
reference, _ = quantize_iq1_s(weight)
|
||||||
|
packed, _ = quantize_iq1_s(weight.cuda())
|
||||||
|
|
||||||
|
assert torch.equal(reference, packed.cpu())
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_cuda_saturates_finite_values_above_the_float32_range():
|
||||||
|
"""The extension saturates such values rather than dropping them to zero.
|
||||||
|
|
||||||
|
Byte parity with the reference encoder is not asserted here: at these magnitudes the
|
||||||
|
squared-error objective overflows to infinity in float32, so every codebook candidate ties
|
||||||
|
and the two search implementations break that tie differently. The saturation policy is
|
||||||
|
what both paths must agree on.
|
||||||
|
"""
|
||||||
|
weight = torch.randn(1, 256, dtype=torch.float64, device="cuda")
|
||||||
|
weight[0, 7] = 1e100
|
||||||
|
saturated = weight.clone()
|
||||||
|
saturated[0, 7] = torch.finfo(torch.float32).max
|
||||||
|
zeroed = weight.clone()
|
||||||
|
zeroed[0, 7] = 0.0
|
||||||
|
|
||||||
|
packed, _ = quantize_iq1_s(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, quantize_iq1_s(saturated)[0])
|
||||||
|
assert not torch.equal(packed, quantize_iq1_s(zeroed)[0])
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
# 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 torch
|
||||||
|
|
||||||
|
import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
|
||||||
|
from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
|
||||||
|
from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, iq2_xs_grid, quantize_iq2_xs
|
||||||
|
|
||||||
|
|
||||||
|
def _extension():
|
||||||
|
extension = get_cuda_ext_ggml(raise_if_failed=True)
|
||||||
|
assert extension is not None
|
||||||
|
return extension
|
||||||
|
|
||||||
|
|
||||||
|
def _pack(weight):
|
||||||
|
blocks = weight.contiguous().reshape(-1, 256)
|
||||||
|
scales = iq2_xs_module._predict_iq2_xs_scales(blocks)
|
||||||
|
return _extension().iq2_xs_pack(weight, iq2_xs_grid("cuda"), scales)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch):
|
||||||
|
generator = torch.Generator(device="cuda").manual_seed(1234)
|
||||||
|
weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed = _pack(weight).reshape(8, 2, 74)
|
||||||
|
packed_again = _pack(weight).reshape(8, 2, 74)
|
||||||
|
monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
reference, shape = quantize_iq2_xs(weight)
|
||||||
|
reconstructed = dequantize_iq2_xs(packed, shape)
|
||||||
|
|
||||||
|
assert packed.shape == (8, 2, 74)
|
||||||
|
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 < 0.1
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout():
|
||||||
|
weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16)
|
||||||
|
packed = _pack(weight).reshape(1, 1, 74)
|
||||||
|
shape = torch.tensor(weight.shape, device="cuda")
|
||||||
|
|
||||||
|
assert not packed.any()
|
||||||
|
assert torch.equal(dequantize_iq2_xs(packed, shape), weight)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_underflowed_scale_has_canonical_zero_encoding():
|
||||||
|
weight = torch.full((1, 256), -1e-6, device="cuda", dtype=torch.bfloat16)
|
||||||
|
packed = _pack(weight).reshape(1, 1, 74)
|
||||||
|
|
||||||
|
assert not packed.any()
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch):
|
||||||
|
weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16)
|
||||||
|
weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda")
|
||||||
|
|
||||||
|
packed = _pack(weight).reshape(1, 1, 74)
|
||||||
|
monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
reference, _ = quantize_iq2_xs(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, reference)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch):
|
||||||
|
monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
|
||||||
|
generator = torch.Generator(device="cuda").manual_seed(1234)
|
||||||
|
weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq2_xs(weight)
|
||||||
|
reconstructed = dequantize_iq2_xs(packed, shape)
|
||||||
|
normalized_mse = (
|
||||||
|
reconstructed.float() - weight.float()
|
||||||
|
).square().mean() / weight.float().square().mean()
|
||||||
|
|
||||||
|
assert packed.shape == (2, 1, 74)
|
||||||
|
assert normalized_mse < 0.1
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_float64_matches_pytorch_encoder():
|
||||||
|
"""float64 weights inside the float32 range must pack identically on both paths."""
|
||||||
|
weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7))
|
||||||
|
|
||||||
|
reference, _ = quantize_iq2_xs(weight)
|
||||||
|
packed, _ = quantize_iq2_xs(weight.cuda())
|
||||||
|
|
||||||
|
assert torch.equal(reference, packed.cpu())
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_cuda_saturates_finite_values_above_the_float32_range():
|
||||||
|
"""The extension saturates such values rather than dropping them to zero.
|
||||||
|
|
||||||
|
Byte parity with the reference encoder is not asserted here: at these magnitudes the
|
||||||
|
squared-error objective overflows to infinity in float32, so every codebook candidate ties
|
||||||
|
and the two search implementations break that tie differently. The saturation policy is
|
||||||
|
what both paths must agree on.
|
||||||
|
"""
|
||||||
|
weight = torch.randn(1, 256, dtype=torch.float64, device="cuda")
|
||||||
|
weight[0, 7] = 1e100
|
||||||
|
saturated = weight.clone()
|
||||||
|
saturated[0, 7] = torch.finfo(torch.float32).max
|
||||||
|
zeroed = weight.clone()
|
||||||
|
zeroed[0, 7] = 0.0
|
||||||
|
|
||||||
|
packed, _ = quantize_iq2_xs(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, quantize_iq2_xs(saturated)[0])
|
||||||
|
assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0])
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
import modelopt.torch.quantization as mtq
|
||||||
|
import modelopt.torch.quantization.ggml.backend as backend_module
|
||||||
|
import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
|
||||||
|
import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
|
||||||
|
from modelopt.torch.quantization.ggml.backend import ggml_fake_quant
|
||||||
|
from modelopt.torch.quantization.ggml.common import narrow_to_float32
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
|
||||||
|
def test_ggml_backend_via_quantize(num_bits):
|
||||||
|
torch.manual_seed(1234)
|
||||||
|
model = torch.nn.Linear(256, 2, bias=False)
|
||||||
|
inputs = torch.randn(2, 256)
|
||||||
|
unquantized_output = model(inputs).detach()
|
||||||
|
config = {
|
||||||
|
"quant_cfg": [
|
||||||
|
{"quantizer_name": "*", "enable": False},
|
||||||
|
{
|
||||||
|
"quantizer_name": "*weight_quantizer",
|
||||||
|
"cfg": {"num_bits": num_bits, "backend": "ggml"},
|
||||||
|
"enable": True,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
"algorithm": "max",
|
||||||
|
}
|
||||||
|
|
||||||
|
mtq.quantize(model, config, forward_loop=lambda module: module(inputs))
|
||||||
|
output = model(inputs)
|
||||||
|
|
||||||
|
assert model.weight_quantizer.backend == "ggml"
|
||||||
|
assert model.weight_quantizer.num_bits == num_bits
|
||||||
|
assert output.shape == (2, 2)
|
||||||
|
assert torch.isfinite(output).all()
|
||||||
|
assert not torch.equal(output, unquantized_output)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ggml_backend_rejects_unknown_format():
|
||||||
|
with pytest.raises(ValueError, match="requires num_bits"):
|
||||||
|
ggml_fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits="unknown"))
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
def test_ggml_backend_forwards_block_chunk_size(monkeypatch):
|
||||||
|
received = {}
|
||||||
|
|
||||||
|
def fake_quant(inputs, _quantizer, *, block_chunk_size):
|
||||||
|
received["block_chunk_size"] = block_chunk_size
|
||||||
|
return inputs
|
||||||
|
|
||||||
|
monkeypatch.setattr(backend_module, "iq1_s_fake_quant", fake_quant)
|
||||||
|
inputs = torch.ones(1, 256)
|
||||||
|
quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"block_chunk_size": 17})
|
||||||
|
|
||||||
|
assert ggml_fake_quant(inputs, quantizer) is inputs
|
||||||
|
assert received == {"block_chunk_size": 17}
|
||||||
|
|
||||||
|
|
||||||
|
def test_ggml_backend_rejects_unknown_extra_arg():
|
||||||
|
quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"unknown": 1})
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Unsupported ggml backend_extra_args"):
|
||||||
|
ggml_fake_quant(torch.ones(1, 256), quantizer)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("num_bits", "module", "fake_quant_name", "quantize_name"),
|
||||||
|
[
|
||||||
|
("iq1_s", iq1_s_module, "iq1_s_fake_quant", "quantize_iq1_s"),
|
||||||
|
("iq2_xs", iq2_xs_module, "iq2_xs_fake_quant", "quantize_iq2_xs"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_ggml_backend_caches_packed_weight_and_invalidates_on_change(
|
||||||
|
monkeypatch, num_bits, module, fake_quant_name, quantize_name
|
||||||
|
):
|
||||||
|
weight = torch.randn(1, 256)
|
||||||
|
quantizer = SimpleNamespace(num_bits=num_bits, _quantizer_cache=None)
|
||||||
|
original_quantize = getattr(module, quantize_name)
|
||||||
|
call_count = 0
|
||||||
|
|
||||||
|
def counted_quantize(*args, **kwargs):
|
||||||
|
nonlocal call_count
|
||||||
|
call_count += 1
|
||||||
|
return original_quantize(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(module, quantize_name, counted_quantize)
|
||||||
|
fake_quant = getattr(module, fake_quant_name)
|
||||||
|
|
||||||
|
fake_quant(weight, quantizer, block_chunk_size=1)
|
||||||
|
fake_quant(weight, quantizer, block_chunk_size=1)
|
||||||
|
assert call_count == 1
|
||||||
|
|
||||||
|
fake_quant(weight, quantizer, block_chunk_size=2)
|
||||||
|
assert call_count == 2
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
weight.add_(0.01)
|
||||||
|
fake_quant(weight, quantizer, block_chunk_size=2)
|
||||||
|
assert call_count == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_narrow_to_float32_matches_the_cuda_load_float_policy():
|
||||||
|
"""Non-finite elements become zero; finite out-of-range elements saturate."""
|
||||||
|
largest = torch.finfo(torch.float32).max
|
||||||
|
values = torch.tensor(
|
||||||
|
[torch.nan, torch.inf, -torch.inf, 1e100, -1e100, 1.5], dtype=torch.float64
|
||||||
|
)
|
||||||
|
|
||||||
|
narrowed = narrow_to_float32(values)
|
||||||
|
|
||||||
|
assert narrowed.dtype is torch.float32
|
||||||
|
assert torch.equal(
|
||||||
|
narrowed, torch.tensor([0.0, 0.0, 0.0, largest, -largest, 1.5], dtype=torch.float32)
|
||||||
|
)
|
||||||
@@ -0,0 +1,183 @@
|
|||||||
|
# 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.iq1_s as iq1_s_module
|
||||||
|
from modelopt.torch.quantization.ggml.iq1_s import (
|
||||||
|
IQ1_S_BLOCK_BYTES,
|
||||||
|
dequantize_iq1_s,
|
||||||
|
iq1_s_fake_quant,
|
||||||
|
iq1_s_grid,
|
||||||
|
quantize_iq1_s,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_canonical_grid():
|
||||||
|
grid = iq1_s_grid()
|
||||||
|
|
||||||
|
assert grid.shape == (2048, 8)
|
||||||
|
assert grid.dtype == torch.float32
|
||||||
|
assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0}
|
||||||
|
assert grid[0].tolist() == [-1.0] * 8
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_grid_normalizes_unindexed_cuda_device(monkeypatch):
|
||||||
|
cached = torch.empty(0)
|
||||||
|
indexed_device = torch.device("cuda", 7)
|
||||||
|
monkeypatch.setattr(torch.cuda, "current_device", lambda: 7)
|
||||||
|
monkeypatch.setitem(iq1_s_module._GRID_CACHE, indexed_device, cached)
|
||||||
|
|
||||||
|
assert iq1_s_grid("cuda") is cached
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_zero_block_has_canonical_zero_encoding():
|
||||||
|
weight = torch.zeros((2, 256), dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq1_s(weight)
|
||||||
|
|
||||||
|
assert packed.shape == (2, 1, IQ1_S_BLOCK_BYTES)
|
||||||
|
assert packed.dtype == torch.uint8
|
||||||
|
assert not packed.any()
|
||||||
|
assert shape.tolist() == [2, 256]
|
||||||
|
assert torch.equal(dequantize_iq1_s(packed, shape), weight)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_dequantizes_ggml_metadata_bit_fields():
|
||||||
|
packed = torch.zeros((1, 1, 50), dtype=torch.uint8)
|
||||||
|
d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8)
|
||||||
|
packed[0, 0, :2] = d
|
||||||
|
entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64)
|
||||||
|
packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8)
|
||||||
|
qh = (
|
||||||
|
((entries[0] >> 8) & 7)
|
||||||
|
| (((entries[1] >> 8) & 7) << 3)
|
||||||
|
| (((entries[2] >> 8) & 7) << 6)
|
||||||
|
| (((entries[3] >> 8) & 7) << 9)
|
||||||
|
| (3 << 12)
|
||||||
|
| (1 << 15)
|
||||||
|
)
|
||||||
|
packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8)
|
||||||
|
packed[0, 0, 35] = (qh >> 8).to(torch.uint8)
|
||||||
|
|
||||||
|
decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32)
|
||||||
|
expected = (iq1_s_grid()[entries] - 0.125) * 14.0
|
||||||
|
|
||||||
|
assert torch.equal(decoded[0, :32].reshape(4, 8), expected)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_round_trip_and_payload_fields():
|
||||||
|
generator = torch.Generator().manual_seed(1234)
|
||||||
|
weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq1_s(weight, block_chunk_size=1)
|
||||||
|
reconstructed = dequantize_iq1_s(packed, shape, block_chunk_size=1)
|
||||||
|
default_reconstructed = dequantize_iq1_s(packed, shape)
|
||||||
|
|
||||||
|
assert packed.shape == (2, 1, 50)
|
||||||
|
assert reconstructed.shape == weight.shape
|
||||||
|
assert reconstructed.dtype == torch.bfloat16
|
||||||
|
assert torch.equal(reconstructed, default_reconstructed)
|
||||||
|
normalized_mse = (
|
||||||
|
reconstructed.float() - weight.float()
|
||||||
|
).square().mean() / weight.float().square().mean()
|
||||||
|
assert normalized_mse < 0.25
|
||||||
|
|
||||||
|
blocks = packed.reshape(-1, 50)
|
||||||
|
qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8)
|
||||||
|
assert torch.all(((qh >> 12) & 0x7) < 8)
|
||||||
|
assert torch.all((qh & 0xFFF) < 0x1000)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_search_is_independent_of_default_dtype():
|
||||||
|
generator = torch.Generator().manual_seed(0)
|
||||||
|
weight = torch.randn((8, 256), generator=generator, dtype=torch.float32)
|
||||||
|
expected, _ = quantize_iq1_s(weight)
|
||||||
|
|
||||||
|
default_dtype = torch.get_default_dtype()
|
||||||
|
try:
|
||||||
|
torch.set_default_dtype(torch.bfloat16)
|
||||||
|
actual, _ = quantize_iq1_s(weight)
|
||||||
|
finally:
|
||||||
|
torch.set_default_dtype(default_dtype)
|
||||||
|
|
||||||
|
assert torch.equal(actual, expected)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_requires_complete_last_dimension_blocks():
|
||||||
|
with pytest.raises(ValueError, match="last weight dimension"):
|
||||||
|
quantize_iq1_s(torch.ones(2, 257))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_treats_nonfinite_values_as_zero():
|
||||||
|
weight = torch.randn(1, 256)
|
||||||
|
weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf])
|
||||||
|
|
||||||
|
packed, _ = quantize_iq1_s(weight)
|
||||||
|
expected, _ = quantize_iq1_s(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0))
|
||||||
|
|
||||||
|
assert torch.equal(packed, expected)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"weight_shape",
|
||||||
|
[
|
||||||
|
torch.tensor(256),
|
||||||
|
torch.tensor([[1, 256]]),
|
||||||
|
torch.tensor([1.0, 256.0]),
|
||||||
|
torch.tensor([0, 256]),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_iq1_s_rejects_invalid_shape_metadata(weight_shape):
|
||||||
|
packed = torch.zeros((1, 1, 50), dtype=torch.uint8)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"):
|
||||||
|
dequantize_iq1_s(packed, weight_shape)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_rejects_scalar_packed_payload():
|
||||||
|
with pytest.raises(ValueError, match="packed_weights"):
|
||||||
|
dequantize_iq1_s(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_fake_quant_has_pass_through_gradient():
|
||||||
|
class Quantizer:
|
||||||
|
num_bits = "iq1_s"
|
||||||
|
|
||||||
|
weight = torch.randn(1, 256, requires_grad=True)
|
||||||
|
output = iq1_s_fake_quant(weight, Quantizer())
|
||||||
|
output.sum().backward()
|
||||||
|
|
||||||
|
assert torch.equal(weight.grad, torch.ones_like(weight))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq1_s_saturates_finite_values_above_the_float32_range():
|
||||||
|
"""float64 weights are accepted, so a finite value too large for float32 must saturate.
|
||||||
|
|
||||||
|
Converting before sanitizing would turn it into infinity and then zero, which silently
|
||||||
|
encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy.
|
||||||
|
"""
|
||||||
|
weight = torch.randn(1, 256, dtype=torch.float64)
|
||||||
|
weight[0, 7] = 1e100
|
||||||
|
saturated = weight.clone()
|
||||||
|
saturated[0, 7] = torch.finfo(torch.float32).max
|
||||||
|
zeroed = weight.clone()
|
||||||
|
zeroed[0, 7] = 0.0
|
||||||
|
|
||||||
|
packed, _ = quantize_iq1_s(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, quantize_iq1_s(saturated)[0])
|
||||||
|
assert not torch.equal(packed, quantize_iq1_s(zeroed)[0])
|
||||||
@@ -0,0 +1,168 @@
|
|||||||
|
# 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.iq2_xs as iq2_xs_module
|
||||||
|
from modelopt.torch.quantization.ggml.iq2_xs import (
|
||||||
|
IQ2_XS_BLOCK_BYTES,
|
||||||
|
dequantize_iq2_xs,
|
||||||
|
iq2_xs_fake_quant,
|
||||||
|
iq2_xs_grid,
|
||||||
|
quantize_iq2_xs,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_canonical_grid():
|
||||||
|
grid = iq2_xs_grid()
|
||||||
|
|
||||||
|
assert grid.shape == (512, 8)
|
||||||
|
assert grid.dtype == torch.float32
|
||||||
|
assert set(grid.unique().tolist()) == {8.0, 25.0, 43.0}
|
||||||
|
assert grid[0].tolist() == [8.0] * 8
|
||||||
|
assert grid[-1].tolist() == [43.0] * 8
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_grid_normalizes_unindexed_cuda_device(monkeypatch):
|
||||||
|
cached = torch.empty(0)
|
||||||
|
indexed_device = torch.device("cuda", 7)
|
||||||
|
monkeypatch.setattr(torch.cuda, "current_device", lambda: 7)
|
||||||
|
monkeypatch.setitem(iq2_xs_module._GRID_CACHE, indexed_device, cached)
|
||||||
|
|
||||||
|
assert iq2_xs_grid("cuda") is cached
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_zero_block_has_canonical_zero_encoding():
|
||||||
|
weight = torch.zeros((2, 256), dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq2_xs(weight)
|
||||||
|
|
||||||
|
assert packed.shape == (2, 1, IQ2_XS_BLOCK_BYTES)
|
||||||
|
assert packed.dtype == torch.uint8
|
||||||
|
assert not packed.any()
|
||||||
|
assert shape.tolist() == [2, 256]
|
||||||
|
assert torch.equal(dequantize_iq2_xs(packed, shape), weight)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_underflowed_scale_has_canonical_zero_encoding():
|
||||||
|
weight = torch.full((1, 256), -1e-6, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq2_xs(weight)
|
||||||
|
|
||||||
|
assert not packed.any()
|
||||||
|
assert torch.equal(dequantize_iq2_xs(packed, shape), torch.zeros_like(weight))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_round_trip_and_payload_fields():
|
||||||
|
generator = torch.Generator().manual_seed(1234)
|
||||||
|
weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16)
|
||||||
|
|
||||||
|
packed, shape = quantize_iq2_xs(weight, block_chunk_size=2)
|
||||||
|
reconstructed = dequantize_iq2_xs(packed, shape, block_chunk_size=1)
|
||||||
|
default_reconstructed = dequantize_iq2_xs(packed, shape)
|
||||||
|
|
||||||
|
assert packed.shape == (2, 2, 74)
|
||||||
|
assert reconstructed.shape == weight.shape
|
||||||
|
assert reconstructed.dtype == torch.bfloat16
|
||||||
|
assert torch.equal(reconstructed, default_reconstructed)
|
||||||
|
normalized_mse = (
|
||||||
|
reconstructed.float() - weight.float()
|
||||||
|
).square().mean() / weight.float().square().mean()
|
||||||
|
assert normalized_mse < 0.1
|
||||||
|
|
||||||
|
blocks = packed.reshape(-1, 74)
|
||||||
|
codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8)
|
||||||
|
assert torch.all((codes & 0x1FF) < 512)
|
||||||
|
assert torch.all((codes >> 9) < 128)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_dequantizes_pinned_scale_factor():
|
||||||
|
packed = torch.zeros((1, 1, 74), dtype=torch.uint8)
|
||||||
|
packed[0, 0, :2] = torch.tensor([1.0], dtype=torch.float16).view(torch.uint8)
|
||||||
|
packed[0, 0, 2:66:2] = 0xFF
|
||||||
|
packed[0, 0, 3:66:2] = 0x01
|
||||||
|
packed[0, 0, 66:] = 0xFF
|
||||||
|
|
||||||
|
decoded = dequantize_iq2_xs(packed, torch.tensor([1, 256]), dtype=torch.float32)
|
||||||
|
|
||||||
|
# Entry 511 contains eight 43s and local code 15 gives (2 * 15 + 1) / 8.
|
||||||
|
assert torch.equal(decoded, torch.full((1, 256), 43 * 31 / 8, dtype=torch.float32))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_requires_complete_last_dimension_blocks():
|
||||||
|
with pytest.raises(ValueError, match="last weight dimension"):
|
||||||
|
quantize_iq2_xs(torch.ones(2, 257))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_treats_nonfinite_values_as_zero():
|
||||||
|
weight = torch.randn(1, 256)
|
||||||
|
weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf])
|
||||||
|
|
||||||
|
packed, _ = quantize_iq2_xs(weight)
|
||||||
|
expected, _ = quantize_iq2_xs(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0))
|
||||||
|
|
||||||
|
assert torch.equal(packed, expected)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"weight_shape",
|
||||||
|
[
|
||||||
|
torch.tensor(256),
|
||||||
|
torch.tensor([[1, 256]]),
|
||||||
|
torch.tensor([1.0, 256.0]),
|
||||||
|
torch.tensor([0, 256]),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_iq2_xs_rejects_invalid_shape_metadata(weight_shape):
|
||||||
|
packed = torch.zeros((1, 1, 74), dtype=torch.uint8)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"):
|
||||||
|
dequantize_iq2_xs(packed, weight_shape)
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_rejects_scalar_packed_payload():
|
||||||
|
with pytest.raises(ValueError, match="packed_weights"):
|
||||||
|
dequantize_iq2_xs(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_fake_quant_has_pass_through_gradient():
|
||||||
|
class Quantizer:
|
||||||
|
num_bits = "iq2_xs"
|
||||||
|
|
||||||
|
weight = torch.randn(1, 256, requires_grad=True)
|
||||||
|
output = iq2_xs_fake_quant(weight, Quantizer())
|
||||||
|
output.sum().backward()
|
||||||
|
|
||||||
|
assert torch.equal(weight.grad, torch.ones_like(weight))
|
||||||
|
|
||||||
|
|
||||||
|
def test_iq2_xs_saturates_finite_values_above_the_float32_range():
|
||||||
|
"""float64 weights are accepted, so a finite value too large for float32 must saturate.
|
||||||
|
|
||||||
|
Converting before sanitizing would turn it into infinity and then zero, which silently
|
||||||
|
encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy.
|
||||||
|
"""
|
||||||
|
weight = torch.randn(1, 256, dtype=torch.float64)
|
||||||
|
weight[0, 7] = 1e100
|
||||||
|
saturated = weight.clone()
|
||||||
|
saturated[0, 7] = torch.finfo(torch.float32).max
|
||||||
|
zeroed = weight.clone()
|
||||||
|
zeroed[0, 7] = 0.0
|
||||||
|
|
||||||
|
packed, _ = quantize_iq2_xs(weight)
|
||||||
|
|
||||||
|
assert torch.equal(packed, quantize_iq2_xs(saturated)[0])
|
||||||
|
assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0])
|
||||||
Reference in New Issue
Block a user