compressed-tensors-nightly 0.3.3.20240514__py3-none-any.whl

compressed_tensors/quantization/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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
+from .helpers import *

compressed_tensors/quantization/utils/helpers.py

@@ -0,0 +1,131 @@
+# 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 Tuple
+
+import torch
+from compressed_tensors.quantization.observers.base import Observer
+from torch.nn import Module
+from tqdm import tqdm
+
+
+__all__ = [
+    "is_module_quantized",
+    "is_model_quantized",
+    "iter_named_leaf_modules",
+    "module_type",
+    "calculate_compression_ratio",
+]
+
+
+def is_module_quantized(module: Module) -> bool:
+    """
+    Check if a module is quantized, based on the existence of a non-empty quantization
+    scheme
+
+    :param module: pytorch module to check
+    :return: True if module is quantized, False otherwise
+    """
+    if not hasattr(module, "quantization_scheme"):
+        return False
+
+    if module.quantization_scheme.weights is not None:
+        return True
+
+    if module.quantization_scheme.input_activations is not None:
+        return True
+
+    if module.quantization_scheme.output_activations is not None:
+        return True
+
+    return False
+
+
+def is_model_quantized(model: Module) -> bool:
+    """
+    Check if any modules in a model are quantized, based on the existence of a non-empty
+    quantization scheme in at least one module
+
+    :param model: pytorch model
+    :return: True if model is quantized, False otherwise
+    """
+
+    for _, submodule in iter_named_leaf_modules(model):
+        if is_module_quantized(submodule):
+            return True
+
+    return False
+
+
+def module_type(module: Module) -> str:
+    """
+    Gets a string representation of a module type
+
+    :module: pytorch module to get type of
+    :return: module type as a string
+    """
+    return type(module).__name__
+
+
+def iter_named_leaf_modules(model: Module) -> Tuple[str, Module]:
+    """
+    Yields modules that do not have any submodules except observers. The observers
+    themselves are not yielded
+
+    :param model: model to get leaf modules of
+    :returns: generator tuple of (name, leaf_submodule)
+    """
+    for name, submodule in model.named_modules():
+        children = list(submodule.children())
+        if len(children) == 0 and not isinstance(submodule, Observer):
+            yield name, submodule
+        else:
+            has_non_observer_children = False
+            for child in children:
+                if not isinstance(child, Observer):
+                    has_non_observer_children = True
+
+            if not has_non_observer_children:
+                yield name, submodule
+
+
+def calculate_compression_ratio(model: Module) -> float:
+    """
+    Calculates the quantization compression ratio of a pytorch model, based on the
+    number of bits needed to represent the total weights in compressed form. Does not
+    take into account activation quantizatons.
+
+    :param model: pytorch module to calculate compression ratio for
+    :return: compression ratio of the whole model
+    """
+    total_compressed = 0.0
+    total_uncompressed = 0.0
+    for name, submodule in tqdm(
+        iter_named_leaf_modules(model),
+        desc="Calculating quantization compression ratio",
+    ):
+        for parameter in model.parameters():
+            try:
+                uncompressed_bits = torch.finfo(parameter.dtype).bits
+            except TypeError:
+                uncompressed_bits = torch.iinfo(parameter.dtype).bits
+            compressed_bits = uncompressed_bits
+            if is_module_quantized(submodule):
+                compressed_bits = submodule.quantization_scheme.weights.num_bits
+
+            num_weights = parameter.numel()
+            total_compressed += compressed_bits * num_weights
+            total_uncompressed += uncompressed_bits * num_weights
+
+    return total_uncompressed / total_compressed

compressed_tensors/registry/__init__.py

@@ -0,0 +1,17 @@
+# flake8: noqa
+
+# 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 .registry import *

compressed_tensors/registry/registry.py

