[Transform] Hadamard Permutations (#329)

* add utilities

Signed-off-by: Kyle Sayers <[email protected]>

* add tests

Signed-off-by: Kyle Sayers <[email protected]>

* add additional tests

Signed-off-by: Kyle Sayers <[email protected]>

* add utils and tests

Signed-off-by: Kyle Sayers <[email protected]>

* Implement transform factories

Signed-off-by: Kyle Sayers <[email protected]>

* add permutations

Signed-off-by: Kyle Sayers <[email protected]>

* add delete_offload_module

Signed-off-by: Kyle Sayers <[email protected]>

* key inverses by weight

Signed-off-by: Kyle Sayers <[email protected]>

* fix tests

Signed-off-by: Kyle Sayers <[email protected]>

* standardize random hadamard

Signed-off-by: Kyle Sayers <[email protected]>

* prepend input hooks

Signed-off-by: Kyle Sayers <[email protected]>

* apply sqrt division first

Signed-off-by: Kyle Sayers <[email protected]>

* use divided hadamards

Signed-off-by: Kyle Sayers <[email protected]>

* fix typo

Signed-off-by: Kyle Sayers <[email protected]>

* add random option

Signed-off-by: Kyle Sayers <[email protected]>

* use random seeds, rename matrix multiply

Signed-off-by: Kyle Sayers <[email protected]>

* add deterministic generation to random matrix

Signed-off-by: Kyle Sayers <[email protected]>

* fix perm math

Signed-off-by: Kyle Sayers <[email protected]>

* update docstrings

Signed-off-by: Kyle Sayers <[email protected]>

* update docstrings

Signed-off-by: Kyle Sayers <[email protected]>

* cleanup

Signed-off-by: Kyle Sayers <[email protected]>

* cleanup 2

Signed-off-by: Kyle Sayers <[email protected]>

* make seed optional

Signed-off-by: Kyle Sayers <[email protected]>

* remove iterable check and missing return value

Signed-off-by: Kyle Sayers <[email protected]>

* Remove unrelated changes

* simplify code

Signed-off-by: Kyle Sayers <[email protected]>

* fix style

Signed-off-by: Kyle Sayers <[email protected]>

---------

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

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
@@ -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):
         """
@@ -56,24 +57,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[weight] if self.scheme.randomize else None
+        return HadamardTransform(weight, perm, args)
 
     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)
         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
+    ):
         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
+        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)

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)

tests/test_transform/test_transform_scheme.py

@@ -24,7 +24,7 @@ def test_basic_scheme():
         type="hadamard",
         apply=[basic_args],
     )
-    assert not scheme.randomize_modules
+    assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 1
     assert isinstance(scheme.apply[0], TransformArgs)
@@ -43,10 +43,10 @@ def test_multiple_groups_global():
     scheme = TransformScheme(
         type="hadamard",
         apply=[embedding_args, linear_args],
-        randomize_modules=True,
+        randomize=True,
     )
 
-    assert scheme.randomize_modules
+    assert scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 2
     assert isinstance(scheme.apply[0], TransformArgs)
@@ -69,6 +69,6 @@ def test_multiple_groups():
         apply=apply,
     )
 
-    assert not scheme.randomize_modules
+    assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 20