compressed-tensors 0.3.0__py3-none-any.whl

compressed_tensors/config/sparse_bitmask.py

@@ -0,0 +1,36 @@
+# 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 typing import Optional
+
+from compressed_tensors.config import CompressionConfig, CompressionFormat
+
+
+__all__ = ["BitmaskConfig"]
+
+
+@CompressionConfig.register(name=CompressionFormat.sparse_bitmask.value)
+class BitmaskConfig(CompressionConfig):
+    """
+    Configuration for storing a sparse model using
+    bitmask compression
+
+    :param global_sparsity: average sparsity of the entire model
+    :param sparsity_structure: structure of the sparsity, such as
+    "unstructured", "2:4", "8:16" etc
+    """
+
+    format: str = CompressionFormat.sparse_bitmask.value
+    global_sparsity: Optional[float] = 0.0
+    sparsity_structure: Optional[str] = "unstructured"

compressed_tensors/quantization/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+# flake8: noqa
+# isort: skip_file
+
+from .quant_args import *
+from .quant_config import *
+from .quant_scheme import *
+from .lifecycle import *

compressed_tensors/quantization/lifecycle/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+
+# flake8: noqa
+# isort: skip_file
+
+from .calibration import *
+from .forward import *
+from .frozen import *
+from .initialize import *
+from .apply import *

compressed_tensors/quantization/lifecycle/apply.py

@@ -0,0 +1,173 @@
+# 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.
+
+import re
+from collections import OrderedDict
+from typing import Dict, Iterable, Optional
+
+from compressed_tensors.quantization.lifecycle.calibration import (
+    set_module_for_calibration,
+)
+from compressed_tensors.quantization.lifecycle.frozen import freeze_module_quantization
+from compressed_tensors.quantization.lifecycle.initialize import (
+    initialize_module_for_quantization,
+)
+from compressed_tensors.quantization.quant_config import (
+    QuantizationConfig,
+    QuantizationStatus,
+)
+from compressed_tensors.quantization.utils import iter_named_leaf_modules
+from compressed_tensors.utils.safetensors_load import get_safetensors_folder
+from torch.nn import Module
+
+
+__all__ = [
+    "load_pretrained_quantization",
+    "apply_quantization_config",
+    "apply_quantization_status",
+]
+
+from compressed_tensors.quantization.utils.helpers import is_module_quantized
+from compressed_tensors.utils.safetensors_load import get_quantization_state_dict
+
+
+def load_pretrained_quantization(model: Module, model_name_or_path: str):
+    """
+    Loads the quantization parameters (scale and zero point) from model_name_or_path to
+    a model that has already been initialized with a quantization config
+
+    :param model: model to load pretrained quantization parameters to
+    :param model_name_or_path: Hugging Face stub or local folder containing a quantized
+    model, which is used to load quantization parameters
+    """
+    model_path = get_safetensors_folder(model_name_or_path)
+    state_dict = get_quantization_state_dict(model_path)
+
+    for name, submodule in iter_named_leaf_modules(model):
+        if not is_module_quantized(submodule):
+            continue
+        if submodule.quantization_scheme.weights is not None:
+            base_name = "weight"
+            _load_quant_args_from_state_dict(
+                base_name=base_name,
+                module_name=name,
+                module=submodule,
+                state_dict=state_dict,
+            )
+        if submodule.quantization_scheme.input_activations is not None:
+            base_name = "input"
+            _load_quant_args_from_state_dict(
+                base_name=base_name,
+                module_name=name,
+                module=submodule,
+                state_dict=state_dict,
+            )
+        if submodule.quantization_scheme.output_activations is not None:
+            base_name = "output"
+            _load_quant_args_from_state_dict(
+                base_name=base_name,
+                module_name=name,
+                module=submodule,
+                state_dict=state_dict,
+            )
+
+
+def apply_quantization_config(model: Module, config: QuantizationConfig):
+    """
+    Initializes the model for quantization in-place based on the given config
+
+    :param model: model to apply quantization config to
+    :param config: quantization config
+    """
+    # build mapping of targets to schemes for easier matching
+    # use ordered dict to preserve target ordering in config
+    target_to_scheme = OrderedDict()
+    for scheme in config.config_groups.values():
+        for target in scheme.targets:
+            target_to_scheme[target] = scheme
+
+    # mark appropriate layers for quantization by setting their quantization schemes
+    for name, submodule in iter_named_leaf_modules(model):
+        if _find_first_name_or_class_match(name, submodule, config.ignore):
+            continue  # layer matches ignore list, continue
+        target = _find_first_name_or_class_match(name, submodule, target_to_scheme)
+        if target is not None:
+            # target matched - add layer and scheme to target list
+            submodule.quantization_scheme = target_to_scheme[target]
+
+    # apply current quantization status across all targeted layers
+    apply_quantization_status(model, config.quantization_status)
+
+
+def apply_quantization_status(model: Module, status: QuantizationStatus):
+    """
+    Applies in place the quantization lifecycle up to the given status
+
+    :param model: model to apply quantization to
+    :param status: status to update the module to
+    """
+    if status >= QuantizationStatus.INITIALIZED:
+        model.apply(initialize_module_for_quantization)
+    if status >= QuantizationStatus.CALIBRATION:
+        model.apply(set_module_for_calibration)
+    if status >= QuantizationStatus.FROZEN:
+        model.apply(freeze_module_quantization)
+
+
+def _find_first_name_or_class_match(
+    name: str,
+    module: Module,
+    targets: Iterable[str],
+) -> Optional[str]:
+    # first element of targets that matches the given name
+    # if no name matches returns first target that matches the class name
+    # returns None otherwise
+    return _find_first_match(name, targets) or _find_first_match(
+        module.__class__.__name__, targets
+    )
+
+
+def _find_first_match(value: str, targets: Iterable[str]) -> Optional[str]:
+    # returns first element of target that matches value either
+    # exactly or as a regex after 're:'
+    for target in targets:
+        if target.startswith("re:"):
+            pattern = target[3:]
+            if re.match(pattern, value):
+                return target
+        elif target == value:
+            return target
+    return None
+
+
+def _load_quant_args_from_state_dict(
+    base_name: str, module_name: str, module: Module, state_dict: Dict
+):
+    """
+    Loads scale and zero point from a state_dict into the specified module
+
+    :param base_name: quantization target, one of: weights, input_activations or
+    output_activations
+    :param module_name: pytorch module name to look up in state_dict
+    :module: pytorch module associated with module_name
+    :state_dict: state_dict to search for matching quantization parameters
+    """
+    scale_name = f"{base_name}_scale"
+    zp_name = f"{base_name}_zero_point"
+    device = next(module.parameters()).device
+
+    scale = getattr(module, scale_name)
+    zp = getattr(module, zp_name)
+    scale.data = state_dict[f"{module_name}.{scale_name}"].to(device)
+    zp.data = state_dict[f"{module_name}.{zp_name}"].to(device)

