mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
9e3d555aa1
## 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>
36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 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.
|
|
|
|
"""Quantization package."""
|
|
|
|
from importlib import import_module as _import_module
|
|
|
|
# Initialize mode and plugins
|
|
from . import mode, plugins, utils
|
|
|
|
# Add methods to mtq namespace
|
|
from .compress import *
|
|
from .config import *
|
|
from .conversion import *
|
|
from .model_quant import *
|
|
from .nn.modules.quant_module import QuantModuleRegistry
|
|
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
|