| # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. | |
| # | |
| # 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 enum import Enum | |
| from compressed_tensors.utils import delete_offload_parameter | |
| from torch.nn import Module | |
| __all__ = ["QuantizationMetadata", "KVCacheScaleType"] | |
| class KVCacheScaleType(Enum): | |
| KEY = "k_scale" | |
| VALUE = "v_scale" | |
| class QuantizationMetadata: | |
| """ | |
| Container class for metadata related to quantization | |
| """ | |
| def all_qparam_names(): | |
| """ | |
| All quantization parameter names that might be registered | |
| onto a module during lifecycle (excluding serialized parameters) | |
| """ | |
| return [KVCacheScaleType.KEY.value, KVCacheScaleType.VALUE.value] + [ | |
| f"{base_name}_{suffix}" | |
| for base_name in ("input", "weight", "output") | |
| for suffix in ( | |
| "global_scale", | |
| "scale", | |
| "zero_point", | |
| "g_idx", | |
| ) | |
| ] | |
| def clear_all_qparams(cls, module: Module): | |
| """ | |
| Remove all parameters related to quantization that might have | |
| been registered onto a module previously in lifecycle (excluding | |
| serialized parameters) | |
| :param module: Module to clear | |
| """ | |
| for key in cls.all_qparam_names(): | |
| if hasattr(module, key): | |
| delete_offload_parameter(module, key) | |