add permutations

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

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Union
 
 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 (
+    apply_permutation,
     apply_transform_weight,
     get_matrix_size,
 )
@@ -41,6 +42,7 @@ class HadamardFactory(TransformFactory):
     def __init__(self, name: str, scheme: TransformScheme, seed: int = 42):
         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):
         """
@@ -56,24 +58,35 @@ def create_transform(self, module: Module, args: TransformArgs):
         device = get_offloaded_device(module)
 
         weight = self.weights[size, dtype, device]
-        return HadamardTransform(weight, args)
+        perm = self.perms[module, weight] if self.scheme.randomize_modules else None
+        return HadamardTransform(weight, perm, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
         data = torch.tensor(deterministic_hadamard_matrix(size))  # TODO: seed=self.seed
         data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 
+    def _create_permutation(self, module: Module, weight: Parameter) -> Parameter:
+        data = torch.randperm(weight.size(0))
+        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
+    ):
         super().__init__()
         self.weight = weight
+        self.perm = perm
         self.args = args
 
     def forward(self, value: Tensor) -> Tensor:
-        if not self.args.inverse:
-            weight = self.weight
-        else:
-            weight = self.weight.T / self.weight.size(0)
+        weight = self.weight
+
+        if self.perm is not None:
+            weight = apply_permutation(weight, self.perm)
+
+        if self.args.inverse:
+            weight = weight.T / weight.size(0)
 
         return apply_transform_weight(weight, value, self.args.location)

src/compressed_tensors/transform/utils/utils.py

@@ -16,7 +16,7 @@
 from compressed_tensors.transform import TransformLocation
 
 
-__all__ = ["get_matrix_size", "apply_transform_weight"]
+__all__ = ["get_matrix_size", "apply_transform_weight", "apply_permutation"]
 
 
 def get_matrix_size(module: torch.nn.Module, location: TransformLocation) -> int:
@@ -83,3 +83,10 @@ def apply_transform_weight(
 
     elif location == TransformLocation.OUTPUT:
         return value @ weight
+
+
+def apply_permutation(weight: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
+    weight = weight.clone()
+    diag_indices = torch.arange(weight.size(0))
+    weight[diag_indices, diag_indices] = weight.diagonal()[perm]
+    return weight

tests/test_transform/factory/test_correctness.py

@@ -19,10 +19,18 @@
     TransformFactory,
     TransformScheme,
 )
-from compressed_tensors.utils import align_modules, force_cpu_offload
+from compressed_tensors.utils import force_cpu_offload
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 
+_test_schemes = [
+    TransformScheme(type=name) for name in TransformFactory.registered_names()
+] + [
+    TransformScheme(type=name, randomize_modules=True)
+    for name in TransformFactory.registered_names()
+]
+
+
 class TransformableModel(torch.nn.Module):
     def __init__(self, *sizes):
         super().__init__()
@@ -37,10 +45,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_correctness_linear(scheme):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=True)
@@ -68,10 +73,7 @@ def test_correctness_linear(scheme):
     torch.allclose(true_output, output, atol=1e-7, rtol=0.0)
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_correctness_model(scheme, offload=False):
     # load model
     model = TransformableModel(2, 4, 8, 16)
@@ -99,9 +101,6 @@ def test_correctness_model(scheme, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_correctness_model_offload(scheme):
     test_correctness_model(scheme, offload=True)

tests/test_transform/factory/test_memory.py

@@ -26,6 +26,14 @@
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 
+_test_schemes = [
+    TransformScheme(type=name) for name in TransformFactory.registered_names()
+] + [
+    TransformScheme(type=name, randomize_modules=True)
+    for name in TransformFactory.registered_names()
+]
+
+
 class TransformableModel(torch.nn.Module):
     def __init__(self, *sizes):
         super().__init__()
@@ -40,10 +48,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_memory_sharing(scheme, offload=False):
     # load scheme and factory
     scheme = TransformScheme(
@@ -93,20 +98,12 @@ def test_memory_sharing(scheme, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize(
-    "scheme",
-    [TransformScheme(type=name) for name in TransformFactory.registered_names()],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_memory_sharing_offload(scheme):
     test_memory_sharing(scheme, offload=True)
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [
-        TransformScheme(type=name, requires_grad=True)
-        for name in TransformFactory.registered_names()
-    ],
-)
+@pytest.mark.parametrize("scheme", _test_schemes)
 def test_memory_sharing_training(scheme):
+    scheme.requires_grad = True
     test_memory_sharing(scheme, offload=False)