[transforms] SpinQuantModifier & QuIPModifier transform_block_size field (#1806)

SUMMARY:
Resolves `INFERENG-1882`

The research community [has pointed
out](https://github.com/IST-DASLab/FP-Quant?tab=readme-ov-file#fp-format-quantization-harness)
that the rotation/transform block size is important when performing
transforms:

> Key to efficiency is that the Hadamard block size matches the
microscaling format group size (16 or 32)

This exposes a new field on SpinQuantModifier and QuIPModifier to allow
the user to set it to an arbitrary value, as long as the model's
hidden_size and head_dim are both evenly divisible by it.

- [x] Add to SpinQuant Modifier. Option to allow for different
`transform_block_size`s for R1 vs. R2 can be added at a future time.
- [x] Add to QuIPModifier. Option to allow for different
`transform_block_size`s for U vs. V can be added at a future time.

Merge in conjunction with:
 * neuralmagic/compressed-tensors#466

TEST PLAN:
`transform_block_size` added to parameterized
`tests/llmcompressor/modifiers/transform/(test_correctness.py|test_serialization.py)`

---------

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

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

@@ -54,6 +54,11 @@ class QuIPModifier(Modifier):
     :param learnable: If true, attach gradients to transform weights for training
     :param precision: Precision at which all transforms should be applied. This applies
         to both weight fusing and online rotations
+    :param transform_block_size: Block size to use for rotation matrices. The model's
+        hidden_size must be evenly divisible by transform_block_size.
+        Layers will be transformed by a block-diagonal matrix where each block is a
+        matrix of this size.
+        If None is provided, model's hidden_size will be used
     :param ignore: Modules to ignore when attaching transforms
     :param transform_config: Optional transform config for overriding provided arguments
     """  # noqa: E501
@@ -66,6 +71,7 @@ class QuIPModifier(Modifier):
     randomize: bool = Field(default=False)
     learnable: bool = Field(default=False)
     precision: TorchDtype = Field(default=torch.float64)
+    transform_block_size: Optional[int] = Field(default=None)
     ignore: Union[str, List[str]] = Field(default="lm_head")
 
     # optional override for more fine-grained control
@@ -129,6 +135,7 @@ def _create_config(self) -> TransformConfig:
     def _create_v_scheme(self) -> TransformScheme:
         return TransformScheme(
             type=self.transform_type,
+            block_size=self.transform_block_size,
             apply=[
                 TransformArgs(
                     targets=self.targets,
@@ -150,6 +157,7 @@ def _create_v_scheme(self) -> TransformScheme:
     def _create_u_scheme(self) -> TransformScheme:
         return TransformScheme(
             type=self.transform_type,
+            block_size=self.transform_block_size,
             apply=[
                 TransformArgs(
                     targets=self.targets,

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

@@ -69,6 +69,12 @@ class SpinQuantModifier(Modifier, use_enum_values=True):
     :param learnable: if True, attach gradients to transform weights for training
     :param precision: Precision at which all transforms should be applied. This applies
         to both weight fusing and online rotations
+    :param transform_block_size: Block size to use for rotation matrices. The model's
+        hidden_size and head_dim must be evenly divisible by transform_block_size.
+        Layers will be transformed by a block-diagonal matrix where each block is a
+        matrix of this size.
+        If None is provided, model's hidden_size will be used for R1, R3, and R4
+        and model's head_dim will be used for R2
     :param mappings: Specifies layers within a model to target for transforms.
         A mapping will be inferred if None is provided
     :param norm_mappings: Specifies layers within a model to target for norm fusing.
@@ -83,6 +89,7 @@ class SpinQuantModifier(Modifier, use_enum_values=True):
     randomize: bool = Field(default=False)
     learnable: bool = Field(default=False)
     precision: TorchDtype = Field(default=torch.float64)
+    transform_block_size: Optional[int] = Field(default=None)
 
     # norm mappings separate from spinquant mappings to allow users to
     # override spinquant mappings with transform_config without overriding norms
@@ -186,6 +193,7 @@ def _create_r1_scheme(self) -> TransformScheme:
             randomize=self.randomize,
             requires_grad=self.learnable,
             precision=self.precision,
+            block_size=self.transform_block_size,
             apply=[
                 TransformArgs(
                     targets=[
@@ -219,12 +227,20 @@ def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
         else:
             raise NotImplementedError()
 
+        if self.transform_block_size:
+            if head_dim % self.transform_block_size != 0:
+                raise ValueError(
+                    f"transform_block_size {self.transform_block_size} must be set "
+                    f"such that model's head_dim {head_dim} is evenly divisible by it"
+                )
+            head_dim = self.transform_block_size
+
         return TransformScheme(
             type=self.transform_type,
             randomize=self.randomize,
             requires_grad=self.learnable,
             precision=self.precision,
-            head_dim=head_dim,
+            block_size=head_dim,
             apply=[
                 TransformArgs(targets=[self.mappings.attn_v], location="weight_output"),
                 TransformArgs(
@@ -246,6 +262,7 @@ def _create_r4_scheme(self) -> TransformScheme:
             randomize=self.randomize,
             requires_grad=self.learnable,
             precision=self.precision,
+            block_size=self.transform_block_size,
             apply=[
                 TransformArgs(
                     targets=[*self.mappings.mlp_out],

tests/llmcompressor/modifiers/transform/test_correctness.py

@@ -18,26 +18,32 @@
     reason="Skipping correctness tests requiring gated model access",
 )
 @pytest.mark.parametrize(
-    "modifier,model_dtype,precision,exp_mse",
+    "modifier,model_dtype,precision,transform_block_size,exp_mse",
     [
-        (QuIPModifier, torch.bfloat16, torch.bfloat16, 5e-3),  # 0.0019
-        (QuIPModifier, torch.bfloat16, torch.float32, 5e-3),  # 0.0022
-        (QuIPModifier, torch.float32, torch.float32, 5e-10),  # 1.0e-10
-        (QuIPModifier, torch.float32, torch.float64, 5e-11),  # 2.7e-11
-        (SpinQuantModifier, torch.bfloat16, torch.bfloat16, 5e-3),  # 0.0030
-        (SpinQuantModifier, torch.bfloat16, torch.float32, 5e-3),  # 0.0029
-        (SpinQuantModifier, torch.float32, torch.float32, 5e-4),  # 4e-4
-        (SpinQuantModifier, torch.float32, torch.float64, 5e-4),  # 4e-4
+        (QuIPModifier, torch.bfloat16, torch.bfloat16, None, 5e-3),  # 0.0019
+        (QuIPModifier, torch.bfloat16, torch.float32, 16, 5e-3),  # 0.0022
+        (QuIPModifier, torch.float32, torch.float32, 32, 5e-10),  # 1.0e-10
+        (QuIPModifier, torch.float32, torch.float64, 64, 5e-11),  # 2.7e-11
+        (SpinQuantModifier, torch.bfloat16, torch.bfloat16, None, 5e-3),  # 0.0030
+        (SpinQuantModifier, torch.bfloat16, torch.float32, 16, 5e-3),  # 0.0029
+        (SpinQuantModifier, torch.float32, torch.float32, 32, 5e-4),  # 4e-4
+        (SpinQuantModifier, torch.float32, torch.float64, 64, 5e-4),  # 4e-4
     ],
 )
-def test_apply_correctness(modifier, model_dtype, precision, exp_mse):
+def test_apply_correctness(
+    modifier, model_dtype, precision, transform_block_size, exp_mse
+):
     model = AutoModelForCausalLM.from_pretrained(
         "meta-llama/Llama-3.2-1B-Instruct", device_map="cuda", torch_dtype=model_dtype
     )
     untie_word_embeddings(model)
 
     state = State(model=model)
-    modifier = modifier(transform_type="random-hadamard", precision=precision)
+    modifier = modifier(
+        transform_type="random-hadamard",
+        precision=precision,
+        transform_block_size=transform_block_size,
+    )
 
     input = {k: v.to("cuda") for k, v in model.dummy_inputs.items()}
     with torch.no_grad():

tests/llmcompressor/modifiers/transform/test_serialization.py

@@ -4,7 +4,10 @@
 
 
 @pytest.mark.parametrize("modifier", [SpinQuantModifier, QuIPModifier])
-def test_reload(modifier):
-    instance = modifier(transform_type="hadamard")
+@pytest.mark.parametrize("transform_block_size", [16, 32])
+def test_reload(modifier, transform_block_size):
+    instance = modifier(
+        transform_type="hadamard", transform_block_size=transform_block_size
+    )
     dump = instance.model_dump()
     assert modifier.model_validate(dump) == instance