add rotation args

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

Author: kylesayrs

examples/transform/quip_example.py

@@ -14,14 +14,14 @@
 # NOTE: because the datafree pipeline is being used in this
 # example, you can use additional GPUs to support larger models
 MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype="auto")
 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

@@ -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,
         )