compressed-tensors 0.10.2a20250606__tar.gz → 0.10.2a20250611__tar.gz

{compressed_tensors-0.10.2a20250606/src/compressed_tensors.egg-info → compressed_tensors-0.10.2a20250611}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250606
+Version: 0.10.2a20250611
 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.2a20250606 → compressed_tensors-0.10.2a20250611}/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.2a20250606 → compressed_tensors-0.10.2a20250611}/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.2a20250606 → compressed_tensors-0.10.2a20250611}/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.2a20250606 → compressed_tensors-0.10.2a20250611}/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.2a20250611/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.2a20250611/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)
+        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.2a20250611/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.2a20250611/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, self.generator)
+        data = data.to(dtype=dtype, device=device)
+        return Parameter(data, requires_grad=self.scheme.requires_grad)

{compressed_tensors-0.10.2a20250606 → compressed_tensors-0.10.2a20250611}/src/compressed_tensors/utils/offload.py

@@ -87,12 +87,15 @@ def check_accelerate(fallback: Any):
         if not _has_accelerate:
 
             if fallback == "error":
+
                 @wraps(func)
                 def fallback_fn(*args, **kwargs):
                     raise ValueError(
                         "Please install `accelerate` in order to use this function"
                     )
+
             else:
+
                 @wraps(func)
                 def fallback_fn(*args, **kwargs):
                     return fallback

{compressed_tensors-0.10.2a20250606 → compressed_tensors-0.10.2a20250611}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.10.2.a20250606'
+__version__ = version = '0.10.2.a20250611'
 __version_tuple__ = version_tuple = (0, 10, 2)

{compressed_tensors-0.10.2a20250606 → compressed_tensors-0.10.2a20250611/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250606
+Version: 0.10.2a20250611
 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.2a20250606 → compressed_tensors-0.10.2a20250611}/src/compressed_tensors.egg-info/SOURCES.txt

@@ -75,6 +75,11 @@ src/compressed_tensors/transform/__init__.py
 src/compressed_tensors/transform/transform_args.py
 src/compressed_tensors/transform/transform_config.py
 src/compressed_tensors/transform/transform_scheme.py
+src/compressed_tensors/transform/factory/__init__.py
+src/compressed_tensors/transform/factory/base.py
+src/compressed_tensors/transform/factory/hadamard.py
+src/compressed_tensors/transform/factory/matrix_multiply.py
+src/compressed_tensors/transform/factory/random_hadamard.py
 src/compressed_tensors/transform/utils/__init__.py
 src/compressed_tensors/transform/utils/hadamard.py
 src/compressed_tensors/transform/utils/utils.py
@@ -127,6 +132,8 @@ tests/test_quantization/test_utils/test_helpers.py
 tests/test_transform/test_transform_args.py
 tests/test_transform/test_transform_config.py
 tests/test_transform/test_transform_scheme.py
+tests/test_transform/factory/test_correctness.py
+tests/test_transform/factory/test_memory.py
 tests/test_transform/utils/test_hadamard.py
 tests/test_utils/__init__.py
 tests/test_utils/test_helpers.py

compressed_tensors-0.10.2a20250611/tests/test_transform/factory/test_correctness.py

@@ -0,0 +1,116 @@
+# 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 pytest
+import torch
+from compressed_tensors.transform import (
+    TransformArgs,
+    TransformFactory,
+    TransformScheme,
+)
+from compressed_tensors.utils import align_modules, force_cpu_offload
+from tests.testing_utils import requires_accelerate, requires_gpu
+
+
+class TransformableModel(torch.nn.Module):
+    def __init__(self, *sizes):
+        super().__init__()
+        self.fcs = torch.nn.ModuleList([])
+        self.fcs.append(torch.nn.Linear(sizes[0], sizes[1], bias=False))
+        for index in range(1, len(sizes) - 1):
+            self.fcs.append(torch.nn.Linear(sizes[index], sizes[index + 1], bias=False))
+
+    def forward(self, x):
+        for layer in self.fcs:
+            x = layer(x)
+        return x
+
+
+@pytest.mark.parametrize(
+    "scheme",
+    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
+)
+def test_correctness_linear(scheme):
+    size = (4, 8)
+    module = torch.nn.Linear(*size, bias=True)
+    factory = TransformFactory.from_scheme(scheme, name="")
+
+    input_tfm = factory.create_transform(
+        module, TransformArgs(targets="Linear", location="input", inverse=True)
+    )
+    w_in_tfm = factory.create_transform(
+        module, TransformArgs(targets="Linear", location="weight_input")
+    )
+    w_out_tfm = factory.create_transform(
+        module, TransformArgs(targets="Linear", location="weight_output")
+    )
+    output_tfm = factory.create_transform(
+        module, TransformArgs(targets="Linear", location="output", inverse=True)
+    )
+
+    input = torch.rand((17, size[0]))
+    true_output = input @ module.weight.T
+    input_transformed = input_tfm(input)
+    weight_transformed = w_out_tfm(w_in_tfm(module.weight))
+    output = output_tfm(input_transformed @ weight_transformed.T)
+    assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0)
+
+
+@pytest.mark.parametrize(
+    "scheme",
+    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
+)
+def test_correctness_model(scheme, offload=False):
+    # load model
+    model = TransformableModel(2, 4, 8, 16, 32, 64)
+    if offload:
+        model = force_cpu_offload(model, torch.device("cuda"))
+
+    # create factory
+    scheme.apply = [
+        # weight output -> input
+        TransformArgs(targets="fcs.0", location="weight_output"),
+        TransformArgs(targets="fcs.1", location="input", inverse=True),
+        # output -> weight input
+        TransformArgs(targets="fcs.1", location="output"),
+        TransformArgs(targets="fcs.2", location="weight_input", inverse=True),
+        # output -> input
+        TransformArgs(targets="fcs.2", location="output"),
+        TransformArgs(targets="fcs.3", location="input", inverse=True),
+        # weight output -> weight input
+        TransformArgs(targets="fcs.3", location="weight_output"),
+        TransformArgs(targets="fcs.4", location="weight_input", inverse=True),
+    ]
+    factory = TransformFactory.from_scheme(scheme, name="")
+
+    # create inputs
+    input = torch.rand((17, model.fcs[0].in_features))
+    if offload:
+        input = input.to(torch.device("cuda"))
+
+    # compare outputs
+    true_output = model(input)
+    factory.apply_to_model(model)
+    output = model(input)
+    assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0)
+
+
+@requires_gpu
+@requires_accelerate()
+@pytest.mark.parametrize(
+    "scheme",
+    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
+)
+def test_correctness_model_offload(scheme):
+    test_correctness_model(scheme, offload=True)