TransformScheme.block_size, deprecate head_dim (#466)

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

Author: brian-dellabetta

src/compressed_tensors/transform/factory/hadamard.py

@@ -52,7 +52,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.head_dim)
+        size = get_transform_size(module, args.location, self.scheme.block_size)
         exec_device = get_execution_device(module)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -51,7 +51,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.head_dim)
+        size = get_transform_size(module, args.location, self.scheme.block_size)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64
 

src/compressed_tensors/transform/transform_scheme.py

@@ -17,7 +17,7 @@
 import torch
 from compressed_tensors.transform import TransformArgs
 from compressed_tensors.utils import TorchDtype
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 __all__ = ["TransformScheme"]
@@ -36,6 +36,8 @@ class TransformScheme(BaseModel):
     :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
+    :param block_size: If set, the transform matrix will be block diagonal, with each
+        block being a square matrix of this size.
     :param precision: Precision at which this transform should be applied during online
         rotations. Fused (offline) rotations are always performed in float64
     """
@@ -44,7 +46,21 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
-    head_dim: Optional[int] = Field(default=None)
+    block_size: Optional[int] = Field(default=None)
+    head_dim: Optional[int] = Field(
+        default=None, deprecated="head_dim is deprecated, use block_size instead"
+    )
     precision: TorchDtype = Field(default=torch.float32)
 
+    @model_validator(mode="after")
+    def validate_model_after(model: "TransformScheme") -> "TransformScheme":
+        """
+        If head_dim is used instead of block_size, set block_size to head_dim
+        and remove head_dim
+        """
+        if model.block_size is None and model.head_dim is not None:
+            model.block_size = model.head_dim
+            model.head_dim = None
+        return model
+
     model_config = ConfigDict(extra="forbid")

tests/test_transform/factory/test_correctness.py

@@ -33,7 +33,7 @@
 def test_correctness_linear(type, randomize, head_dim, input_batch_size):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=False)
-    scheme = TransformScheme(type=type, randomize=randomize, head_dim=head_dim)
+    scheme = TransformScheme(type=type, randomize=randomize, block_size=head_dim)
     factory = TransformFactory.from_scheme(scheme, name="")
 
     input_tfm = factory.create_transform(
@@ -150,7 +150,7 @@ def test_correctness_attention_heads(type, randomize, head_dim, input_batch_size
             "": TransformScheme(
                 type=type,
                 randomize=randomize,
-                head_dim=head_dim,
+                block_size=head_dim,
                 apply=[
                     TransformArgs(targets="v_proj", location="weight_output"),
                     TransformArgs(

tests/test_transform/test_transform_scheme.py

@@ -72,3 +72,28 @@ def test_multiple_groups():
     assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 20
+
+
+def test_transform_scheme_block_size():
+    """
+    Ensure json with (deprecated) `head_dim` or `block_size`
+    both load up correctly and save with `block_size` field
+    """
+
+    old_scheme = TransformScheme.model_validate_json(
+        '{"type": "hadamard", "head_dim": 128}'
+    )
+    assert old_scheme.block_size == 128
+    assert old_scheme.model_dump()["block_size"] == 128
+    old_scheme = TransformScheme(type="hadamard", head_dim=64)
+    assert old_scheme.block_size == 64
+    assert old_scheme.model_dump()["block_size"] == 64
+
+    new_scheme = TransformScheme.model_validate_json(
+        '{"type": "hadamard", "block_size": 128}'
+    )
+    assert new_scheme.block_size == 128
+    assert new_scheme.model_dump()["block_size"] == 128
+    new_scheme = TransformScheme(type="hadamard", block_size=64)
+    assert new_scheme.block_size == 64
+    assert new_scheme.model_dump()["block_size"] == 64