refactor R1 rotation and partial_hadamard

Author: carrot-o0o

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -266,33 +266,21 @@ def _create_r1_scheme(self) -> TransformScheme:
                 location="weight_output",
             )
         )
-        if getattr(self.mappings, "attn_v_is_kv_combined", False):
-            apply_list.append(
-                TransformArgs(
-                    targets=[
-                        self.mappings.attn_q,
-                        self.mappings.attn_k,
-                        *self.mappings.mlp_in,
-                        self.mappings.lm_head,
-                    ],
-                    location="weight_input",
-                    inverse=True,
-                )
-            )
-        else:
-            apply_list.append(
-                TransformArgs(
-                    targets=[
-                        self.mappings.attn_q,
-                        self.mappings.attn_k,
-                        self.mappings.attn_v,
-                        *self.mappings.mlp_in,
-                        self.mappings.lm_head,
-                    ],
-                    location="weight_input",
-                    inverse=True,
-                )
+        r1_input_targets = [
+            self.mappings.attn_q,
+            self.mappings.attn_k,
+            *self.mappings.mlp_in,
+            self.mappings.lm_head,
+        ]
+        if not getattr(self.mappings, "attn_v_is_kv_combined", False):
+            r1_input_targets.append(self.mappings.attn_v)
+        apply_list.append(
+            TransformArgs(
+                targets=r1_input_targets,
+                location="weight_input",
+                inverse=True,
             )
+        )
         return TransformScheme(
             type=self.transform_type,
             randomize=self.randomize,

src/llmcompressor/modifiers/transform/spinquant/partial_hadamard.py

@@ -96,14 +96,8 @@ def __init__(
         v_head_dim: int = 0,
     ):
         super().__init__(weight, perm, scheme, args, module_type)
-        self.weight = weight
-        self.perm = perm
-        self.scheme = scheme
-        self.args = args
-        self.module_type = module_type
         self.qk_nope_head_dim = qk_nope_head_dim
         self.v_head_dim = v_head_dim
-        self._scale = torch.tensor(weight.size(0), dtype=torch.float64).sqrt()
 
     def forward(self, value: Tensor) -> Tensor:
         weight = self.weight
@@ -154,6 +148,7 @@ def apply_partial_transform_weight(
     :return: value after transform_weight has been applied
     """
     assert transform_weight.shape[0] == transform_weight.shape[1]
+    assert qk_nope_head_dim > 0 and v_head_dim > 0
     if TransformLocation(location).is_online():
         return _multihead_matmul(value, transform_weight)