nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/quantization/calib/max.py

@@ -0,0 +1,110 @@
+# 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.
+
+"""Calibrator that returns the absolute max of all collected tensors."""
+
+import torch
+
+from .. import utils as quant_utils
+from .calibrator import _Calibrator
+
+__all__ = ["MaxCalibrator"]
+
+
+class MaxCalibrator(_Calibrator):
+    """Max calibrator, tracks the maximum value globally.
+
+    Args:
+        calib_desc: A MaxCalibDescriptor.
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see :class:`QuantizerAttributeConfig <..config.QuantizerAttributeConfig>`.
+        unsigned: A boolean. using unsigned quantization.
+
+    Readonly Properties:
+        amaxs: A list of amax. Numpy array is saved as it is likely to be used for some plot.
+    """
+
+    def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False):
+        """Initialize."""
+        super().__init__(num_bits, axis, unsigned)
+        self._track_amax = track_amax
+        if self._track_amax:
+            self._amaxs = []  # shall we have a better name?
+        self._calib_amax = None
+
+    @property
+    def amaxs(self):
+        """Returns the list of amax`s collected so far."""
+        return self._amaxs
+
+    @torch.no_grad()
+    def collect(self, x):
+        """Tracks the absolute max of all tensors.
+
+        Args:
+            x: A tensor
+
+        Raises:
+            RuntimeError: If amax shape changes
+        """
+        # Swap axis to reduce.
+        reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(x, self._axis)
+        local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach()
+        # meta device is used for initialization
+        if x.device.type == "meta":
+            self._calib_amax = local_amax
+            return
+        assert torch.all(local_amax >= 0), (
+            "detected negative values after abs, could be torch or cuda bug"
+        )
+        assert not torch.any(torch.isinf(local_amax)), (
+            f"detected inf values in amax. inf in original tensor: {torch.any(torch.isinf(x))}"
+        )
+        assert not torch.any(torch.isnan(local_amax)), (
+            f"detected nan values in amax. nan in original tensor: {torch.any(torch.isnan(x))}"
+        )
+        if self._calib_amax is None:
+            self._calib_amax = local_amax
+        else:
+            if local_amax.shape != self._calib_amax.shape:
+                raise RuntimeError("amax shape changed!")
+            self._calib_amax = torch.max(self._calib_amax, local_amax)
+
+        if self._track_amax:
+            self._amaxs.append(local_amax.cpu().numpy())
+
+    def reset(self):
+        """Reset the collected absolute max."""
+        self._calib_amax = None
+
+    def compute_amax(self):
+        """Return the absolute max of all tensors collected."""
+        return self._calib_amax
+
+    def __str__(self):
+        s = "MaxCalibrator("
+        s += "track_amax={_track_amax}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+    def __repr__(self):
+        s = "MaxCalibrator("
+        s += super().__repr__()
+        s += " calib_amax={_calib_amax}"
+        s += " track_amax={_track_amax}"
+        if self._track_amax:
+            s += " amaxs={_amaxs}"
+        s += ")"
+        return s.format(**self.__dict__)

modelopt/torch/quantization/compress.py

