[pat] Support fixed head_dim in attention groupers

[lisjin]

Stack from ghstack (oldest at bottom):

• -> [pat] Support fixed head_dim in attention groupers #4982

Allow heterogeneous attention widths to share a pruning configuration while preserving existing num_heads calls. Validate exclusive positive-integer arguments and packed-axis divisibility in the groupers, and reject missing or conflicting head arguments at optimizer construction before warmup. Adapted from qpat ebc76a1 (#57).

Test Plan:

source ~/.ao/bin/activate
python -m pytest test/prototype/pat/test_pat_grouper.py test/prototype/pat/test_pat_pruneopt.py -q
# 24 passed, 28 subtests passed
ruff check --select F,I torchao/prototype/pat/group/attention.py torchao/prototype/pat/optim/pruneopt.py test/prototype/pat/test_pat_grouper.py
ruff format --check torchao/prototype/pat/group/attention.py torchao/prototype/pat/optim/pruneopt.py test/prototype/pat/test_pat_grouper.py
git diff --check

The new construction-time regression failed for all six missing/conflicting argument cases before the fix.

Tags: stack_created_with_devmate

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 33be250 with merge base cff77b4 ():
💚 Looks good so far! There are no failures yet. 💚

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

[lisjin]

@claude review

[claude]

Claude finished @lisjin's task in 1m 14s —— View job

---

Review: [pat] Support fixed head_dim in attention groupers

• Read the PR diff against origin/gh/lisjin/1/base
• Check how the optimizer builds and calls the groupers (pruneopt.py, prox_executor.py)
• Post review

Overall this looks good. _resolve_heads is short and handles the edge cases (bools, floats, non-positive values, axes that don't divide evenly). Positional num_heads calls still work. The old code would fail inside .view() when the axis didn't divide evenly (Dim0), or produce a confusing error (Dim1), so the explicit divisibility error is an improvement. Calling the grouper on p.full_tensor() for DTensors means head_dim is resolved against the global shape, which is what you want. I didn't find any correctness bugs. Some suggestions:

1. Optimizer config errors only appear at the first prune step, not at construction

_get_grouper_kwargs only runs from _build_group_artifacts, which runs inside step() after should_prune passes. With warmup_steps=N, a group with no head argument, or with both num_heads and head_dim, raises only after N steps. A shape that doesn't divide evenly also surfaces only then. The PR description says config validation happens in the optimizer, but nothing is checked in __init__. A cheap fix is to run the check in the existing __init__ loop (pruneopt.py:60-68):

for group in self.regularized_param_groups():
    ...
    self._get_grouper_kwargs(group)  # fail fast on bad grouper config

You could go further and build the grouper once per param there, so divisibility errors also show up early. That costs some construction work, so it's optional.
Fix this →

2. The exactly-one check is written twice

pruneopt.py:189-195 repeats the (num_heads is None) == (head_dim is None) check that _resolve_heads already does. Once (1) is in place, keep it in the optimizer so the error comes early. Otherwise it could shrink to:

grouper_kwargs.update(
    {k: group[k] for k in ("num_heads", "head_dim") if group.get(k) is not None}
)

Then the grouper is the only place the rule is enforced. This is minor; either way works.

3. One README line still mentions only num_heads

README.md:90 (the "Attention head sparsity" recipe row) still says "Configure matching num_heads for the selected projections". It should probably say "num_heads or head_dim", or point to the new paragraph.

4. Nits

• attention.py:17-19: the name, value = (...) tuple swap is a bit dense. An if/else that sets name and value would read more easily, but it's fine as is.
• The tests use parenthesized multi-item with (...), which requires Python 3.10+. That's fine if 3.10 is already the minimum.
• test_optimizer_config doesn't run the "packed axis doesn't divide evenly" case through the optimizer, e.g. head_dim: 5 on a (6, 4) param raising ValueError from step(). If you do (1), a construction-time version of that test would cover it.
• The MinSparsityConstraint assertion of 12 zeros relies on the 2 heads of arange data having clearly different norms, which they do. That's correct but a little implicit; a one-line comment would help.

Tests

I didn't run the test suite in this environment. The new tests cover:

• grouping and write-back being identical for num_heads and head_dim on both axes
• each invalid-argument case
• the optimizer accepting head_dim

[z-a-f]

Suggested change

	if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
	if (
	not isinstance(value, Integral)
	or isinstance(value, bool)
	or value <= 0
	):
	raise ...
	value = int(value)

Otherwise this will fail on np.int32, np.int64

[z-a-f]

Suggested change

	if packed_dim % value:
	if packed_dim > 0 and packed_dim % value:

Or add if packed_dim <= 0: raise validation somewhere above.

[lisjin]

@claude review fixes

[claude]

Claude finished @lisjin's task in 41s —— View job

---

Review of fixes ([UPDATE] commit 33be250)

• Diff the update commit against the previous revision (f1c9a53..33be250)
• Check each fix against my earlier review and @z-a-f's inline comments
• Run tests (blocked: running pytest needs approval in this environment)
• Post review

Fixed from my earlier review

#	Item	Status
1	Bad config should fail at construction	✅ pruneopt.py:64 calls self._get_grouper_kwargs(group) in the __init__ loop. test_optimizer_rejects_invalid_config_before_warmup covers {}, explicit Nones, and both args set, with warmup_steps=10, on both axes.
2	Duplicate exactly-one check	✅ The duplicate is kept on purpose, and it's justified now that (1) depends on it.
3	README recipe row	✅ Now says "num_heads or head_dim".
4	Tuple-swap nit	✅ Replaced with if/else.
4	Indivisible case through optimizer	✅ test_optimizer_rejects_indivisible_head_dim added. It raises at step(), which is expected because the shape is only checked when the grouper is built.
4	Comment on the 12-zeros assertion	✅ Added.

All of these look right to me.

Still open: @z-a-f's inline comments

These weren't addressed in this update:

1. np.int64 / np.int32 rejected (attention.py:21). isinstance(np.int64(4), int) is False, so a head_dim that comes from a numpy computation or some config loaders raises "must be a positive integer". The suggested fix is right: check numbers.Integral, still exclude bool, then normalize with value = int(value) so later view() calls get a Python int.
   Fix this →
2. packed_dim == 0 (attention.py:23). 0 % value == 0, so an empty packed axis passes validation and resolves to num_heads=0 (with head_dim) or head_dim=0 (with num_heads). In practice this only matters for zero-size tensors, so it's an edge case. Still, a packed_dim <= 0 guard or a decision to skip such tensors would make the behavior explicit.

Minor observation

The construction-time check only enforces exactly one of the two arguments. Type and positivity checks, such as head_dim: 0 or a YAML string "64", still run only when the grouper is built at the first prune step, which is after warmup. If you want all argument errors to surface at construction, _get_grouper_kwargs could call a value-only check (the type/positivity part of _resolve_heads, factored out). This is optional; the divisibility check has to wait for the shape anyway.

Tests

I couldn't run the suite here because pytest needs approval in this environment. The new tests look correct when read.