From 92dc9d441120f01d379018c859c3bc4e25ea1340 Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 10:49:29 -0700 Subject: [PATCH 1/4] Update (base 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 c83e862c9a6b2eb90f2712870187388c9f17a7ed Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 10:49:29 -0700 Subject: [PATCH 2/4] [INITIAL] Update [ghstack-poisoned] --- .../pat/test_pat_global_min_sparsity.py | 74 +++++++++++++++++++ torchao/prototype/pat/README.md | 2 + torchao/prototype/pat/optim/prox_executor.py | 27 ++++++- torchao/prototype/pat/optim/pruneopt.py | 1 + 4 files changed, 100 insertions(+), 4 deletions(-) diff --git a/test/prototype/pat/test_pat_global_min_sparsity.py b/test/prototype/pat/test_pat_global_min_sparsity.py index eb742a05f11..881b252152a 100644 --- a/test/prototype/pat/test_pat_global_min_sparsity.py +++ b/test/prototype/pat/test_pat_global_min_sparsity.py @@ -95,6 +95,80 @@ 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_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) + + 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) diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 6ed1ceca230..f1532d18b8c 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -96,6 +96,8 @@ Attention-head groupers accept exactly one positive integer `num_heads` or `head `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. diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 3e43639c2b7..43fe8bc1e38 100644 --- a/torchao/prototype/pat/optim/prox_executor.py +++ b/torchao/prototype/pat/optim/prox_executor.py @@ -283,6 +283,7 @@ def apply_global_prox( grouper_kwargs: dict[str, Any], min_sparsity: float, *, + score_group_count_ref: int | None = None, materialization_policy: GlobalDTensorMaterializationPolicy = ( GlobalDTensorMaterializationPolicy.CACHE_FULL_TENSORS ), @@ -304,6 +305,13 @@ def apply_global_prox( f"Unsupported materialization policy: {materialization_policy}" ) + if score_group_count_ref is not None and ( + not isinstance(score_group_count_ref, int) + or isinstance(score_group_count_ref, bool) + or score_group_count_ref <= 0 + ): + raise ValueError("score_group_count_ref must be a positive integer") + dtensor_params = [p for p in params if _is_dtensor(p)] if dtensor_params: if len(dtensor_params) != len(params): @@ -340,6 +348,8 @@ def apply_global_prox( f"for parameter shape {tuple(p.shape)}." ) scores = prox_map.score(view).detach() + if score_group_count_ref is not None: + scores = scores * math.sqrt(scores.numel() / score_group_count_ref) entries.append((p, grouper, view, full, scores, grouper.p.numel())) score_chunks.append(scores) @@ -373,9 +383,7 @@ def apply_global_prox( ) zero_mask[drop_idx] = True - parameter_results = [] - group_zeros = 0 - group_params = 0 + zero_counts = [] offset = 0 for p, grouper, view, full, scores, numel in entries: n_local = scores.numel() @@ -393,7 +401,18 @@ def apply_global_prox( placements=grouper.p.placements, ) ) - zero_elts = int(zeros.item()) if torch.is_tensor(zeros) else int(zeros) + zero_counts.append( + zeros + if torch.is_tensor(zeros) + else torch.tensor(zeros, device=view.device, dtype=torch.long) + ) + + # One device-to-host synchronization for the whole parameter group. + counts = torch.stack(zero_counts).to(torch.long).tolist() + parameter_results = [] + group_zeros = 0 + group_params = 0 + for (p, grouper, view, full, scores, numel), zero_elts in zip(entries, counts): parameter_results.append( ParameterSparsity(parameter=p, zero_elts=zero_elts, numel=numel) ) diff --git a/torchao/prototype/pat/optim/pruneopt.py b/torchao/prototype/pat/optim/pruneopt.py index 0ffbab18918..ae54e9efe47 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.py @@ -403,6 +403,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: grouper_cls, grouper_kwargs, self._effective_min_sparsity(group), + score_group_count_ref=group.get("score_group_count_ref"), ) for param_result in global_result.parameters: state = self.state[param_result.parameter] From 96f5fb31fa4156d4f02263be69a90039b4f7af77 Mon Sep 17 00:00:00 2001 From: lisjin Date: Wed, 14 Oct 2026 16:31:46 -0700 Subject: [PATCH 3/4] Update (base 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; " From 4bea2a51336d7d1ab4636850228bc00b3b93e3be Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 15 Oct 2026 10:36:22 -0700 Subject: [PATCH 4/4] [UPDATE] Update [ghstack-poisoned] --- .../pat/test_pat_global_min_sparsity.py | 39 +++++++++++++++++++ torchao/prototype/pat/optim/prox_executor.py | 4 ++ 2 files changed, 43 insertions(+) diff --git a/test/prototype/pat/test_pat_global_min_sparsity.py b/test/prototype/pat/test_pat_global_min_sparsity.py index b08e3c969e6..80794662679 100644 --- a/test/prototype/pat/test_pat_global_min_sparsity.py +++ b/test/prototype/pat/test_pat_global_min_sparsity.py @@ -142,6 +142,45 @@ def test_equal_group_counts_preserve_default(self): ) 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): diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 43fe8bc1e38..5e4614f0981 100644 --- a/torchao/prototype/pat/optim/prox_executor.py +++ b/torchao/prototype/pat/optim/prox_executor.py @@ -349,6 +349,10 @@ def apply_global_prox( ) scores = prox_map.score(view).detach() if score_group_count_ref is not None: + # Scale low-precision scores in FP32 to avoid overflow and + # additional rounding before global ranking. + if scores.dtype in (torch.float16, torch.bfloat16): + scores = scores.float() scores = scores * math.sqrt(scores.numel() / score_group_count_ref) entries.append((p, grouper, view, full, scores, grouper.p.numel())) score_chunks.append(scores)