@@ -0,0 +1,227 @@
+# 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.
+
+"""Module for compressing the model weights after quantization."""
+
+__all__ = ["compress"]
+
+import fnmatch
+import warnings
+
+import torch
+import torch.nn as nn
+
+from modelopt.torch.opt import apply_mode
+from modelopt.torch.opt.conversion import ModelLikeModule, ModeloptStateManager
+from modelopt.torch.opt.dynamic import _DMRegistryCls
+from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
+
+from .backends.gemm_registry import disable_real_quant_gemm, enable_real_quant_gemm
+from .config import CompressCfgType, CompressConfig
+from .conversion import _replace_quant_module, set_quantizer_attribute
+from .nn.modules.quant_linear import RealQuantLinear
+from .qtensor import QTensorWrapper, pack_real_quantize_weight
+from .utils import is_quantized_linear
+
+try:
+    from .plugins.megatron import (
+        _MegatronColumnParallelLinear,
+        _MegatronRowParallelLinear,
+        _RealQuantMegatronColumnParallelLinear,
+        _RealQuantMegatronRowParallelLinear,
+    )
+
+    mcore_available = True
+except ImportError:
+    mcore_available = False
+
+RealQuantModuleRegistry = _DMRegistryCls("RealQuant")
+
+
+def compress_convert(
+    model,
+    config: CompressConfig,
+    skip_real_quantize_weight: bool = False,
+) -> ConvertReturnType:
+    """Compress entry point.
+
+    This function converts the model to a real quantized model.
+
+    Args:
+        model: The model to compress.
+        config: The compression configuration.
+        skip_real_quantize_weight: Whether to skip the real quantize step. Currently, it is
+            only set to True in the Megatron restore path to unify the restore behavior regardless
+            of whether the model is initialized on meta device or not.
+
+    Returns:
+        The compressed model.
+    """
+    for _, module in model.named_modules():
+        if is_quantized_linear(module) and type(module) not in RealQuantModuleRegistry:
+            class_to_register = RealQuantLinear
+            if mcore_available:
+                if issubclass(type(module), _MegatronRowParallelLinear):
+                    class_to_register = _RealQuantMegatronRowParallelLinear
+                elif issubclass(type(module), _MegatronColumnParallelLinear):
+                    class_to_register = _RealQuantMegatronColumnParallelLinear
+            RealQuantModuleRegistry.register({type(module): module.__class__.__name__})(
+                class_to_register
+            )
+    # Convert QuantLinear to RealQuantLinear
+    _replace_quant_module(
+        model, version=ModeloptStateManager(model).state_version, registry=RealQuantModuleRegistry
+    )
+
+    compress_cfg = config.compress
+    if "default" in compress_cfg and isinstance(compress_cfg["default"], bool):
+        set_quantizer_attribute(
+            model, "*weight_quantizer*", {"fake_quant": not compress_cfg["default"]}
+        )
+
+    for pattern, to_compress in compress_cfg.items():
+        if pattern == "default":
+            continue
+        if isinstance(to_compress, bool):
+
+            def filter_func(name):
+                return fnmatch.fnmatch(name, pattern) and "weight_quantizer" in name
+
+            set_quantizer_attribute(model, filter_func, {"fake_quant": not to_compress})
+        else:
+            raise ValueError(
+                f"Invalid compression configuration: {to_compress}, expected a boolean as value."
+            )
+    # If real quant quantizer is present, real quantize the weights.
+    if not skip_real_quantize_weight:
+        pack_real_quantize_weight(model)
+
+    def _has_qtensorwrapper(module):
+        if hasattr(module, "weight") and isinstance(module.weight, QTensorWrapper):
+            return True
+        return any(_has_qtensorwrapper(submodule) for _, submodule in module.named_children())
+
+    if _has_qtensorwrapper(model):
+        warnings.warn(
+            "Real quantization has been applied to the model. This feature is still "
+            "experimental, and some functionalities may not be supported. For example, "
+            "converting the model back to its original state or saving and restoring "
+            "the quantized model may not be available."
+        )
+
+    # Turn on real quant gemm after compression
+    if config.quant_gemm:
+        enable_real_quant_gemm(model)
+    else:
+        disable_real_quant_gemm(model)
+
+    metadata = {}
+    update_compress_metadata(model, config, metadata)
+
+    return model, metadata
+
+
+def compress_restore(
+    model: ModelLikeModule, config: CompressConfig, metadata: MetadataDict
+) -> nn.Module:
+    """Restore the model from the compressed state.
+
+    Note:
+        When restoring Megatron distributed checkpoint, real_quantizer_state and q_tensor_state
+        have been removed from metadata and stored as a part of QuantModule.extra_state.
+        Restoring happens in set_extra_state when load_state_dict is called. We also skip real
+        quantize weight (skip_real_quantize_weight). All these steps are
+        delayed. For details, see plugins.megatron.quant_module_set_extra_state.
+    """
+    # Compress with dummy weights
+    model, _ = compress_convert(
+        model,
+        config,
+        skip_real_quantize_weight=("q_tensor_state" not in metadata),
+    )
+    # restore scale state in weight quantizer
+    if "real_quantizer_state" in metadata:
+        for name, module in model.named_modules():
+            if isinstance(module, RealQuantLinear) and name in metadata["real_quantizer_state"]:
+                if not metadata["real_quantizer_state"][name].items():
+                    raise ValueError(f"Cannot find real quantizer state for {name}")
+                module.weight_quantizer.set_from_modelopt_state(
+                    metadata["real_quantizer_state"][name]
+                )
+
+    # restore real quant tensor states
+    if "q_tensor_state" in metadata:
+        for name, module in model.named_modules():
+            if isinstance(module, RealQuantLinear) and name in metadata["q_tensor_state"]:
+                module._parameters["weight"] = QTensorWrapper(
+                    qtensor=torch.empty(
+                        metadata["q_tensor_state"][name]["quantized_data.shape"],
+                        dtype=metadata["q_tensor_state"][name]["quantized_data.dtype"],
+                        device=module.weight.device,
+                    ),
+                    metadata=metadata["q_tensor_state"][name]["metadata"],
+                )
+    return model
+
+
+def update_compress_metadata(model: nn.Module, config: CompressConfig, metadata: MetadataDict):
+    # save scales state in weight quantizer
+    real_quantizer_state = {}
+    for name, module in model.named_modules():
+        if isinstance(module, RealQuantLinear):
+            real_quantizer_state[name] = module.weight_quantizer.get_modelopt_state()
+
+    # real quant tensor states
+    q_tensor_state = {}
+    for name, module in model.named_modules():
+        if isinstance(module, RealQuantLinear) and isinstance(module.weight, QTensorWrapper):
+            q_tensor_state[name] = module.weight.get_state()
+
+    metadata["real_quantizer_state"] = real_quantizer_state
+    metadata["q_tensor_state"] = q_tensor_state
+
+
+def compress(model, config: CompressCfgType = None):
+    """Compress model weights of quantized model.
+
+    This function compresses weights in layers that have an enabled `weight_quantizer` with
+    a supported quantization format. The compression is controlled by a pattern-based configuration.
+
+    Args:
+        model: The quantized model to compress.
+        config: Dictionary mapping layer patterns to boolean compression flags.
+            If ``None``, defaults to ``{"default": True}`` which compresses all supported layers.
+
+            Example configuration::
+
+                {
+                    "*.mlp.fc1*": False,  # Skip compression for fc1 layers
+                    "default": True,  # Compress all other layers
+                }
+
+            Note: Each configuration except "default" is applied sequentially; therefore the later
+            configurations will override the previous ones if the same layer is matched.
+
+
+    Note: This function modifies the input model in-place.
+    """
+    if config is None:
+        config = CompressConfig()
+    apply_mode(model, [("real_quantize", config)])
+
+
+def is_real_quantized(model):
+    """Check if the model is real quantized."""
+    return any(isinstance(_module, RealQuantLinear) for _module in model.modules())