[Transform] Support separating v and u transforms of quip (#1782)

## Purpose ##
* Support `rotation` argument with `QuIPModifier` which allows users to
select V and U rotations separately.
* According to research, the default can be changed to V after a basic
recovery evaluation is done

## Changes ##
* Add `rotations` argument
* Updated docstrings and comments to remove references to spinquant

## Testing ##
* Tested model coherence with V and U separately

---------

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

Author: kylesayrs

examples/transform/quip_example.py

@@ -18,10 +18,10 @@
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 # Configure the quantization algorithm to run.
-#   * apply spinquant transforms to model in order to make quantization easier
+#   * apply quip transforms to model in order to make quantization easier
 #   * quantize the weights to 4 bit with a group size 128
 recipe = [
-    QuIPModifier(targets="Linear", transform_type="random-hadamard"),
+    QuIPModifier(rotations=["v", "u"], transform_type="random-hadamard"),
     QuantizationModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]),
 ]
 
@@ -35,7 +35,7 @@
 input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to(
     model.device
 )
-output = model.generate(input_ids, max_new_tokens=100)
+output = model.generate(input_ids, max_new_tokens=50)
 print(tokenizer.decode(output[0]))
 print("==========================================\n\n")
 

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

@@ -32,18 +32,18 @@ class QuIPModifier(Modifier):
 
     Lifecycle:
         - on_initialize
-            - infer SpinQuantMappings & NormMappings
-            - as needed, create transform schemes for R1, R2, R3, & R4
+            - as needed, create transform schemes for V (input) and U (output)
         - on_start
-            - normalize embeddings
-            - fuse norm layers into subsequent Linear layers
             - apply TransformConfig
                 - fuse transforms into weights for mergeable transforms
                 - add hooks for online transforms
         - on sequential epoch end
         - on_end
         - on_finalize
 
+    :param rotations: which rotation schemes to apply to the model. Including `"v"` will
+        rotate the input side of weights, and including `"u"` will rotate the output
+        side of weights (note that v does not require u and vice-versa)
     :param transform_type: The type of transform to apply to the model.
         `"hadamard"` has the least performance cost but only supports sizes which are
         powers of power of two.
@@ -58,6 +58,7 @@ class QuIPModifier(Modifier):
     :param transform_config: Optional transform config for overriding provided arguments
     """  # noqa: E501
 
+    rotations: List[Literal["v", "u"]] = Field(default_factory=lambda: ["v", "u"])
     transform_type: Literal["hadamard", "random-hadamard", "random-matrix"] = Field(
         default="random-hadamard"
     )
@@ -77,6 +78,12 @@ def validate_not_implemented(cls, value, info: ValidationInfo):
             raise NotImplementedError(f"{info.field_name} is not supported right now")
         return value
 
+    @field_validator("rotations", mode="before")
+    def validate_lowercase_list(cls, value):
+        if isinstance(value, list):
+            value = [v.lower() if isinstance(v, str) else v for v in value]
+        return value
+
     def on_initialize(self, state: State, **kwargs) -> bool:
         if self.transform_config is not None:
             return True
@@ -111,45 +118,52 @@ def on_finalize(self, state: State, **kwargs) -> bool:
         return True
 
     def _create_config(self) -> TransformConfig:
-        return TransformConfig(
-            config_groups={
-                "v": TransformScheme(
-                    type=self.transform_type,
-                    apply=[
-                        TransformArgs(
-                            targets=self.targets,
-                            location="input",  # non-mergable
-                            ignore=self.ignore,
-                        ),
-                        TransformArgs(
-                            targets=self.targets,
-                            location="weight_input",
-                            inverse=True,
-                            ignore=self.ignore,
-                        ),
-                    ],
-                    randomize=self.randomize,
-                    requires_grad=self.learnable,
-                    precision=self.precision,
+        config_groups = dict()
+        if "v" in self.rotations:
+            config_groups["v"] = self._create_v_scheme()
+        if "u" in self.rotations:
+            config_groups["u"] = self._create_u_scheme()
+
+        return TransformConfig(config_groups=config_groups)
+
+    def _create_v_scheme(self) -> TransformScheme:
+        return TransformScheme(
+            type=self.transform_type,
+            apply=[
+                TransformArgs(
+                    targets=self.targets,
+                    location="input",  # non-mergable
+                    ignore=self.ignore,
+                ),
+                TransformArgs(
+                    targets=self.targets,
+                    location="weight_input",
+                    inverse=True,
+                    ignore=self.ignore,
+                ),
+            ],
+            randomize=self.randomize,
+            requires_grad=self.learnable,
+            precision=self.precision,
+        )
+
+    def _create_u_scheme(self) -> TransformScheme:
+        return TransformScheme(
+            type=self.transform_type,
+            apply=[
+                TransformArgs(
+                    targets=self.targets,
+                    location="weight_output",
+                    ignore=self.ignore,
                 ),
-                "u": TransformScheme(
-                    type=self.transform_type,
-                    apply=[
-                        TransformArgs(
-                            targets=self.targets,
-                            location="weight_output",
-                            ignore=self.ignore,
-                        ),
-                        TransformArgs(
-                            targets=self.targets,
-                            location="output",  # non-mergable
-                            inverse=True,
-                            ignore=self.ignore,
-                        ),
-                    ],
-                    randomize=self.randomize,
-                    requires_grad=self.learnable,
-                    precision=self.precision,
+                TransformArgs(
+                    targets=self.targets,
+                    location="output",  # non-mergable
+                    inverse=True,
+                    ignore=self.ignore,
                 ),
-            }
+            ],
+            randomize=self.randomize,
+            requires_grad=self.learnable,
+            precision=self.precision,
         )