mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
183 lines
8.0 KiB
Python
183 lines
8.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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.
|
|
|
|
"""Custom plugin base modules and utilities for quantization."""
|
|
|
|
import warnings
|
|
from collections.abc import Callable, Iterator
|
|
from functools import partial
|
|
from types import ModuleType
|
|
|
|
import torch
|
|
from packaging.version import Version
|
|
|
|
from modelopt import __version__
|
|
from modelopt.torch.utils.distributed import ParallelState
|
|
|
|
from ..nn import QuantModule, SequentialQuantizer, TensorQuantizer
|
|
from ..nn.modules.quant_linear import _QuantLinear
|
|
from ..utils import multi_context, replace_function
|
|
|
|
CUSTOM_MODEL_PLUGINS = set()
|
|
|
|
|
|
# TODO: This is a temporary solution
|
|
# In future implement a decorator to register methods updating QUANT_MODULE on the fly
|
|
def register_custom_model_plugins_on_the_fly(model):
|
|
"""Registers custom modules as QUANT_MODULE on the fly."""
|
|
for callback in CUSTOM_MODEL_PLUGINS:
|
|
callback(model)
|
|
|
|
|
|
class _QuantFunctionalMixin(QuantModule):
|
|
"""Mixin class for quantized functionals.
|
|
|
|
Often we need to replace a functional with a quantized version. This class provides a way to do that.
|
|
"""
|
|
|
|
# List of functionals to replace with quantized versions, e.g. [(package, func_name, quantized_func), ...]
|
|
_functionals_to_replace: list[tuple[ModuleType, str, Callable]] = []
|
|
|
|
@property
|
|
def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]:
|
|
return (
|
|
(package, func_name, quantized_func)
|
|
for package, func_name, quantized_func in self._functionals_to_replace
|
|
if hasattr(package, func_name)
|
|
)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
with multi_context(
|
|
*(
|
|
replace_function(package, func_name, quantized_func)
|
|
for package, func_name, quantized_func in self.functionals_to_replace
|
|
)
|
|
):
|
|
return super().forward(*args, **kwargs)
|
|
|
|
|
|
class _ParallelLinear(_QuantFunctionalMixin, QuantModule):
|
|
"""Quantized base class for ParallelLinear type classes.
|
|
|
|
For this type of modules, we need to quantize the inputs and weights just before calling the actual linear
|
|
functional. This is accomplished by replacing the linear functional with a custom one that quantizes the inputs
|
|
and weights before calling the original functional.
|
|
"""
|
|
|
|
# List of functionals to replace [(package, func_name), ...]
|
|
_functionals_to_replace: list[tuple[ModuleType, str]] = []
|
|
|
|
_parallel_state: ParallelState
|
|
_is_column_parallel = False
|
|
_is_row_parallel = False
|
|
_quantized_linear_fn: Callable = _QuantLinear.quantized_linear_fn
|
|
|
|
@property
|
|
def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]:
|
|
for package, func_name in self._functionals_to_replace:
|
|
if not hasattr(package, func_name):
|
|
continue
|
|
quantized_func = partial(
|
|
self.__class__._quantized_linear_fn,
|
|
package,
|
|
"_" + func_name,
|
|
self,
|
|
)
|
|
if hasattr(getattr(package, func_name), "__dict__"):
|
|
quantized_func.__dict__.update(getattr(package, func_name).__dict__)
|
|
yield package, func_name, quantized_func
|
|
|
|
def _setup(self):
|
|
self.input_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_input)
|
|
self.weight_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_weight)
|
|
self.output_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_output)
|
|
self.output_quantizer.disable()
|
|
|
|
def modelopt_post_restore(self, prefix: str = ""):
|
|
"""Post restore to correctly configure the TensorQuantizer states for MCore/distributed frameworks.
|
|
|
|
ModelOpt restores the TensorQuantizer states such as `_amax` and `_pre_quant_scale` to their
|
|
shape before saving. However this is not enough for MCore/distributed frameworks since the tensor parallelism
|
|
could change between saving and restoring. If the tensor parallelism changes, the shape of the quantizer
|
|
states also changes. So we need to re-calculate the quantizer states.
|
|
"""
|
|
from modelopt.torch.quantization.model_calib import max_calibrate
|
|
|
|
def _check_unsupported_states(quantizer: TensorQuantizer):
|
|
for k in quantizer.state_dict():
|
|
if k not in ["_amax", "_pre_quant_scale"]:
|
|
warnings.warn(
|
|
f"Restore of {k} for {prefix} is not supported. The restore of this layer might be "
|
|
f"incorrect. Please implement a custom restore for {k}."
|
|
)
|
|
|
|
def _has_state(quantizer, name):
|
|
# Handling for SequentialQuantizer
|
|
quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer
|
|
|
|
if self.is_version_less_than("0.29") and "dev" not in __version__:
|
|
# For backward compatibility, previously we used to save a boolean attribute "_has_amax"
|
|
# to indicate if the quantizer has amax.
|
|
return hasattr(quantizer, "_has" + name)
|
|
return hasattr(quantizer, name)
|
|
|
|
if self.weight is None:
|
|
return
|
|
|
|
for quantizer in [self.weight_quantizer, self.input_quantizer, self.output_quantizer]:
|
|
_check_unsupported_states(
|
|
quantizer if isinstance(quantizer, TensorQuantizer) else quantizer[0]
|
|
)
|
|
if _has_state(self.weight_quantizer, "_amax"):
|
|
self.weight_quantizer.reset_amax()
|
|
max_calibrate(self.weight_quantizer, lambda wq: wq(self.weight), distributed_sync=False)
|
|
if _has_state(self.input_quantizer, "_pre_quant_scale"):
|
|
if hasattr(self.input_quantizer, "_pre_quant_scale"):
|
|
delattr(self.input_quantizer, "_pre_quant_scale")
|
|
pqs = torch.zeros(
|
|
(self.weight.shape[1]), device=self.weight.device, dtype=self.weight.dtype
|
|
)
|
|
self.input_quantizer.register_buffer("_pre_quant_scale", pqs)
|
|
if _has_state(self.input_quantizer, "_amax"):
|
|
self.input_quantizer.reset_amax()
|
|
dummy_input = torch.ones(
|
|
(1, 1, self.weight.shape[1]), device=self.weight.device, dtype=self.weight.dtype
|
|
)
|
|
max_calibrate(self.input_quantizer, lambda iq: iq(dummy_input), distributed_sync=False)
|
|
if _has_state(self.output_quantizer, "_amax"):
|
|
self.output_quantizer.reset_amax()
|
|
dummy_input = torch.ones(
|
|
(1, 1, self.weight.shape[0]), device=self.weight.device, dtype=self.weight.dtype
|
|
)
|
|
max_calibrate(self.output_quantizer, lambda oq: oq(dummy_input), distributed_sync=False)
|
|
# If there are any other states, lets move them to the correct device
|
|
super().modelopt_post_restore(prefix=prefix)
|
|
|
|
def is_version_less_than(self, version: str) -> bool:
|
|
self_version = (
|
|
Version(self.mopt_ckpt_versn)
|
|
if self.mopt_ckpt_versn is not None
|
|
else Version(__version__)
|
|
)
|
|
|
|
# version in NeMo container is 0.0.0 if installed from source without history
|
|
if self_version < Version(version) and self_version != Version("0.0.0"):
|
|
warnings.warn(
|
|
f"Checkpoint version {self_version} is less than {version}. "
|
|
"Please re-save model to avoid this warning."
|
|
)
|
|
return True
|
|
return False
|