compressed-tensors 0.10.1a20250605__tar.gz → 0.10.2__tar.gz

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/.github/workflows/report.yml

@@ -120,7 +120,7 @@ jobs:
               shell: bash
 
             - name: report to reportportal
-              uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.15.0
+              uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.22.0
               with:
                 droute_username: ${{ secrets.DROUTE_USERNAME }}
                 droute_password: ${{ secrets.DROUTE_PASSWORD }}

{compressed_tensors-0.10.1a20250605/src/compressed_tensors.egg-info → compressed_tensors-0.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.1a20250605
+Version: 0.10.2
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/setup.py

@@ -113,5 +113,6 @@ setup(
     extras_require=_setup_extras(),
     install_requires=_setup_install_requires(),
     package_dir={"": "src"},
+    package_data={"": ["transform/utils/hadamards.safetensors"]},
     packages=_setup_packages(),
 )

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -68,6 +68,10 @@ from transformers import AutoConfig
 from transformers.file_utils import CONFIG_NAME
 
 
+if TYPE_CHECKING:
+    from compressed_tensors.compressors import BaseQuantizationCompressor
+
+
 __all__ = ["ModelCompressor", "map_module_to_scheme"]
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -257,7 +261,9 @@ class ModelCompressor:
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
         self.sparsity_compressor = None
-        self.quantization_compressor = None
+        self.quantization_compressor: Optional[
+            Union[BaseQuantizationCompressor, DenseCompressor]
+        ] = None
 
         if sparsity_config is not None:
             self.sparsity_compressor = BaseCompressor.load_from_registry(

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/src/compressed_tensors/compressors/sparse_compressors/dense.py

@@ -12,13 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Generator, Tuple
+from typing import TYPE_CHECKING, Dict, Generator, Tuple
 
+import torch
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.config import CompressionFormat
 from torch import Tensor
 
 
+if TYPE_CHECKING:
+    from compressed_tensors.quantization import QuantizationScheme
+
+
 @BaseCompressor.register(name=CompressionFormat.dense.value)
 class DenseCompressor(BaseCompressor):
     """
@@ -47,3 +52,16 @@ class DenseCompressor(BaseCompressor):
     ) -> Generator[Tuple[str, Dict[str, Tensor]], None, None]:
         for key, value in state_dict.items():
             yield key, value
+
+    def decompress_module_from_state_dict(
+        self,
+        prefix: str,
+        state_dict: Dict[str, torch.Tensor],
+        scheme: "QuantizationScheme",
+    ) -> Dict[str, torch.Tensor]:
+        """
+        This function is implemented as a workaround because of how
+        `ModelCompressor.quantization_compressor` can be set to either
+        an instance of `BaseQuantizationCompressor` or `DenseCompressor`.
+        """
+        return state_dict.copy()

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/src/compressed_tensors/quantization/lifecycle/apply.py

@@ -183,9 +183,7 @@ def apply_quantization_config(
                         replace_module(model, name, compressed_linear)
 
             # target matched - add layer and scheme to target list
-            submodule.quantization_scheme = _scheme_from_targets(
-                target_to_scheme, targets, name
-            )
+            submodule.quantization_scheme = scheme
 
             names_to_scheme[name] = submodule.quantization_scheme
 

{compressed_tensors-0.10.1a20250605 → compressed_tensors-0.10.2}/src/compressed_tensors/transform/__init__.py

@@ -18,3 +18,8 @@
 from .transform_args import *
 from .transform_scheme import *
 from .transform_config import *
+
+from .factory.base import *
+from .factory.hadamard import *
+from .factory.matrix_multiply import *
+from .factory.random_hadamard import *

compressed_tensors-0.10.2/src/compressed_tensors/transform/factory/base.py

@@ -0,0 +1,164 @@
+# 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 abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+import torch.nn.utils.parametrize as P
+from compressed_tensors.quantization.lifecycle import is_target  # TODO: move to utils
+from compressed_tensors.registry.registry import RegistryMixin, T
+from compressed_tensors.transform import (
+    TransformArgs,
+    TransformLocation,
+    TransformScheme,
+)
+from compressed_tensors.utils import (
+    align_module_device,
+    has_offloaded_params,
+    patch_attr,
+    register_offload_module,
+    update_offload_parameter,
+)
+from torch import Tensor
+from torch.nn import Module, Parameter
+
+
+__all__ = ["TransformFactory", "TransformBase"]
+
+
+class TransformFactory(RegistryMixin, ABC):
+    """
+    Abstract factory base used to create and apply transforms to a model
+
+    :param name: name associated with transform scheme
+    :param scheme: transform scheme which defines how transforms should be created
+    :param seed: random seed used to transform weight randomization
+    """
+
+    def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
+        self.name = name
+        self.scheme = scheme
+        self.generator = torch.Generator()
+        if seed is not None:
+            self.generator.manual_seed(seed)
+
+    @classmethod
+    def from_scheme(cls: type[T], scheme: TransformScheme, **kwargs) -> T:
+        """
+        Create a transform factory from a scheme
+
+        :param scheme: defines how transforms should be created
+        :param kwargs: TransformFactory constructor arguments
+        :return: subclass of `TransformFactory` corresponding to the scheme type
+        """
+        constructor = cls.get_value_from_registry(name=scheme.type)
+        return constructor(scheme=scheme, **kwargs)
+
+    @abstractmethod
+    def create_transform(self, module: Module, args: TransformArgs) -> "TransformBase":
+        """
+        Abstract method which defines how a transform should be created. May utilize
+        caching to maximize shared memory
+
+        :param module: parent module that transform will be applied to
+        :param args: defines how the transform will be applied to the module
+        :return: instance of TransformBase
+        """
+        raise NotImplementedError()
+
+    def apply_to_model(self, model: Module):
+        """
+        Create transforms and apply them to the model
+
+        :param model: module to apply transforms to
+        """
+        for arg in self.scheme.apply:
+            for name, module in list(model.named_modules()):
+                if is_target(name, module, arg.targets, arg.ignore):
+                    self._apply_to_module(module, arg)
+
+    def _apply_to_module(self, module: Module, args: TransformArgs):
+        """
+        Create transforms and apply them to the module
+
+        :param module: target module to apply transforms to
+        :param args: defines how the transform will be applied to the target module
+        """
+        # create transform as submodule
+        transform_name = f"{self.name}_{args.location.value}"
+        transform = self.create_transform(module, args)
+        register_offload_module(module, transform_name, transform)  # (1)
+
+        # register input transformation hook
+        if args.location == TransformLocation.INPUT:
+
+            def input_hook(_, args):
+                input = args[0]
+                return transform(input)
+
+            module.register_forward_pre_hook(input_hook, prepend=True)
+
+        # eagerly apply transformation to weight
+        elif args.location in (
+            TransformLocation.WEIGHT_INPUT,
+            TransformLocation.WEIGHT_OUTPUT,
+        ):
+            assert isinstance(module, torch.nn.Linear)
+            assert module.bias is None
+
+            with torch.no_grad(), align_module_device(module):
+                update_offload_parameter(module, "weight", transform(module.weight))
+
+            if self.scheme.requires_grad:
+                # for training, the weight changes with every forward pass
+                # so we can leverage parametrization to propagate the gradient
+                if has_offloaded_params(module):
+                    raise ValueError("Offloaded training is not supported")
+                P.register_parametrization(module, "weight", transform)
+
+        # register output transformation hook
+        elif args.location == TransformLocation.OUTPUT:
+
+            def output_hook(_, _input, output):
+                return transform(output)
+
+            module.register_forward_hook(output_hook)
+
+        # other locations such as q_attn and k_attn have not been implemented
+        else:
+            raise NotImplementedError()
+
+        # (1) even in the `weight` cases, this submodule attachment is needed in order
+        # to support saving in the frozen state
+
+
+class TransformBase(Module, ABC):
+    """
+    Represents the application of a transform accord to TransformArgs
+    """
+
+    args: TransformArgs
+    weight: Parameter
+
+    @abstractmethod
+    def forward(self, value: Tensor) -> Tensor:
+        raise NotImplementedError()
+
+    def right_inverse(self, value: Tensor) -> Tensor:
+        with patch_attr(self.args, "inverse", not self.args.inverse):
+            return self.forward(value)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(inverse={self.args.inverse})"

compressed_tensors-0.10.2/src/compressed_tensors/transform/factory/hadamard.py

@@ -0,0 +1,79 @@
+# 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
+
+import torch
+from compressed_tensors.transform import TransformArgs, TransformScheme
+from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
+from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
+from compressed_tensors.transform.utils.utils import (
+    apply_transform_weight,
+    get_matrix_size,
+)
+from compressed_tensors.utils import get_offloaded_device
+from compressed_tensors.utils.helpers import ParameterizedDefaultDict
+from torch import Tensor, device, dtype
+from torch.nn import Linear, Module, Parameter
+
+
+@TransformFactory.register("hadamard")
+class HadamardFactory(TransformFactory):
+    """
+    Factory used to apply hadamard transforms to a model
+
+    :param name: name associated with transform scheme
+    :param scheme: transform scheme which defines how transforms should be created
+    :param seed: random seed used to transform weight randomization
+    """
+
+    def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
+        super().__init__(name, scheme, seed)
+        self.weights = ParameterizedDefaultDict(self._create_weight)
+
+    def create_transform(self, module: Module, args: TransformArgs):
+        """
+        Create a HadamardTransform for applying to a module. Transforms with the same
+        size, dtype, and device are cached
+
+        :param module: parent module that transform will be applied to
+        :param args: defines how the transform will be applied to the module
+        """
+        assert isinstance(module, Linear)
+        size = get_matrix_size(module, args.location)
+        dtype = module.weight.dtype
+        device = get_offloaded_device(module)
+
+        weight = self.weights[size, dtype, device]
+        return HadamardTransform(weight, args)
+
+    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        data = deterministic_hadamard_matrix(size, dtype, device)
+        data = data.to(dtype=dtype, device=device)
+        return Parameter(data, requires_grad=self.scheme.requires_grad)
+
+
+class HadamardTransform(TransformBase):
+    def __init__(self, weight: Parameter, args: TransformArgs):
+        super().__init__()
+        self.weight = weight
+        self.args = args
+
+    def forward(self, value: Tensor) -> Tensor:
+        if not self.args.inverse:
+            weight = self.weight
+        else:
+            weight = self.weight.T
+
+        return apply_transform_weight(weight, value, self.args.location)

compressed_tensors-0.10.2/src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -0,0 +1,90 @@
+# 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
+
+import torch
+from compressed_tensors.transform import TransformArgs, TransformScheme
+from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
+from compressed_tensors.transform.utils.utils import (
+    apply_transform_weight,
+    get_matrix_size,
+)
+from compressed_tensors.utils import get_offloaded_device
+from compressed_tensors.utils.helpers import ParameterizedDefaultDict
+from torch import Tensor, device, dtype
+from torch.nn import Linear, Module, Parameter
+
+
+@TransformFactory.register("random-matrix")
+class RandomMatrixFactory(TransformFactory):
+    """
+    Factory used to apply random matrix transforms to a model
+
+    :param name: name associated with transform scheme
+    :param scheme: transform scheme which defines how transforms should be created
+    :param seed: random seed used to transform weight randomization
+    """
+
+    def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
+        super().__init__(name, scheme, seed)
+        self.weights = ParameterizedDefaultDict(self._create_weight)
+        self.inverses = ParameterizedDefaultDict(self._create_inverse)
+
+    def create_transform(self, module: Module, args: TransformArgs):
+        """
+        Create a RandomMatrixTransform for applying to a module. Transforms with the
+        same size, dtype, and device are cached
+
+        :param module: parent module that transform will be applied to
+        :param args: defines how the transform will be applied to the module
+        """
+        assert isinstance(module, Linear)
+        size = get_matrix_size(module, args.location)
+        dtype = module.weight.dtype
+        device = get_offloaded_device(module)
+
+        weight = self.weights[size, dtype, device]
+        if args.inverse:
+            weight = self.inverses[weight]
+
+        return RandomMatrixTransform(weight, args)
+
+    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        data = torch.rand(
+            (size, size), generator=self.generator, dtype=dtype, device=device
+        )
+        return Parameter(data, requires_grad=self.scheme.requires_grad)
+
+    def _create_inverse(self, weight: Parameter) -> Parameter:
+        data = high_precision_invert(weight.data)
+        return Parameter(data, requires_grad=False)
+
+
+class RandomMatrixTransform(TransformBase):
+    def __init__(self, weight: Tensor, args: TransformArgs):
+        super().__init__()
+        self.weight = weight  # is an inverse if args.inverse
+        self.args = args
+
+    def forward(self, value: Tensor) -> Parameter:
+        return apply_transform_weight(self.weight, value, self.args.location)
+
+    def right_inverse(self, value: Tensor) -> Tensor:
+        inverse = high_precision_invert(self.weight)
+        return apply_transform_weight(inverse, value, self.args.location)
+
+
+def high_precision_invert(weight: Tensor) -> Tensor:
+    return torch.linalg.inv(weight.to(torch.float32)).to(weight.dtype)

compressed_tensors-0.10.2/src/compressed_tensors/transform/factory/random_hadamard.py

@@ -0,0 +1,34 @@
+# 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.transform import HadamardFactory, TransformFactory
+from compressed_tensors.transform.utils.hadamard import random_hadamard_matrix
+from torch import device, dtype
+from torch.nn import Parameter
+
+
+@TransformFactory.register("random-hadamard")
+class RandomHadamardFactory(HadamardFactory):
+    """
+    Factory used to apply random hadamard transforms to a model
+
+    :param name: name associated with transform scheme
+    :param scheme: transform scheme which defines how transforms should be created
+    :param seed: random seed used to transform weight randomization
+    """
+
+    def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        data = random_hadamard_matrix(size, dtype, device, self.generator)
+        data = data.to(dtype=dtype, device=device)
+        return Parameter(data, requires_grad=self.scheme.requires_grad)

compressed_tensors-0.10.2/src/compressed_tensors/transform/utils/hadamard.py

@@ -0,0 +1,160 @@
+# 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 math
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors import safe_open
+
+
+REPO_PATH = Path(__file__).parent / "hadamards.safetensors"
+
+
+__all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"]
+
+
+# note that hadamard matrix multiplication can be accelerated using a library such as
+# https://github.com/Dao-AILab/fast-hadamard-transform/tree/master
+
+
+def deterministic_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """
+    Construct an n-by-n Hadamard matrix, using Sylvester's construction.
+    `n` must be a power of 2.
+
+    Adapated from https://github.com/scipy/scipy/blob/v1.15.2/scipy/linalg/_special_matrices.py  # noqa: E501
+
+    :param size: order of the matrix, must be a power of 2
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :return: hadamard matrix of size `size`
+    """
+    if size <= 0:
+        raise ValueError("Cannot construct deterministic hadamard of size <= 0")
+
+    log2 = int(math.log2(size))
+    if size != 2**log2:
+        raise ValueError("Cannot construct deterministic hadamard of size != 2^n")
+
+    H = torch.tensor([[1]], dtype=dtype, device=device)
+
+    # Sylvester's construction
+    for _ in range(log2):
+        H = torch.vstack((torch.hstack((H, H)), torch.hstack((H, -H))))
+
+    return H / math.sqrt(size)
+
+
+def random_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+    gen: Optional[torch.Generator] = None,
+) -> torch.Tensor:
+    """
+    Produces a randomly generated Hadamard matrix. Differs from
+    `deterministic_hadamard_matrix` in that this function supports non powers of 2
+    and randomization using a seeded generator
+
+    Adapated from https://github.com/facebookresearch/SpinQuant/blob/main/utils/hadamard_utils.py  # noqa: E501
+    Known matrices were retrieved from N. J. A. Sloane's Library of Hadamard Matrices http://www.neilsloane.com/hadamard/  # noqa: E501
+
+    :param size: The dimension of the hamadard matrix
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :param gen: Optional generator random values
+    :return: randomly generated hadamard matrix
+    """
+    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=dtype)  # cpu
+    Q = Q.to(device=device)
+    Q = Q * 2 - 1
+    Q = torch.diag(Q)
+    return _matmul_hadU(Q) / math.sqrt(size)
+
+
+def is_pow2(n: int) -> bool:
+    """
+    Check if a number is a power of 2
+
+    :param n: number to check
+    :return: True iff `n` is a power of 2
+    """
+    return n > 0 and (n & (n - 1) == 0)
+
+
+def _fetch_hadamard_divisor(
+    n: int,
+    dtype: torch.dtype,
+    device: torch.device = torch.device("cpu"),
+    file_path: str = REPO_PATH,
+) -> Optional[torch.Tensor]:
+    """
+    Fetch a known hadamard matrix from the given file path. The returned matrix will
+    be of of size `k` such that `n / k` is a power of two. Return None if no such
+    matrix exists.
+
+    Note: This function reopens the safetensors file every time it is called.
+    This is technically inefficient, but a very small runtime cost and simpler
+    than forcing callers to manage the file open context
+
+    :param n: size of known hadamard matrix
+    :return: a known hadamard matrix of size `n` if one exists, else None
+    """
+    with safe_open(file_path, framework="pt", device=str(device)) as file:
+        divisors = sorted((int(key) for key in file.keys()), reverse=True)
+        for divisor in divisors:
+            if n % divisor == 0 and is_pow2(n // divisor):
+                return file.get_tensor(str(divisor)).to(dtype=dtype)
+
+    return None
+
+
+def _matmul_hadU(X: torch.Tensor) -> torch.Tensor:
+    size = X.size(0)
+    dtype = X.dtype
+    device = X.device
+
+    # Check if we have the determined hadamard matrix
+    hadK = _fetch_hadamard_divisor(size, dtype, device=device)
+    if hadK is None:
+        raise ValueError(f"Cannot construct random hadamard matrix of size {size}")
+    K = hadK.size(0)
+
+    # Reshape diag matrix with randomized -1/+1
+    input = X.clone().view(-1, size, 1)
+    output = input.clone()
+    while input.shape[1] > K:
+        input = input.view(input.shape[0], input.shape[1] // 2, 2, input.shape[2])
+        output = output.view(input.shape)
+        output[:, :, 0, :] = input[:, :, 0, :] + input[:, :, 1, :]
+        output[:, :, 1, :] = input[:, :, 0, :] - input[:, :, 1, :]
+        output = output.view(input.shape[0], input.shape[1], -1)
+        (input, output) = (output, input)
+    assert input.shape[1] == K
+    del output
+
+    # Do not explicitly repeat - OOM
+    # input = torch.bmm(
+    #     hadK.repeat(len(input), 1, 1).to(input.device).to(input.dtype), input)
+    # Use bcast instead
+    input = hadK.view(1, K, K).to(input) @ input
+
+    # normalize
+    return input.view(X.shape)