From f1c9a53285a7011f6a2f3d5b3d213e6be5fae4a9 Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 10:49:29 -0700 Subject: [PATCH 1/2] [INITIAL] Update [ghstack-poisoned] --- test/prototype/pat/test_pat_grouper.py | 68 ++++++++++++++++++++++++ torchao/prototype/pat/README.md | 6 ++- torchao/prototype/pat/group/attention.py | 50 ++++++++++++----- torchao/prototype/pat/optim/pruneopt.py | 8 ++- 4 files changed, 115 insertions(+), 17 deletions(-) diff --git a/test/prototype/pat/test_pat_grouper.py b/test/prototype/pat/test_pat_grouper.py index 8a11eb8c17c..3ce6843fc98 100644 --- a/test/prototype/pat/test_pat_grouper.py +++ b/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,72 @@ 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() + self.assertEqual(int(params[0].eq(0).sum()), 12) + for kwargs in ({}, {"num_heads": 2, "head_dim": 3}): + invalid = {"group_type": group["group_type"], **kwargs} + with self.assertRaisesRegex(ValueError, "exactly one"): + opt._get_grouper_kwargs(invalid) + + class TestDimGrouper(common_utils.TestCase): def test_noncontiguous_flatten_writes_back(self): param = torch.arange(24.0).reshape(2, 3, 4).transpose(1, 2) diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 93192522072..6ed1ceca230 100644 --- a/torchao/prototype/pat/README.md +++ b/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 | diff --git a/torchao/prototype/pat/group/attention.py b/torchao/prototype/pat/group/attention.py index 5faa7165f30..cea1b667f76 100644 --- a/torchao/prototype/pat/group/attention.py +++ b/torchao/prototype/pat/group/attention.py @@ -4,25 +4,42 @@ # 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") + name, value = ( + ("head_dim", head_dim) if head_dim is not None else ("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) @@ -30,14 +47,19 @@ def __enter__(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 diff --git a/torchao/prototype/pat/optim/pruneopt.py b/torchao/prototype/pat/optim/pruneopt.py index ebed7dccc85..0ffbab18918 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.py @@ -186,7 +186,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": From 33be25032fe4e57e3cc4fe98f82f938a402b787f Mon Sep 17 00:00:00 2001 From: lisjin Date: Wed, 14 Oct 2026 16:17:08 -0700 Subject: [PATCH 2/2] [UPDATE] Update [ghstack-poisoned] --- test/prototype/pat/test_pat_grouper.py | 42 +++++++++++++++++++++--- torchao/prototype/pat/README.md | 2 +- torchao/prototype/pat/group/attention.py | 7 ++-- torchao/prototype/pat/optim/pruneopt.py | 1 + 4 files changed, 44 insertions(+), 8 deletions(-) diff --git a/test/prototype/pat/test_pat_grouper.py b/test/prototype/pat/test_pat_grouper.py index 3ce6843fc98..e40712f80e5 100644 --- a/test/prototype/pat/test_pat_grouper.py +++ b/test/prototype/pat/test_pat_grouper.py @@ -144,11 +144,45 @@ def test_optimizer_config(self): 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) - for kwargs in ({}, {"num_heads": 2, "head_dim": 3}): - invalid = {"group_type": group["group_type"], **kwargs} - with self.assertRaisesRegex(ValueError, "exactly one"): - opt._get_grouper_kwargs(invalid) + + 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): diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 6ed1ceca230..d0332d7bec9 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -87,7 +87,7 @@ Attention-head groupers accept exactly one positive integer `num_heads` or `head | **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 | diff --git a/torchao/prototype/pat/group/attention.py b/torchao/prototype/pat/group/attention.py index cea1b667f76..caf3766323a 100644 --- a/torchao/prototype/pat/group/attention.py +++ b/torchao/prototype/pat/group/attention.py @@ -14,9 +14,10 @@ def _resolve_heads( ) -> tuple[int, int]: if (num_heads is None) == (head_dim is None): raise ValueError("specify exactly one of num_heads or head_dim") - name, value = ( - ("head_dim", head_dim) if head_dim is not None else ("num_heads", num_heads) - ) + 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: diff --git a/torchao/prototype/pat/optim/pruneopt.py b/torchao/prototype/pat/optim/pruneopt.py index 0ffbab18918..e22d9fc29a7 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/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; "