[pat] Low-rank from internal repo

[lisjin]

Stack from ghstack (oldest at bottom):

• [pat] Add GlobalMinSparsityConstraint for cross-layer budget allocation #4587
• -> [pat] Low-rank from internal repo #4586

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4586

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Unrelated MacOS MPS test jobs time out (30 min) due to stale cpp_extension lock file poisoning self-hosted runners

✅ No Failures

As of commit 783de37 with merge base b7f01f1 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lisjin]

Duplicate of #4579 created by ghstack. Refer to that PR for past claude reviews

[andrewor14]

@claude review

[claude]

Claude finished @andrewor14's task in 5m 22s —— View job

---

Claude Code Review

• Read PR diff
• Analyze changes
• Post review feedback

This PR ports the low-rank (MinRankConstraint) proximal map from the internal repo, refactors PruneOptimizer.step() to share prox/grouper construction between pruning and healing, adds prox_freq gating and a prox_through_heal policy, removes the unused QKGrouper, and adds a solid test suite + README rewrite. Overall this is clean, well-tested work. A few observations below.

---

Correctness

1. Packed SVD metrics fixes look correct — good catches. The changes to _apply_prox and _run_prox_on_param fix two latent packed-SVD bugs:

• pruneopt.py:376 — dim = -1 if grouper.p.dim() > 1 else None now counts nonzero singular values per pack (shape (npack,)), matching sv_count. The old dim = 0 if sv_count.dim() > 1 else None collapsed the wrong axis for PackedSVDGrouper (where sv_count.dim() == 1).
• _run_prox_on_param uses grouper.U.size(-2) / grouper.Vh.size(-1) and unfactored_size = p.numel() instead of U.size(0) * Vh.size(1), which was wrong for the packed 3-D U/Vh.

The new test_packed_svd_grouper and test_full_step_uses_total_packed_storage correctly pin this down. 👍

2. prox_freq > 1 during healing can let low-rank structure drift. For non-through-heal groups, healing re-zeros literal zeros every step via healing_masks. But through-heal groups (the MinRankConstraint default) only re-apply the prox when should_prune(group, self.num_steps) returns True (pruneopt.py:679). With prox_freq > 1, the base optimizer refills singular values on the skipped steps and they are not re-zeroed until the next multiple — so the effective rank can exceed the target between reapplications. The README claims "the base optimizer cannot refill removed singular values," which only strictly holds for the default prox_freq = 1. This is an edge case (needs both prox_freq > 1 and a healing phase), but worth either documenting or forcing prox_freq = 1 during healing for through-heal groups. Fix this →

3. state["latent"] copy in the skipped-group branch assumes latent exists. At pruneopt.py:588, when should_prune is False the code does self.state[p]["latent"].copy_(p) for every requires_grad param. This is safe today because the first PAT step (num_steps == warmup_steps) always prunes (offset = 0, 0 % freq == 0) and save_latent_params() populates latent there. Just noting the invariant is load-bearing — if warmup/scheduling logic changes so the first post-warmup step could be skipped, this would KeyError. A state.get("latent") guard (mirroring _init_latent_state's "latent" in state check) would be more defensive.

Code quality

4. MinRankConstraint and MinSparsityConstraint are near-duplicates. They differ only in _get_norm and the 1-D-vs-2-D handling in apply_. Not blocking, but a shared base (both already share _TopKZeroMixin) would reduce drift risk between the two constraint validations/docstrings.

5. Minor: relative_factored_frac init. pruneopt.py:90 initializes self.relative_factored_frac = 0 (int) while the computed value is a float; self.relative_sparsity = 0 likewise. Harmless, but 0.0 would be tidier and type-consistent.

Tests

The test file is comprehensive and reads well — direct constraint tests, grouper integration, prox_freq gating (nice use of patch.object to count __enter__ calls), the mixed-frequency stale-metrics guarantee, through-heal policy defaults/overrides/rejection, and the packed factorization metric. One suggestion: TestProxThroughHeal.test_singular_values_stay_zero_during_healing asserts effective_rank <= k - n_killed using default prox_freq = 1; adding a case with prox_freq > 1 would surface finding #2 if you decide it matters.

Cleanups (no action needed)

• QKGrouper removal is clean — no dangling references remain in torchao/, test/, or docs/ (grepped), and the grouper_kwargs branch in _get_grouper_kwargs was removed accordingly.
• utils.py:154 import path fix (pat.group.SVDGrouper → torchao.prototype.pat.group.SVDGrouper) fixes a real bug in insert_svd_modules_, and the README's from pat.optim import ... was likewise corrected.

Nothing here is a blocker; #2 is the only substantive item and it's an edge-case/doc concern. Nice PR.
• gh/lisjin/1/head

[andrewor14]

@claude please pick it up from the reviews in #4579, it's the same changes I think

[claude]

Claude finished @andrewor14's task in 5m 43s —— View job

---

Claude Code Review

• Read prior reviews from [pat] Low-rank pruning support #4579
• Read PR diff
• Analyze changes
• Post review feedback

Note: I wasn't able to fetch the review threads from #4579 in this environment (the GitHub API / web-fetch calls needed approvals that aren't granted here). Rather than rely on those, I re-reviewed these changes directly — the author confirmed they're identical. If a specific point from #4579 is still unaddressed, point me at it and I'll dig in.

