[pat] Add opt-in global group-count score correction

test/prototype/pat/test_pat_global_min_sparsity.py

@@ -95,6 +95,160 @@ def test_rejects_non_2d_grouped_view(self):
             apply_global_prox([param], prox, Grouper, {}, min_sparsity=0.5)
 
 
+class TestGroupCountCorrection(common_utils.TestCase):
+    @staticmethod
+    def _params():
+        return [torch.full((2, 4), 1.5), torch.ones(8, 4)]
+
+    def test_changes_selection_not_budget(self):
+        for score_type in ("rms", "l2", "param_cost"):
+            with self.subTest(score_type=score_type):
+                plain, corrected = self._params(), self._params()
+                prox = GlobalMinSparsityConstraint(0.0, 0.1, score_type)
+                a = apply_global_prox(plain, prox, Dim0Grouper, {}, 0.1)
+                b = apply_global_prox(
+                    corrected, prox, Dim0Grouper, {}, 0.1, score_group_count_ref=2
+                )
+                self.assertEqual([_count_zero_groups(p) for p in plain], [0, 1])
+                self.assertEqual([_count_zero_groups(p) for p in corrected], [1, 0])
+                self.assertEqual(a.zero_elts, b.zero_elts)
+                self.assertEqual(b.zero_elts, 4)
+                self.assertEqual([p.zero_elts for p in b.parameters], [4, 0])
+                self.assertEqual(b.numel, 40)
+
+    def test_reference_magnitude_preserves_selection(self):
+        results = []
+        for ref in (2, 16, 128):
+            params = self._params()
+            apply_global_prox(
+                params,
+                GlobalMinSparsityConstraint(0.0, 0.1),
+                Dim0Grouper,
+                {},
+                0.1,
+                score_group_count_ref=ref,
+            )
+            results.append([p.eq(0) for p in params])
+        self.assertEqual(results[0], results[1])
+        self.assertEqual(results[0], results[2])
+
+    def test_equal_group_counts_preserve_default(self):
+        plain = [torch.arange(1.0, 17.0).reshape(4, 4), torch.full((4, 4), 10.0)]
+        corrected = [p.clone() for p in plain]
+        prox = GlobalMinSparsityConstraint(0.0, 0.5)
+        apply_global_prox(plain, prox, Dim0Grouper, {}, 0.5)
+        apply_global_prox(
+            corrected, prox, Dim0Grouper, {}, 0.5, score_group_count_ref=16
+        )
+        self.assertEqual(plain, corrected)
+
+    def test_low_precision_scaling_preserves_finite_scores_and_selection(self):
+        for dtype in (torch.float16, torch.bfloat16):
+            for ref in (1, 16):
+                with self.subTest(dtype=dtype, reference=ref):
+                    params = [
+                        torch.tensor(values, dtype=dtype).reshape(-1, 1)
+                        for values in ([40000, 45000, 50000, 55000], [60000, 62000])
+                    ]
+                    reference = [p.float() for p in params]
+                    prox = GlobalMinSparsityConstraint(0.0, 1 / 3, "l2")
+                    apply_global_prox(
+                        reference,
+                        prox,
+                        Dim0Grouper,
+                        {},
+                        1 / 3,
+                        score_group_count_ref=ref,
+                    )
+                    with mock.patch("torch.topk", wraps=torch.topk) as topk:
+                        result = apply_global_prox(
+                            params,
+                            prox,
+                            Dim0Grouper,
+                            {},
+                            1 / 3,
+                            score_group_count_ref=ref,
+                        )
+                    scores = topk.call_args.args[0]
+                    self.assertTrue(scores.isfinite().all())
+                    self.assertEqual(scores.dtype, torch.float32)
+                    self.assertEqual(
+                        [p.eq(0) for p in params], [p.eq(0) for p in reference]
+                    )
+                    self.assertEqual(
+                        params[0].eq(0).flatten().tolist(), [True, False, False, False]
+                    )
+                    self.assertEqual(params[1].eq(0).flatten().tolist(), [True, False])
+                    self.assertEqual(result.zero_elts, 2)
+
+    def test_invalid_reference_before_mutation(self):
+        for ref in (0, -1, 1.5, True):
+            with self.subTest(reference=ref):
+                params = self._params()
+                originals = [p.clone() for p in params]
+                with self.assertRaisesRegex(ValueError, "positive integer"):
+                    apply_global_prox(
+                        params,
+                        GlobalMinSparsityConstraint(0.0, 0.1),
+                        Dim0Grouper,
+                        {},
+                        0.1,
+                        score_group_count_ref=ref,
+                    )
+                self.assertEqual(params, originals)
+
+    def test_optimizer_forwards_reference(self):
+        params = [torch.nn.Parameter(p) for p in self._params()]
+        group = {**_global_group(params, 0.1), "score_group_count_ref": 2}
+        optimizer = PruneOptimizer(torch.optim.SGD([group], lr=0.0))
+        for p in params:
+            p.grad = torch.zeros_like(p)
+        optimizer.step()
+        self.assertEqual([_count_zero_groups(p) for p in params], [1, 0])
+        self.assertEqual(optimizer.relative_sparsity, 0.1)
+
+    def test_optimizer_rejects_invalid_reference_before_warmup(self):
+        for ref in (0, -1, 1.5, True):
+            with self.subTest(reference=ref):
+                params = [torch.nn.Parameter(p) for p in self._params()]
+                group = {**_global_group(params, 0.1), "score_group_count_ref": ref}
+                with self.assertRaisesRegex(ValueError, "positive integer"):
+                    PruneOptimizer(torch.optim.SGD([group], lr=0.1), warmup_steps=10)
+
+    def test_optimizer_rejects_reference_on_non_global_groups(self):
+        for prox_type in ("MinSparsityConstraint", "ProxGroupLasso", None):
+            with self.subTest(prox_type=prox_type):
+                params = [torch.nn.Parameter(p) for p in self._params()]
+                group = {
+                    **_global_group(params, 0.1),
+                    "prox_type": prox_type,
+                    "score_group_count_ref": 2,
+                }
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "score_group_count_ref requires GlobalMinSparsityConstraint",
+                ):
+                    PruneOptimizer(torch.optim.SGD([group], lr=0.1), warmup_steps=10)
+
+    def test_optimizer_disabled_reference_preserves_default(self):
+        for prox_type in ("GlobalMinSparsityConstraint", "MinSparsityConstraint"):
+            with self.subTest(prox_type=prox_type):
+                results = []
+                for kwargs in ({}, {"score_group_count_ref": None}):
+                    params = [torch.nn.Parameter(p) for p in self._params()]
+                    group = {
+                        **_global_group(params, 0.1),
+                        "prox_type": prox_type,
+                        **kwargs,
+                    }
+                    optimizer = PruneOptimizer(torch.optim.SGD([group], lr=0.1))
+                    for p in params:
+                        p.grad = torch.ones_like(p)
+                    optimizer.step()
+                    results.append([p.detach().clone() for p in params])
+                self.assertEqual(results[0], results[1])
+
+
 class TestGlobalMinSparsityOptimizer(common_utils.TestCase):
     def _run(self, params, min_sparsity, score_type="rms", steps=3, lr=0.1):
         base = torch.optim.SGD([_global_group(params, min_sparsity, score_type)], lr=lr)

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 |
@@ -94,6 +96,8 @@ A pruning entry pairs a **grouper** with a **proximal map**. The grouper reshape
 
 `GlobalMinSparsityConstraint` computes one budget as `ceil(min_sparsity * total_groups)` for each optimizer parameter group, not one budget per tensor and not a parameter-count budget. It jointly ranks the groups exposed by the configured grouper, and all parameters in that optimizer group must produce scores on the same device. Use `score_type: rms` by default when tensors have different group sizes because it normalizes L2 magnitude by `sqrt(group_size)`. Raw `l2` tends to favor retaining larger groups because their norms grow with group size, while `param_cost` divides by the full group size and more strongly favors removing groups that save more parameters. Padded `KElementGrouper` views are rejected because padding would distort both scoring and accounting; choose a `k` that divides every grouped dimension.
 