compressed_tensors/quantization/lifecycle/calibration.py

@@ -0,0 +1,51 @@
+# 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.
+
+
+import logging
+
+from compressed_tensors.quantization.quant_config import QuantizationStatus
+from torch.nn import Module
+
+
+__all__ = [
+    "set_module_for_calibration",
+]
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def set_module_for_calibration(module: Module):
+    """
+    marks a layer as ready for calibration which activates observers
+    to update scales and zero points on each forward pass
+
+    apply to full model with `model.apply(set_module_for_calibration)`
+
+    :param module: module to set for calibration
+    """
+    if not getattr(module, "quantization_scheme", None):
+        # no quantization scheme nothing to do
+        return
+    status = getattr(module, "quantization_status", None)
+    if not status or status != QuantizationStatus.INITIALIZED:
+        raise _LOGGER.warning(
+            f"Attempting set module with status {status} to calibration mode. "
+            f"but status is not {QuantizationStatus.INITIALIZED} - you may "
+            "be calibrating an uninitialized module which may fail or attempting "
+            "to re-calibrate a frozen module"
+        )
+
+    module.quantization_status = QuantizationStatus.CALIBRATION

compressed_tensors/quantization/lifecycle/forward.py

@@ -0,0 +1,136 @@
+# 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 functools import wraps
+
+import torch
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from torch.nn import Module
+
+
+__all__ = ["wrap_module_forward_quantized"]
+
+
+@torch.no_grad()
+def quantize(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+    q_min: torch.Tensor,
+    q_max: torch.Tensor,
+) -> torch.Tensor:
+    return torch.clamp(
+        torch.round(
+            x / scale + zero_point,
+        ),
+        q_min,
+        q_max,
+    )
+
+
+@torch.no_grad()
+def dequantize(
+    x_q: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+) -> torch.Tensor:
+    return (x_q - zero_point) * scale
+
+
+@torch.no_grad()
+def fake_quantize(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+    args: QuantizationArgs,
+) -> torch.Tensor:
+    bit_range = 2**args.num_bits
+    max_q = torch.tensor(bit_range / 2 - 1, device=x.device)
+    min_q = torch.tensor(-bit_range / 2, device=x.device)
+    Q = torch.zeros_like(x)
+    Q = quantize(x, scale, zero_point, min_q, max_q)
+    return dequantize(Q, scale, zero_point)
+
+
+def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
+    # expects a module already initialized and injected with the parameters in
+    # initialize_module_for_quantization
+    forward_func_orig = module.forward.__func__
+
+    @wraps(forward_func_orig)  # ensures docstring, names, etc are propagated
+    def wrapped_forward(self, *args, **kwargs):
+        input_ = args[0]
+
+        if scheme.input_activations is not None:
+            # calibrate and (fake) quantize input activations when applicable
+            input_ = _maybe_calibrate_or_quantize(
+                module, input_, "input", scheme.input_activations
+            )
+
+        if scheme.weights is not None:
+            # calibrate and (fake) quantize weights when applicable
+            unquantized_weight = self.weight.data.clone()
+            self.weight.data = _maybe_calibrate_or_quantize(
+                module, self.weight, "weight", scheme.weights
+            )
+
+        # perform wrapped forward call
+        output = forward_func_orig.__get__(module, module.__class__)(
+            input_, *args[1:], **kwargs
+        )
+
+        if scheme.output_activations is not None:
+            # calibrate and (fake) quantize output activations when applicable
+            output = _maybe_calibrate_or_quantize(
+                module, output, "output", scheme.output_activations
+            )
+
+        # restore back to unquantized_value
+        if scheme.weights is not None:
+            self.weight.data = unquantized_weight
+
+        return output
+
+    # bind wrapped forward to module class so reference to `self` is correct
+    bound_wrapped_forward = wrapped_forward.__get__(module, module.__class__)
+    # set forward to wrapped forward
+    setattr(module, "forward", bound_wrapped_forward)
+
+
+def _maybe_calibrate_or_quantize(
+    module: Module, value: Module, base_name: str, args: "QuantizationArgs"
+) -> torch.Tensor:
+    # only run quantized for the included stages
+    if module.quantization_status not in {
+        QuantizationStatus.CALIBRATION,
+        QuantizationStatus.FROZEN,
+    }:
+        return value
+
+    device = next(module.parameters()).device
+    scale = getattr(module, f"{base_name}_scale")
+    zero_point = getattr(module, f"{base_name}_zero_point")
+
+    if module.quantization_status == QuantizationStatus.CALIBRATION:
+        # get observer and get new quant params from observation
+        observer = getattr(module, f"{base_name}_observer")
+        updated_scale, updated_zero_point = observer(value)
+
+        # update scale and zero point
+        scale.data = updated_scale.to(device)
+        zero_point.data = updated_zero_point.to(device)
+
+    return fake_quantize(value, scale, zero_point, args)

