diff --git a/test/prototype/pat/test_pat_global_min_sparsity.py b/test/prototype/pat/test_pat_global_min_sparsity.py index eb742a05f11..80794662679 100644 --- a/test/prototype/pat/test_pat_global_min_sparsity.py +++ b/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) 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..36fedafbf55 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 | @@ -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. 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/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 3e43639c2b7..5e4614f0981 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,12 @@ 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: + # 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) @@ -373,9 +387,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 +405,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 ebed7dccc85..7e9ec0da6f3 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.py @@ -53,6 +53,8 @@ def __init__( assert warmup_steps < healing_start_step, ( f"Invalid {warmup_steps=} >= {healing_start_step=}" ) + for group in self.param_groups: + self._validate_score_group_count_ref(group) self.num_steps = 0 self.warmup_steps = warmup_steps self.healing_start_step = healing_start_step @@ -61,6 +63,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 +189,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": @@ -270,6 +279,18 @@ def should_prune(self, group: dict[str, Any], step: int) -> bool: offset = step - self.warmup_steps return offset >= 0 and offset % freq == 0 + @staticmethod + def _validate_score_group_count_ref(group: dict[str, Any]) -> None: + ref = group.get("score_group_count_ref") + if ref is None: + return + if group.get("prox_type") != "GlobalMinSparsityConstraint": + raise ValueError( + "score_group_count_ref requires GlobalMinSparsityConstraint" + ) + if not isinstance(ref, int) or isinstance(ref, bool) or ref <= 0: + raise ValueError("score_group_count_ref must be a positive integer") + @staticmethod def _validate_prox_through_heal(group: dict[str, Any]) -> bool: is_svd_grouper = str(group.get("group_type", "")).endswith("SVDGrouper") @@ -397,6 +418,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]