---

Overall this is a clean, well-structured change. The step() body is refactored into reusable helpers (_build_prox_artifacts, _run_prox_on_param), and the new MinRankConstraint + prox_freq + prox_through_heal features are all backed by focused tests. No blocking issues found.

Correctness — verified fixes ✅

A few pre-existing bugs in the SVD metrics path are quietly fixed here and are worth calling out as good catches:

• Packed unfactored_size (pruneopt.py:446): changed from grouper.U.size(0) * grouper.Vh.size(1) to p.numel(). For PackedSVDGrouper, U is (npack, m, k), so the old U.size(0) was npack and Vh.size(1) was k — the product was meaningless for packed tensors. p.numel() is correct, and TestPackedFactorizationMetrics confirms it.
• Matrix dims (pruneopt.py:444-445): U.size(-2) / Vh.size(-1) correctly pick per-matrix rows/cols for both packed and unpacked, vs. the old U.size(0)/Vh.size(1).
• sv_count reduction dim (pruneopt.py:376): dim = -1 if grouper.p.dim() > 1 else None. The old dim = 0 if sv_count.dim() > 1 ... was always None (sv_count is 1-D), which collapsed per-pack counts into a single scalar. test_packed_svd_grouper covers the per-pack count now.

Feature review

• should_prune / prox_freq — gamma is advanced every step (_set_gamma runs before the gate), and skipped steps still sync latent, so the base optimizer stays consistent. The all_groups_ran guard freezing global metrics on partial steps is a reasonable choice and is documented + tested (test_mixed_frequencies_keep_complete_metrics).
• prox_through_heal — the default policy (hard MinRankConstraint reapplies during healing; soft maps don't unless opted in) is sensible, since dense low-rank weights have no literal zeros for the gradient mask to freeze. The _validate_prox_through_heal guard rejecting the flag on non-SVD groupers both at construction and on later mutation is a nice touch.
• MinRankConstraint — zeros the smallest ceil(min_sparsity * k) singular values; the per-matrix loop for the packed case correctly sums zero counts. Semantics are order-independent (topk-smallest), so it's robust to SVD ordering.

Minor observations (non-blocking)

1. _apply_prox_to_through_heal_groups doesn't refresh sparsity/factored metrics. During healing, reapplying the prox updates sv_count but leaves relative_sparsity / relative_factored_frac at their last pre-healing values. This is almost certainly intentional (rank is held constant during healing), but a one-line comment saying so would save the next reader a double-take.

2. MinRankConstraint packed loop is Python-level per matrix (min_sparsity.py:120-124). Fine for typical npack (2–3), and it correctly bypasses vmap via whole_tensor = True. Just flagging that it's sequential if npack ever grows large.

3. QKGrouper removal is clean — I grep'd the whole tree and confirmed no remaining references to QKGrouper, qk_pack_dim, or qk_reg_index, and the corresponding grouper-kwargs branch in pruneopt.py was removed too. PackedGrouperMixin is still used by low_rank.py, so leaving it in packed.py is correct.

Tests

test_pat_min_rank.py has good coverage: direct prox behavior, SVD/packed-SVD integration, the prox_freq gate (via an __enter__ counter), the prox_through_heal policy matrix, healing invariance, and packed factorization metrics. I was unable to execute the suite here (test-run commands require approval in this environment), so I reviewed it statically — the assertions line up with the implementation.
· branch gh/lisjin/1/head