compressed_tensors/quantization/lifecycle/frozen.py

@@ -0,0 +1,46 @@
+# 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 compressed_tensors.quantization.quant_config import QuantizationStatus
+from torch.nn import Module
+
+
+__all__ = [
+    "freeze_module_quantization",
+]
+
+
+def freeze_module_quantization(module: Module):
+    """
+    deletes observers so static quantization is completed.
+
+    apply to full model with `model.apply(freeze_module_quantization)`
+
+    :param module: module to freeze quantization for
+    """
+    if not getattr(module, "quantization_scheme", None):
+        # no quantization scheme nothing to do
+        return
+
+    # delete observers from module
+    observer_names = []
+    for submodule_name, _ in module.named_modules():
+        if "." not in submodule_name and submodule_name.endswith("_observer"):
+            # delete any observers that belong directly to this module
+            observer_names.append(submodule_name)
+    for observer_name in observer_names:
+        delattr(module, observer_name)
+
+    module.quantization_status = QuantizationStatus.FROZEN

compressed_tensors/quantization/lifecycle/initialize.py

@@ -0,0 +1,96 @@
+# 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.
+
+
+import logging
+from typing import Optional
+
+import torch
+from compressed_tensors.quantization.lifecycle.forward import (
+    wrap_module_forward_quantized,
+)
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from torch.nn import Module, Parameter
+
+
+__all__ = [
+    "initialize_module_for_quantization",
+]
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def initialize_module_for_quantization(
+    module: Module,
+    scheme: Optional[QuantizationScheme] = None,
+):
+    """
+    attaches appropriate scales, zero points, and observers to a layer
+    given its target quantization scheme
+
+    apply to full model with `model.apply(initialize_module_for_quantization)`
+
+    :param module: module to set for calibration
+    :param scheme: scheme to use for quantization. if None is provided,
+        will attempt to use scheme stored in the module under `quantization_scheme`,
+        if not provided, the layer will be skipped
+    """
+    scheme = scheme or getattr(module, "quantization_scheme", None)
+    if scheme is None:
+        # no scheme passed and layer not targeted for quantization - skip
+        return
+
+    if scheme.input_activations is not None:
+        _initialize_scale_zero_point_observer(module, "input", scheme.input_activations)
+    if scheme.weights is not None:
+        if hasattr(module, "weight"):
+            _initialize_scale_zero_point_observer(module, "weight", scheme.weights)
+        else:
+            _LOGGER.warning(
+                f"module type {type(module)} targeted for weight quantization but "
+                "has no attribute weight, skipping weight quantization "
+                f"for {type(module)}"
+            )
+    if scheme.output_activations is not None:
+        _initialize_scale_zero_point_observer(
+            module, "output", scheme.output_activations
+        )
+
+    module.quantization_scheme = scheme
+    module.quantization_status = QuantizationStatus.INITIALIZED
+
+    # wrap forward call of module to perform quantized actions based on calltime status
+    wrap_module_forward_quantized(module, scheme)
+
+
+def _initialize_scale_zero_point_observer(
+    module: Module, base_name: str, quantization_args: QuantizationArgs
+):
+    device = next(module.parameters()).device
+
+    # initializes empty scale and zero point parameters for the module
+    init_scale = Parameter(torch.empty(0, device=device), requires_grad=False)
+    module.register_parameter(f"{base_name}_scale", init_scale)
+
+    init_zero_point = Parameter(
+        torch.empty(0, device=device, dtype=int), requires_grad=False
+    )
+    module.register_parameter(f"{base_name}_zero_point", init_zero_point)
+
+    # initialize observer module and attach as submodule
+    observer = quantization_args.get_observer()
+    module.register_module(f"{base_name}_observer", observer)

