compressed-tensors 0.10.3a20250715__py3-none-any.whl → 0.10.3a20250716__py3-none-any.whl

compressed_tensors/transform/factory/base.py

@@ -117,10 +117,8 @@ class TransformFactory(RegistryMixin, ABC):
             TransformLocation.WEIGHT_INPUT,
             TransformLocation.WEIGHT_OUTPUT,
         ):
-            assert isinstance(module, torch.nn.Linear)
-            assert module.bias is None
-
             # fuse transform into weight
+            assert hasattr(module, "weight")
             with torch.no_grad(), align_module_device(module):
                 update_offload_parameter(module, "weight", transform(module.weight))
 

compressed_tensors/transform/factory/hadamard.py

@@ -14,13 +14,14 @@
 
 from typing import Optional, Union
 
+import math
 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 (
+from compressed_tensors.transform.utils.matrix import (
     apply_transform_weight,
-    get_matrix_size,
+    get_transform_size,
 )
 from compressed_tensors.utils import get_execution_device, get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
@@ -51,8 +52,8 @@ class HadamardFactory(TransformFactory):
         :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)
+        assert hasattr(module, "weight")
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
@@ -60,7 +61,7 @@ class HadamardFactory(TransformFactory):
         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)
+        return HadamardTransform(weight, perm, args, type(module))
 
     def _create_weight(
         self,
@@ -81,12 +82,18 @@ class HadamardFactory(TransformFactory):
 
 class HadamardTransform(TransformBase):
     def __init__(
-        self, weight: Parameter, perm: Union[Parameter, None], args: TransformArgs
+        self,
+        weight: Parameter,
+        perm: Optional[Parameter],
+        args: TransformArgs,
+        module_type: type[torch.nn.Module],
     ):
         super().__init__()
         self.weight = weight
         self.perm = perm
         self.args = args
+        self.module_type = module_type
+        self._scale = math.sqrt(weight.size(0))
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -96,5 +103,7 @@ class HadamardTransform(TransformBase):
 
         if self.args.inverse:
             weight = weight.T
-
-        return apply_transform_weight(weight, value, self.args.location)
+ 
+        return apply_transform_weight(
+            weight, value, self.args.location, self.module_type
+        ) / self._scale

compressed_tensors/transform/factory/matrix_multiply.py

@@ -17,9 +17,9 @@ 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 (
+from compressed_tensors.transform.utils.matrix import (
     apply_transform_weight,
-    get_matrix_size,
+    get_transform_size,
 )
 from compressed_tensors.utils import get_offloaded_device
 from compressed_tensors.utils.helpers import ParameterizedDefaultDict
@@ -50,8 +50,8 @@ class RandomMatrixFactory(TransformFactory):
         :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)
+        assert hasattr(module, "weight")
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
 
@@ -59,7 +59,7 @@ class RandomMatrixFactory(TransformFactory):
         if args.inverse:
             weight = self.inverses[weight]
 
-        return RandomMatrixTransform(weight, args)
+        return RandomMatrixTransform(weight, args, type(module))
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
         # TODO: verify that weight is invertible (has non-zero determinant)
@@ -74,17 +74,27 @@ class RandomMatrixFactory(TransformFactory):
 
 
 class RandomMatrixTransform(TransformBase):
-    def __init__(self, weight: Tensor, args: TransformArgs):
+    def __init__(
+        self,
+        weight: Tensor,
+        args: TransformArgs,
+        module_type: type[torch.nn.Module],
+    ):
         super().__init__()
         self.weight = weight  # is an inverse if args.inverse
         self.args = args
+        self.module_type = module_type
 
     def forward(self, value: Tensor) -> Parameter:
-        return apply_transform_weight(self.weight, value, self.args.location)
+        return apply_transform_weight(
+            self.weight, value, self.args.location, self.module_type
+        )
 
     def right_inverse(self, value: Tensor) -> Tensor:
         inverse = high_precision_invert(self.weight)
-        return apply_transform_weight(inverse, value, self.args.location)
+        return apply_transform_weight(
+            inverse, value, self.args.location, self.module_type
+        )
 
 
 def high_precision_invert(weight: Tensor) -> Tensor:

compressed_tensors/transform/transform_scheme.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import List, Optional
 
 from compressed_tensors.transform import TransformArgs
 from pydantic import BaseModel, Field
@@ -40,3 +40,4 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
+    head_dim: Optional[int] = Field(default=None)

compressed_tensors/transform/utils/hadamard.py

@@ -59,7 +59,7 @@ def deterministic_hadamard_matrix(
     for _ in range(log2):
         H = torch.vstack((torch.hstack((H, H)), torch.hstack((H, -H))))
 
-    return H / math.sqrt(size)
+    return H
 
 
 def random_hadamard_matrix(
@@ -86,7 +86,7 @@ def random_hadamard_matrix(
     Q = Q.to(device=device)
     Q = Q * 2 - 1
     Q = torch.diag(Q)
-    return _matmul_hadU(Q) / math.sqrt(size)
+    return _matmul_hadU(Q)
 
 
 def is_pow2(n: int) -> bool:

compressed_tensors/transform/utils/matrix.py

@@ -0,0 +1,179 @@
+# 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 Callable, Optional, Tuple
+
+import torch
+from compressed_tensors.transform import TransformLocation
+
+
+__all__ = ["get_transform_size", "apply_transform_weight"]
+
+
+def get_transform_size(
+    module: torch.nn.Module,
+    location: TransformLocation,
+    head_dim: Optional[int] = None,
+) -> int:
+    """
+    Determine the size of a transform matrix given its location on the module
+
+    :param module: module that matrix will be applied to
+    :param location: location on module
+    :param head_dim: size of head when transform is applied to mha
+    :return: size of matrix
+    """
+    if isinstance(module, torch.nn.Linear):
+        if location in (TransformLocation.INPUT, TransformLocation.WEIGHT_INPUT):
+            size = module.in_features
+        else:
+            size = module.out_features
+    elif isinstance(module, torch.nn.Embedding):
+        if location in (TransformLocation.INPUT, TransformLocation.WEIGHT_INPUT):
+            size = module.num_embeddings
+        else:
+            size = module.embedding_dim
+    else:
+        raise NotImplementedError(f"Transforms on {type(module)} are not supported")
+
+    if head_dim is not None:
+        if size % head_dim != 0:
+            raise ValueError(
+                f"{head_dim} must divide {size} for {type(module)} at {location}"
+            )
+
+        size = head_dim
+
+    return size
+
+
+def apply_transform_weight(
+    transform_weight: torch.Tensor,
+    value: torch.Tensor,
+    location: TransformLocation,
+    module_type: type[torch.nn.Module],
+) -> torch.Tensor:
+    """
+    Using the transform location, apply the transform_weight to the
+    given value wrt linear weights. For more info on input and output transforms,
+    see `TransformLocation`
+
+    The following explains how weights should be applied to values according to location
+
+    let  x          be input activation
+         W          be weight,
+         yh, xh, Wh be transformed output, input, weight
+
+    note that
+         y  = (x W.T)        // torch.nn.Linear
+
+    Choose values for yh, xh, and Wh which incorporate matrix transforms
+
+    let  V, Vi      be transform matrices on input side
+         U, Ui      be transform matrices on output side
+
+    pick xh = (x V)
+         Wh = (U.T W Vi.T)
+         yh = (y U)
+
+    The following shows that `yh = (xh) (Wh).T` for the chosen values of yh, xh, and Wh
+
+    (xh) (Wh).T = (x V) (U.T W Vi.T).T
+                = (x V) (Vi W.T U)        // transpose matrix product identity
+                = (x W.T) U
+                = y U
+                = yh
+
+    :param transform_weight: transform weight to apply
+    :param value: value to apply transform_weight to
+    :param location: determines how weight should be applied
+    :param model_type: result of type(module), passed in to determine application of
+        weight transform
+    :return: value after transform_weight has been applied
+    """
+
+    assert transform_weight.shape[0] == transform_weight.shape[1]
+
+    if module_type == torch.nn.Linear:
+        if location == TransformLocation.INPUT:
+            return _multihead_matmul(value, transform_weight)
+
+        elif location == TransformLocation.WEIGHT_INPUT:
+            # equivalent to (transform_weight @ value.T).T
+            return _multihead_matmul(value, transform_weight.T)
+
+        elif location == TransformLocation.WEIGHT_OUTPUT:
+            # equivalent to (value.T @ transform_weight).T
+            return _multihead_matmul(transform_weight.T, value)
+
+        elif location == TransformLocation.OUTPUT:
+            return _multihead_matmul(value, transform_weight)
+
+    # similar derivation to torch.nn.Linear, but `y = (x W)`
+    elif module_type == torch.nn.Embedding:
+        if location == TransformLocation.INPUT:
+            return _multihead_matmul(value, transform_weight)
+
+        elif location == TransformLocation.WEIGHT_INPUT:
+            return _multihead_matmul(
+                transform_weight,
+                value,
+            )
+
+        elif location == TransformLocation.WEIGHT_OUTPUT:
+            return _multihead_matmul(value, transform_weight)
+
+        elif location == TransformLocation.OUTPUT:
+            return _multihead_matmul(value, transform_weight)
+
+    raise NotImplementedError(
+        f"Applying transforms to {module_type} {location} is not supported"
+    )
+
+
+def _multihead_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """
+    Performs A @ B for last two dims of two matrices A and B that possibly
+    have different shapes, as is the case in multi-headed dimension. If
+    shapes are different, this is equivalent to converting the last two dims
+    of the smaller matrix into a block-diagonal matrix with the same shape as
+    the last two dims of the larger matrix.
+
+    E.g. if A is half the size of B, this function will perform
+    [[A  ]  @ B
+     [  A]]
+
+    If B is a third of the size of A, this function will perform
+    A @ [[B    ]
+         [  B  ]
+         [    B]]
+
+    This function will error out if the shapes are not evenly divisble
+
+    :param A: left-hand tensor
+    :param B: right-hand tensor
+    :return: result
+    """
+    if A.shape[-1] > B.shape[-2]:
+        head_dim = B.shape[-2]
+        num_heads = A.shape[-1] // head_dim
+        A = A.unflatten(-1, (num_heads, head_dim))
+        return (A @ B).flatten(-2, -1)
+    elif A.shape[-1] < B.shape[-2]:
+        head_dim = A.shape[-1]
+        num_heads = B.shape[-2] // head_dim
+        B = B.unflatten(-2, (num_heads, head_dim))
+        return (A @ B).flatten(-3, -2)
+    else:
+        return A @ B

compressed_tensors/version.py

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

{compressed_tensors-0.10.3a20250715.dist-info → compressed_tensors-0.10.3a20250716.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.3a20250715
+Version: 0.10.3a20250716
 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.3a20250715.dist-info → compressed_tensors-0.10.3a20250716.dist-info}/RECORD

@@ -1,6 +1,6 @@
 compressed_tensors/__init__.py,sha256=UtKmifNeBCSE2TZSAfduVNNzHY-3V7bLjZ7n7RuXLOE,812
 compressed_tensors/base.py,sha256=73HYH7HY7O2roC89yG_piPFnZwrBfn_i7HmKl90SKc0,875
-compressed_tensors/version.py,sha256=4fhlyf_Dxsw-6e50R1K8Mn8AWFNP0TiRJ8S21Gge6UA,523
+compressed_tensors/version.py,sha256=DSQFEQZQHt-pmBpCtXPX8Vc2dUeG5ueZMmYLIwONR1c,523
 compressed_tensors/compressors/__init__.py,sha256=smSygTSfcfuujRrAXDc6uZm4L_ccV1tWZewqVnOb4lM,825
 compressed_tensors/compressors/base.py,sha256=nvWsv4xEw1Tkxkxth6TmHplDYXfBeP22xWxOsZERyDY,7204
 compressed_tensors/compressors/helpers.py,sha256=OK6qxX9j3bHwF9JfIYSGMgBJe2PWjlTA3byXKCJaTIQ,5431
@@ -43,16 +43,16 @@ compressed_tensors/transform/__init__.py,sha256=v2wfl4CMfA6KbD7Hxx_MbRev63y_6QLD
 compressed_tensors/transform/apply.py,sha256=Cnc7Q8d8FzpLGtXixvdPzqApfjAXpfShxvVl_7nNJ4E,1259
 compressed_tensors/transform/transform_args.py,sha256=jJY-Qt996w45LWQ10AHd7tUtNrnV9mjD9M5D4SZ5B3E,3199
 compressed_tensors/transform/transform_config.py,sha256=A3RuLNDqBNEByQNeu40Kg7sItwE6kWgnX18Umg1uONI,2128
-compressed_tensors/transform/transform_scheme.py,sha256=JAFQoCiNLg04diXG5KsynRGcLIB0Y0tC5s8U7HoDM7c,1692
+compressed_tensors/transform/transform_scheme.py,sha256=uGLC4avdbhrVqNC3-Eo0p7WzNRQK92Fpg0N9hWiuCRQ,1752
 compressed_tensors/transform/factory/__init__.py,sha256=fH6rjBYAxuwrTzBTlTjTgCYNyh6TCvCqajCz4Im4YrA,617
-compressed_tensors/transform/factory/base.py,sha256=D5s6dOn01Yy_Z2v8YmYIYfzJie2BYV6uxBJKP_9jTCQ,6044
-compressed_tensors/transform/factory/hadamard.py,sha256=oLdDUu1p82lgD7li-sHMSvXZxz1SDjLeYf-EfXqNzvk,3918
-compressed_tensors/transform/factory/matrix_multiply.py,sha256=KYiQRGFSU33TpPWkGTKwNADTmYoU0E3hjQypOMclHbg,3689
+compressed_tensors/transform/factory/base.py,sha256=w9ic5eSxfSNn2Xju-xZvG4_iXAIsJCU56qik8w---aI,5994
+compressed_tensors/transform/factory/hadamard.py,sha256=iJ2OyKitR2Duw0z5Jqj69GTih2C1WtHRXQCTtATaTtw,4180
+compressed_tensors/transform/factory/matrix_multiply.py,sha256=LdoV2E12HTucmUWcw7UKOpRNnL8QhOOIUnNVlpOpGiI,3925
 compressed_tensors/transform/factory/random_hadamard.py,sha256=nUhTlFa4ikSpcl4Umme71pnjMPgwYoGlwjKlU27UHZ4,1634
 compressed_tensors/transform/utils/__init__.py,sha256=fH6rjBYAxuwrTzBTlTjTgCYNyh6TCvCqajCz4Im4YrA,617
-compressed_tensors/transform/utils/hadamard.py,sha256=U27Kvo-eDebKcVt8oXTSIAaQ5DvPQj9tDv2hdXHCPPQ,5584
+compressed_tensors/transform/utils/hadamard.py,sha256=hDJZC0Gw2fKdxqa3f8TmFc5J0eJqxHtFRxswLU_yVJc,5548
 compressed_tensors/transform/utils/hadamards.safetensors,sha256=mFd1GzNodGG-ifA1IoH-0nHYzfraCOvrq_dX2zFI1B4,1436901
-compressed_tensors/transform/utils/utils.py,sha256=PRPTYwPs2nnNaQMq2GEbC4QYKHFKlZwaRyPgdDhl66g,2992
+compressed_tensors/transform/utils/matrix.py,sha256=FIHCUlpWVIIhdr3c6EbQec41JeiPAAjCM9Ejz77wb-w,6181
 compressed_tensors/utils/__init__.py,sha256=QFQzF6MpV3yStajPzYktZkmvZsxvfpKUZq2oGbd1Cvw,832
 compressed_tensors/utils/helpers.py,sha256=Q3iRAa2XSdmmn4vSpUplnvKOmWwn4Clao9ZkPBHXtpI,12604
 compressed_tensors/utils/internal.py,sha256=7SSWgDoNFRnlfadwkoFhLW-T2jOc7Po_WzWv5h32Sa8,982
@@ -61,8 +61,8 @@ compressed_tensors/utils/permutations_24.py,sha256=kx6fsfDHebx94zsSzhXGyCyuC9sVy
 compressed_tensors/utils/permute.py,sha256=V6tJLKo3Syccj-viv4F7ZKZgJeCB-hl-dK8RKI_kBwI,2355
 compressed_tensors/utils/safetensors_load.py,sha256=DMfZBuUbA6qp_BG_zIWT3ckiEE33K9ob34s-OgzReO4,12057
 compressed_tensors/utils/semi_structured_conversions.py,sha256=XKNffPum54kPASgqKzgKvyeqWPAkair2XEQXjkp7ho8,13489
-compressed_tensors-0.10.3a20250715.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors-0.10.3a20250715.dist-info/METADATA,sha256=Pb3ZHNAmiFtJHKdWbAA7m2KvmSttY8EV2huSI71ZBGM,7031
-compressed_tensors-0.10.3a20250715.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-compressed_tensors-0.10.3a20250715.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors-0.10.3a20250715.dist-info/RECORD,,
+compressed_tensors-0.10.3a20250716.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors-0.10.3a20250716.dist-info/METADATA,sha256=zWpNTNT2bbSxFEQ58pvjEqdAE8MzBBZhvvtVVh6AQ14,7031
+compressed_tensors-0.10.3a20250716.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+compressed_tensors-0.10.3a20250716.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.10.3a20250716.dist-info/RECORD,,

compressed_tensors/transform/utils/utils.py

@@ -1,91 +0,0 @@
-# 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 torch
-from compressed_tensors.transform import TransformLocation
-
-
-__all__ = ["get_matrix_size", "apply_transform_weight"]
-
-
-def get_matrix_size(module: torch.nn.Module, location: TransformLocation) -> int:
-    """
-    Determine the size of a matrix given its location on the module
-
-    :param module: module that matrix will be applied to
-    :param location: location on module
-    :return: size of matrix
-    """
-    assert isinstance(module, torch.nn.Linear)
-    if location in ("input", TransformLocation.WEIGHT_INPUT):
-        return module.in_features
-    else:
-        return module.out_features
-
-
-def apply_transform_weight(
-    weight: torch.Tensor,
-    value: torch.Tensor,
-    location: TransformLocation,
-) -> torch.Tensor:
-    """
-    Using the transform location, determine how to apply the transform weight to the
-    given value. For more info on input and output transforms, see `TransformLocation`
-
-    The following explains how weights should be applied to values according to location
-
-    let  x          be input activation
-         W          be weight,
-         yh, xh, Wh be transformed output, input, weight
-
-    note that
-         y  = (x W.T)        // torch.nn.Linear
-
-    Choose values for yh, xh, and Wh which incorporate matrix transforms
-
-    let  V, Vi      be transform matrices on input side
-         U, Ui      be transform matrices on output side
-
-    pick xh = (x V)
-         Wh = (U.T W Vi.T)
-         yh = (y U)
-
-    The following shows that `yh = (xh) (Wh).T` for the chosen values of yh, xh, and Wh
-
-    (xh) (Wh).T = (x V) (U.T W Vi.T).T
-                = (x V) (Vi W.T U)        // transpose matrix product identity
-                = (x W.T) U
-                = y U
-                = yh
-
-    :param weight: transform weight to apply
-    :param value: value to apply weight to
-    :param location: determines how weight should be applied
-    :return: value after transform weight has been applied
-    """
-
-    if location == TransformLocation.INPUT:
-        return value @ weight
-
-    elif location == TransformLocation.WEIGHT_INPUT:
-        return value @ weight.T
-
-    elif location == TransformLocation.WEIGHT_OUTPUT:
-        return weight.T @ value
-
-    elif location == TransformLocation.OUTPUT:
-        return value @ weight
-
-    else:
-        raise NotImplementedError(f"{location} has not been implemented yet")