utils to apply transforms to torch.nn.Embedding modules

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

src/compressed_tensors/transform/factory/base.py

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

src/compressed_tensors/transform/factory/hadamard.py

@@ -51,15 +51,16 @@ 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.get(size, dtype, device, construct_device=exec_device)
         perm = self.perms[weight] if self.scheme.randomize else None
-        return HadamardTransform(weight, perm, args)
+        return HadamardTransform(weight, perm, args, is_linear)
 
     def _create_weight(
         self,
@@ -80,12 +81,17 @@ def _create_permutation(self, weight: Parameter) -> Parameter:
 
 class HadamardTransform(TransformBase):
     def __init__(
-        self, weight: Parameter, perm: Union[Parameter, None], args: TransformArgs
+        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:
         weight = self.weight
@@ -96,4 +102,4 @@ def forward(self, value: Tensor) -> Tensor:
         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/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")