[pat] Support fixed head_dim in attention groupers

test/prototype/pat/test_pat_grouper.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD 3-Clause license found in the
 # LICENSE file in the root directory of this source tree.
 
+import math
 import random
 import unittest
 
@@ -21,6 +22,7 @@
 from torchao.prototype.pat.optim import (
     ProxGroupLasso,
     ProxNuclearNorm,
+    PruneOptimizer,
 )
 
 
@@ -83,6 +85,106 @@ def test_head_grouper(self, dim=16, head_pack_dim=0, head_dim_ratio=8):
         self._test_post_prune(p, p_orig, head_pack_dim, view_shape, reduce_dim, gamma)
 
 
+class TestFixedHeadWidth(common_utils.TestCase):
+    def test_equivalence_and_writeback(self):
+        for cls, axis in ((AttentionHeadGrouperDim0, 0), (AttentionHeadGrouperDim1, 1)):
+            for heads in (2, 4):
+                with self.subTest(grouper=cls.__name__, heads=heads):
+                    shape = (heads * 3, 5) if axis == 0 else (5, heads * 3)
+                    original = torch.arange(1.0, 1 + math.prod(shape)).reshape(shape)
+                    by_count, by_width = original.clone(), original.clone()
+                    with (
+                        cls(by_count, num_heads=heads) as a,
+                        cls(by_width, head_dim=3) as b,
+                    ):
+                        self.assertEqual(a.p, b.p)
+                        self.assertEqual(b.num_heads, heads)
+                        if axis == 0:
+                            a.p[0].zero_()
+                            b.p[0].zero_()
+                        else:
+                            a.p[:, 0].zero_()
+                            b.p[:, 0].zero_()
+                    expected = original.clone()
+                    if axis == 0:
+                        expected[:3].zero_()
+                    else:
+                        expected[:, :3].zero_()
+                    self.assertEqual(by_count, expected)
+                    self.assertEqual(by_width, expected)
+
+    def test_invalid_arguments(self):
+        for cls in (AttentionHeadGrouperDim0, AttentionHeadGrouperDim1):
+            for kwargs, message in (
+                ({}, "exactly one"),
+                ({"num_heads": 2, "head_dim": 3}, "exactly one"),
+                ({"head_dim": 0}, "positive integer"),
+                ({"num_heads": -1}, "positive integer"),
+                ({"head_dim": 1.5}, "positive integer"),
+                ({"num_heads": True}, "positive integer"),
+                ({"head_dim": 4}, "not divisible"),
+                ({"num_heads": 4}, "not divisible"),
+            ):
+                with self.subTest(grouper=cls.__name__, kwargs=kwargs):
+                    with self.assertRaisesRegex(ValueError, message):
+                        cls(torch.ones(6, 6), **kwargs)
+
+    def test_optimizer_config(self):
+        for axis in (0, 1):
+            params = [nn.Parameter(torch.arange(1.0, 25.0).reshape(6, 4))]
+            if axis == 1:
+                params = [nn.Parameter(params[0].detach().t().contiguous())]
+            group = {
+                "params": params,
+                "group_type": f"AttentionHeadGrouperDim{axis}",
+                "prox_type": "MinSparsityConstraint",
+                "min_sparsity": 0.5,
+                "head_dim": 3,
+            }
+            opt = PruneOptimizer(torch.optim.SGD([group], lr=0.0))
+            params[0].grad = torch.zeros_like(params[0])
+            opt.step()
+            # The lower-norm head contains half the elements and is pruned.
+            self.assertEqual(int(params[0].eq(0).sum()), 12)
+
+    def test_optimizer_rejects_invalid_config_before_warmup(self):
+        for axis in (0, 1):
+            for kwargs in (
+                {},
+                {"num_heads": None, "head_dim": None},
+                {"num_heads": 2, "head_dim": 3},
+            ):
+                with self.subTest(axis=axis, kwargs=kwargs):
+                    group = {
+                        "params": [nn.Parameter(torch.ones(6, 6))],
+                        "group_type": f"AttentionHeadGrouperDim{axis}",
+                        "prox_type": "MinSparsityConstraint",
+                        "min_sparsity": 0.5,
+                        **kwargs,
+                    }
+                    with self.assertRaisesRegex(ValueError, "exactly one"):
+                        PruneOptimizer(
+                            torch.optim.SGD([group], lr=0.0), warmup_steps=10
+                        )
+
+    def test_optimizer_rejects_indivisible_head_dim(self):
+        for axis in (0, 1):
+            with self.subTest(axis=axis):
+                shape = (6, 4) if axis == 0 else (4, 6)
+                param = nn.Parameter(torch.ones(shape))
+                group = {
+                    "params": [param],
+                    "group_type": f"AttentionHeadGrouperDim{axis}",
+                    "prox_type": "MinSparsityConstraint",
+                    "min_sparsity": 0.5,
+                    "head_dim": 5,
+                }
+                opt = PruneOptimizer(torch.optim.SGD([group], lr=0.0))
+                param.grad = torch.zeros_like(param)
+                with self.assertRaisesRegex(ValueError, "not divisible"):
+                    opt.step()
+
+
 class TestDimGrouper(common_utils.TestCase):
     def test_noncontiguous_flatten_writes_back(self):
         param = torch.arange(24.0).reshape(2, 3, 4).transpose(1, 2)