@@ -0,0 +1,360 @@
+# 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.
+
+"""
+Universal registry to support registration and loading of child classes and plugins
+of neuralmagic utilities
+"""
+
+import importlib
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Type, Union
+
+
+__all__ = [
+    "RegistryMixin",
+    "register",
+    "get_from_registry",
+    "registered_names",
+    "registered_aliases",
+    "standardize_lookup_name",
+]
+
+
+_ALIAS_REGISTRY: Dict[Type, Dict[str, str]] = defaultdict(dict)
+_REGISTRY: Dict[Type, Dict[str, Any]] = defaultdict(dict)
+
+
+def standardize_lookup_name(name: str) -> str:
+    """
+    Standardize the given name for lookup in the registry.
+    This will replace all underscores and spaces with hyphens and
+    convert the name to lowercase.
+
+    example:
+    ```
+    standardize_lookup_name("Foo_bar baz") == "foo-bar-baz"
+    ```
+
+    :param name: name to standardize
+    :return: standardized name
+    """
+    return name.replace("_", "-").replace(" ", "-").lower()
+
+
+def standardize_alias_name(
+    name: Union[None, str, List[str]]
+) -> Union[None, str, List[str]]:
+    if name is None:
+        return None
+    elif isinstance(name, str):
+        return standardize_lookup_name(name)
+    else:  # isinstance(name, list)
+        return [standardize_lookup_name(n) for n in name]
+
+
+class RegistryMixin:
+    """
+    Universal registry to support registration and loading of child classes and plugins
+    of neuralmagic utilities.
+
+    Classes that require a registry or plugins may add the `RegistryMixin` and use
+    `register` and `load` as the main entrypoints for adding new implementations and
+    loading requested values from its registry.
+
+    If a class should only have its child classes in its registry, the class should
+    set the static attribute `registry_requires_subclass` to True
+
+    example
+    ```python
+    class Dataset(RegistryMixin):
+        pass
+
+
+    # register with default name
+    @Dataset.register()
+    class ImageNetDataset(Dataset):
+        pass
+
+    # load as "ImageNetDataset"
+    imagenet = Dataset.load("ImageNetDataset")
+
+    # register with custom name
+    @Dataset.register(name="cifar-dataset")
+    class Cifar(Dataset):
+        pass
+
+    Note: the name will be standardized for lookup in the registry.
+    For example, if a class is registered as "cifar_dataset" or
+    "cifar dataset", it will be stored as "cifar-dataset". The user
+    will be able to load the class with any of the three name variants.
+
+    # register with multiple aliases
+    @Dataset.register(alias=["cifar-10-dataset", "cifar_100_dataset"])
+    class Cifar(Dataset):
+        pass
+
+    # load as "cifar-dataset"
+    cifar = Dataset.load_from_registry("cifar-dataset")
+
+    # load from custom file that implements a dataset
+    mnist = Dataset.load_from_registry("/path/to/mnnist_dataset.py:MnistDataset")
+    ```
+    """
+
+    # set to True in child class to add check that registered/retrieved values
+    # implement the class it is registered to
+    registry_requires_subclass: bool = False
+
+    @classmethod
+    def register(
+        cls, name: Optional[str] = None, alias: Union[List[str], str, None] = None
+    ):
+        """
+        Decorator for registering a value (ie class or function) wrapped by this
+        decorator to the base class (class that .register is called from)
+
+        :param name: name or list of names to register the wrapped value as,
+            defaults to value.__name__
+        :param alias: alias or list of aliases to register the wrapped value as,
+            defaults to None
+        :return: register decorator
+        """
+
+        def decorator(value: Any):
+            cls.register_value(value, name=name, alias=alias)
+            return value
+
+        return decorator
+
+    @classmethod
+    def register_value(
+        cls, value: Any, name: str, alias: Union[str, List[str], None] = None
+    ):
+        """
+        Registers the given value to the class `.register_value` is called from
+        :param value: value to register
+        :param name: name to register the wrapped value as,
+            defaults to value.__name__
+        :param alias: alias or list of aliases to register the wrapped value as,
+            defaults to None
+        """
+        register(
+            parent_class=cls,
+            value=value,
+            name=name,
+            alias=alias,
+            require_subclass=cls.registry_requires_subclass,
+        )
+
+    @classmethod
+    def load_from_registry(cls, name: str, **constructor_kwargs) -> object:
+        """
+        :param name: name of registered class to load
+        :param constructor_kwargs: arguments to pass to the constructor retrieved
+            from the registry
+        :return: loaded object registered to this class under the given name,
+            constructed with the given kwargs. Raises error if the name is
+            not found in the registry
+        """
+        constructor = cls.get_value_from_registry(name=name)
+        return constructor(**constructor_kwargs)
+
+    @classmethod
+    def get_value_from_registry(cls, name: str):
+        """
+        :param name: name to retrieve from the registry
+        :return: value from retrieved the registry for the given name, raises
+            error if not found
+        """
+        return get_from_registry(
+            parent_class=cls,
+            name=name,
+            require_subclass=cls.registry_requires_subclass,
+        )
+
+    @classmethod
+    def registered_names(cls) -> List[str]:
+        """
+        :return: list of all names registered to this class
+        """
+        return registered_names(cls)
+
+    @classmethod
+    def registered_aliases(cls) -> List[str]:
+        """
+        :return: list of all aliases registered to this class
+        """
+        return registered_aliases(cls)
+
+
+def register(
+    parent_class: Type,
+    value: Any,
+    name: Optional[str] = None,
+    alias: Union[List[str], str, None] = None,
+    require_subclass: bool = False,
+):
+    """
+    :param parent_class: class to register the name under
+    :param value: the value to register
+    :param name: name to register the wrapped value as, defaults to value.__name__
+    :param alias: alias or list of aliases to register the wrapped value as,
+        defaults to None
+    :param require_subclass: require that value is a subclass of the class this
+        method is called from
+    """
+    if name is None:
+        # default name
+        name = value.__name__
+
+    name = standardize_lookup_name(name)
+    alias = standardize_alias_name(alias)
+    register_alias(name=name, alias=alias, parent_class=parent_class)
+
+    if require_subclass:
+        _validate_subclass(parent_class, value)
+
+    if name in _REGISTRY[parent_class]:
+        # name already exists - raise error if two different values are attempting
+        # to share the same name
+        registered_value = _REGISTRY[parent_class][name]
+        if registered_value is not value:
+            raise RuntimeError(
+                f"Attempting to register name {name} as {value} "
+                f"however {name} has already been registered as {registered_value}"
+            )
+    else:
+        _REGISTRY[parent_class][name] = value
+
+
+def get_from_registry(
+    parent_class: Type, name: str, require_subclass: bool = False
+) -> Any:
+    """
+    :param parent_class: class that the name is registered under
+    :param name: name to retrieve from the registry of the class
+    :param require_subclass: require that value is a subclass of the class this
+        method is called from
+    :return: value from retrieved the registry for the given name, raises
+        error if not found
+    """
+    name = standardize_lookup_name(name)
+
+    if ":" in name:
+        # user specifying specific module to load and value to import
+        module_path, value_name = name.split(":")
+        retrieved_value = _import_and_get_value_from_module(module_path, value_name)
+    else:
+        # look up name in alias registry
+        name = _ALIAS_REGISTRY[parent_class].get(name)
+        # look up name in registry
+        retrieved_value = _REGISTRY[parent_class].get(name)
+        if retrieved_value is None:
+            raise KeyError(
+                f"Unable to find {name} registered under type {parent_class}.\n"
+                f"Registered values for {parent_class}: "
+                f"{registered_names(parent_class)}\n"
+                f"Registered aliases for {parent_class}: "
+                f"{registered_aliases(parent_class)}"
+            )
+
+    if require_subclass:
+        _validate_subclass(parent_class, retrieved_value)
+
+    return retrieved_value
+
+
+def registered_names(parent_class: Type) -> List[str]:
+    """
+    :param parent_class: class to look up the registry of
+    :return: all names registered to the given class
+    """
+    return list(_REGISTRY[parent_class].keys())
+
+
+def registered_aliases(parent_class: Type) -> List[str]:
+    """
+    :param parent_class: class to look up the registry of
+    :return: all aliases registered to the given class
+    """
+    registered_aliases_plus_names = list(_ALIAS_REGISTRY[parent_class].keys())
+    registered_aliases = list(
+        set(registered_aliases_plus_names) - set(registered_names(parent_class))
+    )
+    return registered_aliases
+
+
+def register_alias(
+    name: str, parent_class: Type, alias: Union[str, List[str], None] = None
+):
+    """
+    Updates the mapping from the alias(es) to the given name.
+    If the alias is None, the name is used as the alias.
+    ```
+
+    :param name: name that the alias refers to
+    :param parent_class: class that the name is registered under
+    :param alias: single alias or list of aliases that
+        refer to the name, defaults to None
+    """
+    if alias is not None:
+        alias = alias if isinstance(alias, list) else [alias]
+    else:
+        alias = []
+
+    if name in alias:
+        raise KeyError(
+            f"Attempting to register alias {name}, "
+            f"that is identical to the standardized name: {name}."
+        )
+    alias.append(name)
+
+    for alias_name in alias:
+        if alias_name in _ALIAS_REGISTRY[parent_class]:
+            raise KeyError(
+                f"Attempting to register alias {alias_name} as {name} "
+                f"however {alias_name} has already been registered as "
+                f"{_ALIAS_REGISTRY[alias_name]}"
+            )
+        _ALIAS_REGISTRY[parent_class][alias_name] = name
+
+
+def _import_and_get_value_from_module(module_path: str, value_name: str) -> Any:
+    # import the given module path and try to get the value_name if it is included
+    # in the module
+
+    # load module
+    spec = importlib.util.spec_from_file_location(
+        f"plugin_module_for_{value_name}", module_path
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    # get value from module
+    value = getattr(module, value_name, None)
+
+    if not value:
+        raise RuntimeError(
+            f"Unable to find attribute {value_name} in module {module_path}"
+        )
+    return value
+
+
+def _validate_subclass(parent_class: Type, child_class: Type):
+    if not issubclass(child_class, parent_class):
+        raise ValueError(
+            f"class {child_class} is not a subclass of the class it is "
+            f"registered for: {parent_class}."
+        )

compressed_tensors/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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
+
+from .safetensors_load import *

compressed_tensors/utils/helpers.py

@@ -0,0 +1,45 @@
+# 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.base import SPARSITY_CONFIG_NAME
+from compressed_tensors.compressors import ModelCompressor
+from compressed_tensors.config import CompressionConfig
+from transformers import AutoConfig
+
+
+__all__ = ["infer_compressor_from_model_config"]
+
+
+def infer_compressor_from_model_config(
+    pretrained_model_name_or_path: str,
+) -> Optional[ModelCompressor]:
+    """
+    Given a path to a model config, extract a sparsity config if it exists and return
+    the associated ModelCompressor
+
+    :param pretrained_model_name_or_path: path to model config on disk or HF hub
+    :return: matching compressor if config contains a sparsity config
+    """
+    config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+    sparsity_config = getattr(config, SPARSITY_CONFIG_NAME, None)
+    if sparsity_config is None:
+        return None
+
+    format = sparsity_config.get("format")
+    sparsity_config = CompressionConfig.load_from_registry(format, **sparsity_config)
+    compressor = ModelCompressor.load_from_registry(format, config=sparsity_config)
+    return compressor