implement num_heads

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -18,7 +18,7 @@
 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,
 )
@@ -52,15 +52,16 @@ def create_transform(self, module: Module, args: TransformArgs):
         :param args: defines how the transform will be applied to the module
         """
         assert isinstance(module, Linear)
-        size = get_matrix_size(module, args.location)
+        num_heads = self.scheme.num_heads
+        size = get_matrix_size(module, args.location, num_heads)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
 
         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, num_heads)
 
     def _create_weight(
         self,
@@ -81,12 +82,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: Optional[Parameter],
+        args: TransformArgs,
+        num_heads: Optional[int],
     ):
         super().__init__()
         self.weight = weight
         self.perm = perm
         self.args = args
+        self.num_heads = num_heads
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -97,4 +103,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.num_heads)

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -17,7 +17,7 @@
 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,
 )
@@ -51,15 +51,16 @@ def create_transform(self, module: Module, args: TransformArgs):
         :param args: defines how the transform will be applied to the module
         """
         assert isinstance(module, Linear)
-        size = get_matrix_size(module, args.location)
+        num_heads = self.scheme.num_heads
+        size = get_matrix_size(module, args.location, num_heads)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
 
         weight = self.weights[size, dtype, device]
         if args.inverse:
             weight = self.inverses[weight]
 
-        return RandomMatrixTransform(weight, args)
+        return RandomMatrixTransform(weight, args, num_heads)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
         # TODO: verify that weight is invertible (has non-zero determinant)
@@ -74,17 +75,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, num_heads: Optional[int]):
         super().__init__()
         self.weight = weight  # is an inverse if args.inverse
         self.args = args
+        self.num_heads = num_heads
 
     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.num_heads
+        )
 
     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.num_heads
+        )
 
 
 def high_precision_invert(weight: Tensor) -> Tensor:

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
@@ -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)
+    num_heads: Optional[int] = Field(default=None)

src/compressed_tensors/transform/utils/utils.py renamed to src/compressed_tensors/transform/utils/matrix.py

@@ -12,14 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 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:
+def get_matrix_size(
+    module: torch.nn.Module,
+    location: TransformLocation,
+    num_heads: Optional[int] = None,
+) -> int:
     """
     Determine the size of a matrix given its location on the module
 
@@ -28,17 +34,36 @@ def get_matrix_size(module: torch.nn.Module, location: TransformLocation) -> int
     :return: size of matrix
     """
     assert isinstance(module, torch.nn.Linear)
+
     if location in ("input", TransformLocation.WEIGHT_INPUT):
-        return module.in_features
+        size = module.in_features
     else:
-        return module.out_features
+        size = module.out_features
+
+    if num_heads is not None:
+        assert size % num_heads == 0
+        size = size // num_heads
+
+    return size
 
 
 def apply_transform_weight(
     weight: torch.Tensor,
     value: torch.Tensor,
     location: TransformLocation,
+    num_heads: Optional[int] = None,
 ) -> torch.Tensor:
+    if num_heads is not None:
+        weight = weight.repeat((num_heads, num_heads))
+
+    return apply_transform_weight_linear(weight, value, location)
+
+
+def apply_transform_weight_linear(
+    weight: torch.Tensor,
+    value: torch.Tensor,
+    location: TransformLocation,
+):
     """
     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`
@@ -74,7 +99,6 @@ def apply_transform_weight(
     :param location: determines how weight should be applied
     :return: value after transform weight has been applied
     """
-
     if location == TransformLocation.INPUT:
         return value @ weight