torchao/prototype/pat/README.md

@@ -60,11 +60,13 @@ A pruning entry pairs a **grouper** with a **proximal map**. The grouper reshape
 | `LayerGrouper` | Whole tensor as one group for layer-level pruning |
 | `KElementGrouper(k)` | `(numel / k, k)` blocks of `k` consecutive elements |
 | `ConvFilterGrouper` | One group per `(c_out, c_in)` filter slice of a Conv2d kernel |
-| `AttentionHeadGrouperDim0(num_heads)` | One group per attention head along dimension 0 |
-| `AttentionHeadGrouperDim1(num_heads)` | One group per attention head along dimension 1 |
+| `AttentionHeadGrouperDim0(num_heads=..., head_dim=...)` | One group per attention head along dimension 0; specify exactly one argument |
+| `AttentionHeadGrouperDim1(num_heads=..., head_dim=...)` | One group per attention head along dimension 1; specify exactly one argument |
 | `SVDGrouper` | Decompose `W = U diag(s) Vh` and expose its singular values |
 | `PackedSVDGrouper(npack)` | Apply SVD independently to each of `npack` sub-tensors |
 
+Attention-head groupers accept exactly one positive integer `num_heads` or `head_dim`, and the packed dimension must be divisible by it. Existing positional `num_heads` calls remain supported. For heterogeneous layer widths, use a fixed `head_dim` in the pruning config so each tensor derives its own head count; for example, `group_type: AttentionHeadGrouperDim0` with `head_dim: 64`.
+
 ### Proximal maps (`torchao.prototype.pat.optim`)
 
 | Proximal map | Behavior |
