[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
 
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
@@ -22,7 +22,7 @@
     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):
         """
@@ -50,30 +51,55 @@ 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)
+        is_linear = 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, is_linear)
 
-    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,
+        is_linear: bool,
+    ):
         super().__init__()
         self.weight = weight
+        self.perm = perm
         self.args = args
+        self.is_linear = is_linear
 
     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)
+        return apply_transform_weight(weight, value, self.args.location, self.is_linear)

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

@@ -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)

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)
 

src/compressed_tensors/transform/utils/utils.py

@@ -27,17 +27,28 @@ def get_matrix_size(module: torch.nn.Module, location: TransformLocation) -> int
     :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
+    if isinstance(module, torch.nn.Linear):
+        if location in ("input", TransformLocation.WEIGHT_INPUT):
+            return module.in_features
+        else:
+            return module.out_features
+    elif isinstance(module, torch.nn.Embedding):
+        if location in ("input", TransformLocation.WEIGHT_INPUT):
+            return module.num_embeddings
+        else:
+            return module.embedding_dim
+
+    raise ValueError(
+        f"Unsupported module type {type(module)}, "
+        "should be either Linear or Embedding."
+    )
 
 
 def apply_transform_weight(
-    weight: torch.Tensor,
+    transform_weight: torch.Tensor,
     value: torch.Tensor,
     location: TransformLocation,
+    is_linear: bool = True,
 ) -> torch.Tensor:
     """
     Using the transform location, determine how to apply the transform weight to the
@@ -69,23 +80,36 @@ def apply_transform_weight(
                 = 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
+    :param transform_weight: transform weight to apply
+    :param value: value to apply transform_weight to
+    :param location: determines how transform_weight should be applied
+    :param is_linear: if value belongs to the weights of a Linear module
+        This is needed because torch uses convention:
+        Linear(in_features,out_features) has weight shape (out_features, in_features)
+        But other modules (e.g. torch.nn.Embedding) don't:
+        Embedding(num_embeddings, embedding_dim) has weight shape
+         (num_embeddings, embedding_dim)
+    :return: value after transform_weight has been applied
     """
 
     if location == TransformLocation.INPUT:
-        return value @ weight
+        return value @ transform_weight
 
     elif location == TransformLocation.WEIGHT_INPUT:
-        return value @ weight.T
+        if is_linear:
+            return value @ transform_weight.T
+        else:
+            # TODO is this ever needed?
+            raise NotImplementedError()
 
     elif location == TransformLocation.WEIGHT_OUTPUT:
-        return weight.T @ value
+        if is_linear:
+            return transform_weight.T @ value
+        else:
+            return value @ transform_weight
 
     elif location == TransformLocation.OUTPUT:
-        return value @ weight
+        return value @ transform_weight
 
     else:
         raise NotImplementedError(f"{location} has not been implemented yet")

src/compressed_tensors/utils/helpers.py

@@ -373,11 +373,16 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
+        self._kwargs = {}
 
-    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._kwargs)
         else:
-            value = self.default_factory(key)
+            value = self.default_factory(key, **self._kwargs)
         self[key] = value
         return value
+
+    def get(self, *args, **kwargs) -> Any:
+        with patch_attr(self, "_kwargs", kwargs):
+            return self[args]