diff --git a/test/prototype/pat/test_pat_grouper.py b/test/prototype/pat/test_pat_grouper.py index 8a11eb8c17c..e40712f80e5 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,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) diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 93192522072..d0332d7bec9 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 | @@ -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 | diff --git a/torchao/prototype/pat/group/attention.py b/torchao/prototype/pat/group/attention.py index 5faa7165f30..caf3766323a 100644 --- a/torchao/prototype/pat/group/attention.py +++ b/torchao/prototype/pat/group/attention.py @@ -4,25 +4,43 @@ # 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) @@ -30,14 +48,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..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; " @@ -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":