[Transform] Revert deprecation of TransformScheme.head_dim for compatibility with vllm (#472)

* allow for use of head_dim for vllm

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

* cleanup

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

* revert TransformScheme.block_size

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

* test fixes

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

* docstring update

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

---------

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.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         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.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         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, model_validator
+from pydantic import BaseModel, ConfigDict, Field
 
 
 __all__ = ["TransformScheme"]
@@ -36,8 +36,11 @@ 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 head_dim: If set, the transform matrix will be block diagonal with each
+        block being a square matrix of this size. The name head_dim was chosen because
+        some rotations need to be block-diagonal with block size equal to the head_dim,
+        but research has shown value in applying some rotations with smaller block size,
+        irrespective of head_dim.
     :param precision: Precision at which this transform should be applied during online
         rotations. Fused (offline) rotations are always performed in float64
     """
@@ -46,21 +49,7 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
-    block_size: Optional[int] = Field(default=None)
-    head_dim: Optional[int] = Field(
-        default=None, deprecated="head_dim is deprecated, use block_size instead"
-    )
+    head_dim: Optional[int] = Field(default=None)
     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, block_size=head_dim)
+    scheme = TransformScheme(type=type, randomize=randomize, head_dim=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,
-                block_size=head_dim,
+                head_dim=head_dim,
                 apply=[
                     TransformArgs(targets="v_proj", location="weight_output"),
                     TransformArgs(

tests/test_transform/test_transform_scheme.py

@@ -72,28 +72,3 @@ 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