compressed_tensors/quantization/observers/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+# flake8: noqa
+# isort: skip_file
+
+from .helpers import *
+from .base import *
+from .memoryless import *
+from .min_max import *

compressed_tensors/quantization/observers/base.py

@@ -0,0 +1,69 @@
+# 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 typing import Optional, Tuple
+
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.registry.registry import RegistryMixin
+from torch import FloatTensor, IntTensor, Tensor
+from torch.nn import Module
+
+
+__all__ = ["Observer"]
+
+
+class Observer(Module, RegistryMixin):
+    """
+    Base Observer class to be subclassed for specific implementation.
+    Subclasses should override `calculate_qparams` to return a scale, zero_point
+    pair
+    """
+
+    def __init__(self, quantization_args: QuantizationArgs):
+        self.quantization_args: QuantizationArgs = quantization_args
+        super().__init__()
+        self._scale = None
+        self._zero_point = None
+
+    def forward(self, observed: Tensor) -> Tuple[FloatTensor, IntTensor]:
+        """
+        maps directly to get_qparams
+        :param observed: optional observed tensor to calculate quantization parameters
+            from
+        :return: tuple of scale and zero point based on last observed value
+        """
+        return self.get_qparams(observed=observed)
+
+    def calculate_qparams(self, observed: Tensor) -> Tuple[FloatTensor, IntTensor]:
+        """
+        :param observed: observed tensor to calculate quantization parameters for
+        :return: tuple of scale and zero point derived from the observed tensor
+        """
+        raise NotImplementedError(f"{self.__class__} must implement calculate_qparams")
+
+    def get_qparams(
+        self, observed: Optional[Tensor] = None
+    ) -> Tuple[FloatTensor, IntTensor]:
+        """
+        Convenience function to wrap overwritten calculate_qparams
+        adds support to make observed tensor optional and support for tracking latest
+        calculated scale and zero point
+        :param observed: optional observed tensor to calculate quantization parameters
+            from
+        :return: tuple of scale and zero point based on last observed value
+        """
+        if observed is not None:
+            # re-calcualte scale and zero point, update the stored value
+            self._scale, self._zero_point = self.calculate_qparams(observed)
+        return self._scale, self._zero_point