+For recipes whose initialization scales magnitudes by the inverse square root of the number of groups, set the optional `score_group_count_ref` positive integer in a global pruning group. This multiplies each tensor's scores by `sqrt(n_groups / score_group_count_ref)` before joint selection. It is separate from RMS normalization by group size, is not automatically enabled by `head_dim`, and is disabled by default. Different positive reference magnitudes produce the same mathematical ranking because the reference contributes one common scale factor across candidates.
+
 For DTensor parameters, global selection requires full materialization of every grouped tensor on every rank before the shared top-k decision, followed by scattering the selected masks back to the original placements. The current internal `CACHE_FULL_TENSORS` policy gathers each DTensor once and retains all dense copies until selection and write-back finish, so peak dense memory is the sum of the tensors in the optimizer parameter group. A future lower-memory policy could release each copy after scoring, but would need to gather each DTensor again to apply the selected mask. Every rank gathers the same tensors and performs the same deterministic selection, so masks are expected to agree, but current PAT CI covers the DTensor API only at world size one rather than exercising true multi-rank execution.
 
 SVD decompositions can dominate optimizer cost on wide tensors. Set `prox_freq: N` on a parameter group to run its grouper and proximal map every `N` optimizer steps. SVD groups using the hard `MinRankConstraint` reapply the proximal map during healing by default so the base optimizer cannot refill removed singular values. Soft SVD maps such as `ProxNuclearNorm` retain their existing behavior unless `prox_through_heal: true` is set explicitly, while `MinRankConstraint` can opt out with `prox_through_heal: false`. Setting `prox_through_heal: true` on a non-SVD grouper is invalid and rejected during optimizer construction.