[Transforms] apply transforms to torch.nn.Embedding modules

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -152,11 +152,7 @@ def apply_quantization_config(
     # list of submodules to ignore
     ignored_submodules = defaultdict(list)
     # mark appropriate layers for quantization by setting their quantization schemes
-    for name, submodule in iter_named_quantizable_modules(
-        model,
-        include_children=True,
-        include_attn=True,
-    ):  # child modules and attention modules
+    for name, submodule in model.named_modules():  # child modules and attention modules
         # potentially fix module name to remove FSDP wrapper prefix
         name = fix_fsdp_module_name(name)
         if matches := find_name_or_class_matches(name, submodule, config.ignore):

src/compressed_tensors/transform/__init__.py

@@ -23,3 +23,4 @@
 from .factory.hadamard import *
 from .factory.matrix_multiply import *
 from .factory.random_hadamard import *
+from .apply import *

src/compressed_tensors/transform/apply.py

@@ -0,0 +1,25 @@
+# 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 TransformConfig, TransformFactory
+
+
+__all__ = ["apply_transform_config"]
+
+
+def apply_transform_config(model: torch.nn.Module, config: TransformConfig):
+    for name, scheme in config.config_groups.items():
+        factory = TransformFactory.from_scheme(scheme, name=name)
+        factory.apply_to_model(model)

src/compressed_tensors/transform/factory/base.py

@@ -99,10 +99,10 @@ def _apply_to_module(self, module: Module, args: TransformArgs):
         # 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:
+            register_offload_module(module, transform_name, transform)
 
             def input_hook(_, args):
                 input = args[0]
@@ -115,8 +115,8 @@ def input_hook(_, args):
             TransformLocation.WEIGHT_INPUT,
             TransformLocation.WEIGHT_OUTPUT,
         ):
-            assert isinstance(module, torch.nn.Linear)
-            assert module.bias is None
+            assert isinstance(module, (torch.nn.Linear, torch.nn.Embedding))
+            assert not hasattr(module, "bias") or module.bias is None
 
             with torch.no_grad(), align_module_device(module):
                 update_offload_parameter(module, "weight", transform(module.weight))
@@ -130,6 +130,7 @@ def input_hook(_, args):
 
         # register output transformation hook
         elif args.location == TransformLocation.OUTPUT:
+            register_offload_module(module, transform_name, transform)
 
             def output_hook(_, _input, output):
                 return transform(output)
@@ -140,9 +141,6 @@ def output_hook(_, _input, output):
         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):
     """

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, Type
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
@@ -22,10 +22,10 @@
     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
+from torch.nn import Module, Parameter
 
 
 @TransformFactory.register("hadamard")
@@ -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):
         """
@@ -50,30 +51,70 @@ def create_transform(self, module: Module, args: TransformArgs):
         :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)
+        assert hasattr(module, "weight")
         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)
+        weight = self.weights.get(size, dtype, device, construct_device=exec_device)
+        perm = self.perms[weight] if self.scheme.randomize else None
+        return HadamardTransform(weight, perm, args, type(module))
 
-    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
+        if self.scheme.num_heads is None or self.scheme.num_heads <= 1:
+            data = deterministic_hadamard_matrix(size, dtype, construct_device)
+        else:
+            assert size % self.scheme.num_heads == 0
+            data = torch.kron(
+                torch.eye(self.scheme.num_heads, dtype=dtype),
+                deterministic_hadamard_matrix(
+                    self.scheme.head_dim, 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,
+        module_type: Type,
+    ):
         super().__init__()
         self.weight = weight
+        self.perm = perm
         self.args = args
+        self.module_type = module_type
 
     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)
+        # NOTE: SpinQuant code up-converts to float64 for application
+        # of transform, then down-converts
+        return apply_transform_weight(
+            weight.to(torch.float64),
+            value.to(torch.float64),
+            self.args.location,
+            self.module_type,
+        ).to(weight.dtype)

src/compressed_tensors/transform/factory/matrix_multiply.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, Type
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
@@ -59,7 +59,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         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:
         data = torch.rand(
@@ -73,17 +73,22 @@ def _create_inverse(self, weight: Parameter) -> Parameter:
 
 
 class RandomMatrixTransform(TransformBase):
-    def __init__(self, weight: Tensor, args: TransformArgs):
+    def __init__(self, weight: Tensor, args: TransformArgs, module_type: Type):
         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:

src/compressed_tensors/transform/factory/random_hadamard.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
 from compressed_tensors.transform import HadamardFactory, TransformFactory
 from compressed_tensors.transform.utils.hadamard import random_hadamard_matrix
 from torch import device, dtype
@@ -28,7 +29,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)

src/compressed_tensors/transform/transform_config.py

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

src/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
@@ -31,13 +31,14 @@ 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)
+    # TODO infer num_heads
+    num_heads: Optional[int] = None

src/compressed_tensors/transform/utils/hadamard.py

@@ -51,7 +51,9 @@ def deterministic_hadamard_matrix(
 
     log2 = int(math.log2(size))
     if size != 2**log2:
-        raise ValueError("Cannot construct deterministic hadamard of size != 2^n")
+        raise ValueError(
+            f"Cannot construct deterministic hadamard of size {size} != 2^n"
+        )
 
     H = torch.tensor([[1]], dtype=dtype, device=device)