compressed-tensors 0.10.3a20250701__tar.gz → 0.10.3a20250707__tar.gz

{compressed_tensors-0.10.3a20250701/src/compressed_tensors.egg-info → compressed_tensors-0.10.3a20250707}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.3a20250701
+Version: 0.10.3a20250707
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -11,6 +11,7 @@ License-File: LICENSE
 Requires-Dist: torch>=1.7.0
 Requires-Dist: transformers
 Requires-Dist: pydantic>=2.0
+Requires-Dist: frozendict
 Provides-Extra: dev
 Requires-Dist: black==22.12.0; extra == "dev"
 Requires-Dist: isort==5.8.0; extra == "dev"

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/setup.py

@@ -88,7 +88,7 @@ def _setup_packages() -> List:
     )
 
 def _setup_install_requires() -> List:
-    return ["torch>=1.7.0", "transformers", "pydantic>=2.0"]
+    return ["torch>=1.7.0", "transformers", "pydantic>=2.0", "frozendict"]
 
 def _setup_extras() -> Dict:
     return {

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/transform/factory/hadamard.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
@@ -22,7 +22,7 @@ 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 import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
 from torch import Tensor, device, dtype
 from torch.nn import Linear, Module, Parameter
@@ -41,6 +41,7 @@ class HadamardFactory(TransformFactory):
     def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
         super().__init__(name, scheme, seed)
         self.weights = ParameterizedDefaultDict(self._create_weight)
+        self.perms = ParameterizedDefaultDict(self._create_permutation)
 
     def create_transform(self, module: Module, args: TransformArgs):
         """
@@ -54,26 +55,46 @@ class HadamardFactory(TransformFactory):
         size = get_matrix_size(module, args.location)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
+        exec_device = get_execution_device(module)
 
-        weight = self.weights[size, dtype, device]
-        return HadamardTransform(weight, args)
+        factory_kwargs = {"construct_device": exec_device}
+        weight = self.weights.get(size, dtype, device, factory_kwargs=factory_kwargs)
+        perm = self.perms[weight] if self.scheme.randomize else None
+        return HadamardTransform(weight, perm, 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)
+    def _create_weight(
+        self,
+        size: int,
+        dtype: dtype,
+        device: device,
+        construct_device: device,
+    ) -> Parameter:
+        # construct on execution device, cache on offload device
+        data = deterministic_hadamard_matrix(size, dtype, construct_device)
+        data = data.to(device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 
+    def _create_permutation(self, weight: Parameter) -> Parameter:
+        data = torch.randperm(weight.size(0), generator=self.generator)
+        return Parameter(data, requires_grad=False)
+
 
 class HadamardTransform(TransformBase):
-    def __init__(self, weight: Parameter, args: TransformArgs):
+    def __init__(
+        self, weight: Parameter, perm: Union[Parameter, None], args: TransformArgs
+    ):
         super().__init__()
         self.weight = weight
+        self.perm = perm
         self.args = args
 
     def forward(self, value: Tensor) -> Tensor:
-        if not self.args.inverse:
-            weight = self.weight
-        else:
-            weight = self.weight.T
+        weight = self.weight
+
+        if self.perm is not None:
+            weight = weight[self.perm][:, self.perm]
+
+        if self.args.inverse:
+            weight = weight.T
 
         return apply_transform_weight(weight, value, self.args.location)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -62,6 +62,7 @@ class RandomMatrixFactory(TransformFactory):
         return RandomMatrixTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        # TODO: verify that weight is invertible (has non-zero determinant)
         data = torch.rand(
             (size, size), generator=self.generator, dtype=dtype, device=device
         )

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/transform/factory/random_hadamard.py

@@ -28,7 +28,14 @@ class RandomHadamardFactory(HadamardFactory):
     :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)
+    def _create_weight(
+        self,
+        size: int,
+        dtype: dtype,
+        device: device,
+        construct_device: device,
+    ) -> Parameter:
+        # construct on execution device, cache on offload device
+        data = random_hadamard_matrix(size, dtype, construct_device, self.generator)
+        data = data.to(device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/transform/transform_config.py

@@ -49,7 +49,7 @@ QUIP = TransformConfig(
                     inverse=True,
                 ),
             ],
-            randomize_modules=True,
+            randomize=True,
         ),
         "u": TransformScheme(
             type="hadamard",
@@ -62,7 +62,7 @@ QUIP = TransformConfig(
                     targets=["Linear"], location="output", inverse=True  # non-mergable
                 ),
             ],
-            randomize_modules=True,
+            randomize=True,
         ),
     }
 )

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/transform/transform_scheme.py

@@ -31,13 +31,12 @@ class TransformScheme(BaseModel):
         (see `Transforms.registered_names()`)
     :param apply: list of TransformationArgs containing the information about the
         modules that should be targeted by the specified transform
-    :param randomize_modules: True if unique transforms should be applied to each
-        unique module targeted by `apply`, otherwise reuse transform weights where
-        applicable
+    :param randomize: True if uniquely randomized transform weights should be used,
+        otherwise use identical transform weights where applicable
     :param requires_grad: True if weights include gradients for training
     """
 
     type: str
     apply: List[TransformArgs] = Field(default_factory=list)
-    randomize_modules: bool = Field(default=False)
+    randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/utils/helpers.py

@@ -15,10 +15,11 @@
 import contextlib
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
 
 import numpy
 import torch
+from frozendict import frozendict
 from transformers import AutoConfig
 
 
@@ -373,11 +374,23 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
+        self._factory_kwargs = frozendict()
 
-    def __missing__(self, key):
+    def __missing__(self, key: Any) -> Any:
         if isinstance(key, tuple):
-            value = self.default_factory(*key)
+            value = self.default_factory(*key, **self._factory_kwargs)
         else:
-            value = self.default_factory(key)
+            value = self.default_factory(key, **self._factory_kwargs)
         self[key] = value
         return value
+
+    def get(self, *args, factory_kwargs: Mapping = frozendict()) -> Any:
+        """
+        Similar to `__getitem__`, but allows passing kwargs to factory function
+
+        :param \\*args: args whose tuple will value will be treated as key
+        :param factory_kwargs: keyword arguments to pass to `default_factory`
+        :return: dictionary entry for given key
+        """
+        with patch_attr(self, "_factory_kwargs", factory_kwargs):
+            return self[args]

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.10.3.a20250701'
+__version__ = version = '0.10.3.a20250707'
 __version_tuple__ = version_tuple = (0, 10, 3)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.3a20250701
+Version: 0.10.3a20250707
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -11,6 +11,7 @@ License-File: LICENSE
 Requires-Dist: torch>=1.7.0
 Requires-Dist: transformers
 Requires-Dist: pydantic>=2.0
+Requires-Dist: frozendict
 Provides-Extra: dev
 Requires-Dist: black==22.12.0; extra == "dev"
 Requires-Dist: isort==5.8.0; extra == "dev"

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors.egg-info/SOURCES.txt

@@ -130,6 +130,7 @@ tests/test_quantization/test_configs/__init__.py
 tests/test_quantization/test_configs/test_bit_depths.py
 tests/test_quantization/test_configs/test_strategies.py
 tests/test_quantization/test_utils/test_helpers.py
+tests/test_transform/conftest.py
 tests/test_transform/test_transform_args.py
 tests/test_transform/test_transform_config.py
 tests/test_transform/test_transform_scheme.py

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/src/compressed_tensors.egg-info/requires.txt

@@ -1,6 +1,7 @@
 torch>=1.7.0
 transformers
 pydantic>=2.0
+frozendict
 
 [accelerate]
 accelerate

compressed_tensors-0.10.3a20250707/tests/test_transform/conftest.py

@@ -0,0 +1,54 @@
+# 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
+
+
+class TransformableModel(torch.nn.Module):
+    def __init__(self, *sizes):
+        super().__init__()
+        self.fcs = torch.nn.ModuleList(
+            [
+                torch.nn.Linear(sizes[index], sizes[index + 1], bias=False)
+                for index in range(0, len(sizes) - 1)
+            ]
+        )
+
+    def forward(self, x):
+        for layer in self.fcs:
+            x = layer(x)
+        return x
+
+
+@pytest.fixture(scope="function")
+def model_apply():
+    model = TransformableModel(2, 4, 8, 16, 32, 64)
+    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),
+    ]
+
+    return model, apply

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/tests/test_transform/factory/test_correctness.py