@@ -85,7 +87,7 @@ A pruning entry pairs a **grouper** with a **proximal map**. The grouper reshape
 | **Row sparsity** | `Dim0Grouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Use the hard constraint for an exact target or group Lasso for a smooth regularization knob |
 | **Column sparsity** | `Dim1Grouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Drops input channels of a Linear layer |
 | **Conv filter sparsity** | `ConvFilterGrouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Drops complete `(c_out, c_in)` filter slices |
-| **Attention head sparsity** | `AttentionHeadGrouperDim0` and/or `AttentionHeadGrouperDim1` | `MinSparsityConstraint` | Configure matching `num_heads` for the selected projections |
+| **Attention head sparsity** | `AttentionHeadGrouperDim0` and/or `AttentionHeadGrouperDim1` | `MinSparsityConstraint` | Configure matching `num_heads` or `head_dim` for the selected projections |
 | **Global structured sparsity** | Any supported non-SVD grouper | `GlobalMinSparsityConstraint` | Shares one group-count budget across every tensor in the optimizer parameter group; `score_type: rms` is the recommended default |
 | **Low-rank approximation, smooth** | `SVDGrouper` or `PackedSVDGrouper` | `ProxNuclearNorm` | Uses regularization-controlled rank decay |
 | **Low-rank approximation, exact target** | `SVDGrouper` or `PackedSVDGrouper` | `MinRankConstraint(min_sparsity)` | Zeros `ceil(min_sparsity * k)` and retains `k - ceil(min_sparsity * k)` singular values per matrix |

torchao/prototype/pat/group/attention.py

@@ -4,40 +4,63 @@
 # This source code is licensed under the BSD 3-Clause license found in the
 # LICENSE file in the root directory of this source tree.
 
-
 from torch import Tensor
 
 from .dim import Dim0Grouper, Dim1Grouper
 
 
+def _resolve_heads(
+    packed_dim: int, num_heads: int | None, head_dim: int | None
+) -> tuple[int, int]:
+    if (num_heads is None) == (head_dim is None):
+        raise ValueError("specify exactly one of num_heads or head_dim")
+    if head_dim is not None:
+        name, value = "head_dim", head_dim
+    else:
+        name, value = "num_heads", num_heads
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer, got {value!r}")
+    if packed_dim % value:
+        raise ValueError(
+            f"packed dimension {packed_dim} not divisible by {name} {value}"
+        )
+    if head_dim is not None:
+        return packed_dim // value, value
+    return value, packed_dim // value
+
+
 class AttentionHeadGrouperDim0(Dim0Grouper):
-    """Grouper for attention heads.
+    """Grouper for attention heads packed along dimension 0.
 
-    Args:
-        p (Tensor): Input tensor with packed attention heads.
-        num_heads (int): Number of attention heads.
+    Specify exactly one of ``num_heads`` or ``head_dim``. A fixed head width
+    allows tensors with different head counts to share a pruning config.
     """
 
-    def __init__(self, p: Tensor, num_heads: int):
+    def __init__(
+        self, p: Tensor, num_heads: int | None = None, head_dim: int | None = None
+    ):
         super().__init__(p)
-
-        self.num_heads = num_heads
-        self.head_dim = p.size(0) // num_heads
+        self.num_heads, self.head_dim = _resolve_heads(p.size(0), num_heads, head_dim)
 
     def __enter__(self):
         self.p = self.p.view(self.num_heads, -1)
         return self
 
 
 class AttentionHeadGrouperDim1(Dim1Grouper):
-    def __init__(self, p: Tensor, num_heads: int):
+    """Grouper for attention heads packed along dimension 1.
+
+    Specify exactly one of ``num_heads`` or ``head_dim``.
+    """
+
+    def __init__(
+        self, p: Tensor, num_heads: int | None = None, head_dim: int | None = None
+    ):
         self._orig_p = p
-        self.head_dim = p.size(1) // num_heads
-        p = p.view(-1, num_heads, self.head_dim).transpose(1, 2).contiguous()
+        self.num_heads, self.head_dim = _resolve_heads(p.size(1), num_heads, head_dim)
+        p = p.view(-1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
         super().__init__(p)
 
-        self.num_heads = num_heads
-
     def __enter__(self):
         self.p = self.p.view(-1, self.num_heads)
         return self

torchao/prototype/pat/optim/pruneopt.py

@@ -61,6 +61,7 @@ def __init__(
             group.setdefault("gamma", 0.0)
             group.setdefault("reg_lambda", reg_lambda)
             self._validate_prox_through_heal(group)
+            self._get_grouper_kwargs(group)
             if group.get("min_sparsity_schedule", False):
                 assert self.healing_start_step != sys.maxsize, (
                     "min_sparsity_schedule requires a finite healing_start_step; "
@@ -186,7 +187,13 @@ def _effective_min_sparsity(self, group: dict[str, Any]) -> float:
     def _get_grouper_kwargs(group: dict[str, Any]) -> dict[str, Any]:
         grouper_kwargs = {}
         if group["group_type"].startswith("AttentionHeadGrouper"):
-            grouper_kwargs["num_heads"] = group["num_heads"]
+            num_heads, head_dim = group.get("num_heads"), group.get("head_dim")
+            if (num_heads is None) == (head_dim is None):
+                raise ValueError("specify exactly one of num_heads or head_dim")
+            if head_dim is not None:
+                grouper_kwargs["head_dim"] = head_dim
+            else:
+                grouper_kwargs["num_heads"] = num_heads
         elif group["group_type"] == "KElementGrouper":
             grouper_kwargs["k"] = group["k"]
         elif group["group_type"] == "PackedSVDGrouper":