@@ -16,6 +16,7 @@ import pytest
 import torch
 from compressed_tensors.transform import (
     TransformArgs,
+    TransformConfig,
     TransformFactory,
     TransformScheme,
 )
@@ -23,27 +24,18 @@ from compressed_tensors.utils import offloaded_dispatch
 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 scheme_kwargs():
+    all_types = TransformFactory.registered_names()
+    base = [{"type": type} for type in all_types]
+    randomized = [{"type": type, "randomize": True} for type in all_types]
+    return base + randomized
 
-    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):
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_correctness_linear(scheme_kwargs):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=True)
+    scheme = TransformScheme(**scheme_kwargs)
     factory = TransformFactory.from_scheme(scheme, name="")
 
     input_tfm = factory.create_transform(
@@ -67,50 +59,39 @@ def test_correctness_linear(scheme):
     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):
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_correctness_model(scheme_kwargs, model_apply, offload=False):
     # load model
-    model = TransformableModel(2, 4, 8, 16, 32, 64)
+    model = model_apply[0]
     if offload:
         model = offloaded_dispatch(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
+    # get output
     input = torch.rand((17, model.fcs[0].in_features))
     if offload:
         input = input.to(torch.device("cuda"))
+    true_output = model(input)
+
+    # apply transforms
+    config = TransformConfig(
+        config_groups={
+            "": TransformScheme(
+                **scheme_kwargs,
+                apply=model_apply[1],
+            )
+        }
+    )
+    for name, scheme in config.config_groups.items():
+        factory = TransformFactory.from_scheme(scheme, name=name)
+        factory.apply_to_model(model)
 
     # 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)
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_correctness_model_offload(scheme_kwargs, model_apply):
+    test_correctness_model(scheme_kwargs, model_apply, offload=True)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/tests/test_transform/factory/test_memory.py

@@ -19,49 +19,44 @@ import torch
 from compressed_tensors.transform import (
     TransformArgs,
     TransformBase,
+    TransformConfig,
     TransformFactory,
     TransformScheme,
 )
 from compressed_tensors.utils import align_modules, offloaded_dispatch
+from tests.test_transform.conftest import TransformableModel
 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 scheme_kwargs():
+    all_types = TransformFactory.registered_names()
+    base = [{"type": type} for type in all_types]
+    randomized = [{"type": type, "randomize": True} for type in all_types]
+    return base + randomized
 
-    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_memory_sharing(scheme, offload=False):
-    # load scheme and factory
-    scheme = TransformScheme(
-        type="hadamard",
-        apply=[
-            TransformArgs(targets="Linear", location="input"),
-            TransformArgs(targets="Linear", location="output"),
-        ],
-    )
-    factory = TransformFactory.from_scheme(scheme, name="")
 
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_memory_sharing(scheme_kwargs, offload=False):
     # load model (maybe with offloading)
     model = TransformableModel(2, 2, 4, 4, 8, 8)
     if offload:
         offloaded_dispatch(model, torch.device("cuda"))
 
     # add transforms to model
-    factory.apply_to_model(model)
+    config = TransformConfig(
+        config_groups={
+            "": TransformScheme(
+                **scheme_kwargs,
+                apply=[
+                    TransformArgs(targets="Linear", location="input"),
+                    TransformArgs(targets="Linear", location="output"),
+                ],
+            )
+        }
+    )
+    for name, scheme in config.config_groups.items():
+        factory = TransformFactory.from_scheme(scheme, name=name)
+        factory.apply_to_model(model)
 
     # check that memory is shared when onloaded
     with align_modules(model.modules()):
@@ -93,20 +88,12 @@ def test_memory_sharing(scheme, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
-def test_memory_sharing_offload(scheme):
-    test_memory_sharing(scheme, offload=True)
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_memory_sharing_offload(scheme_kwargs):
+    test_memory_sharing(scheme_kwargs, offload=True)
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [
-        TransformScheme(type=name, requires_grad=True)
-        for name in TransformFactory.registered_names()
-    ],
-)
-def test_memory_sharing_training(scheme):
-    test_memory_sharing(scheme, offload=False)
+@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
+def test_memory_sharing_training(scheme_kwargs):
+    scheme_kwargs["requires_grad"] = True
+    test_memory_sharing(scheme_kwargs, offload=False)

{compressed_tensors-0.10.3a20250701 → compressed_tensors-0.10.3a20250707}/tests/test_transform/test_transform_scheme.py

@@ -24,7 +24,7 @@ def test_basic_scheme():
         type="hadamard",
         apply=[basic_args],
     )
-    assert not scheme.randomize_modules
+    assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 1
     assert isinstance(scheme.apply[0], TransformArgs)
@@ -43,10 +43,10 @@ def test_multiple_groups_global():
     scheme = TransformScheme(
         type="hadamard",
         apply=[embedding_args, linear_args],
-        randomize_modules=True,
+        randomize=True,
     )
 
-    assert scheme.randomize_modules
+    assert scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 2
     assert isinstance(scheme.apply[0], TransformArgs)
@@ -69,6 +69,6 @@ def test_multiple_groups():
         apply=apply,
     )
 
-    assert not scheme.randomize_modules
+    assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 20