[lisjin]

Stack from ghstack (oldest at bottom):

• -> [pat] Add coupled shared-axis pruning #4985
• [pat] Accelerate coupled pruning for sharded DTensors #4986

Select one shared channel mask across caller-defined tensor clusters, including nested composition after ordinary projection. Validate effective cluster settings and layouts before optimizer mutation, force the final mask before healing, and account final literal zeros once. Integrate the preceding constructor validations: coupled constraints remain supported in direct mode, missing coupled fields raise clear errors, and composed soft maps require latent mode. Adapted from qpat ebc76a1 (#57), preserving TorchAO-specific SVD and recursion fixes.

Establish dense execution and a materializing DTensor reference, including uneven layouts, mesh-local redistribution, and cached subset-mesh nonparticipant metrics. This PR contains the coupled API, optimizer integration, and core correctness coverage; sharded acceleration is separated into the next PR.

Test Plan:

source ~/.ao/bin/activate
OMP_NUM_THREADS=1 python -m pytest test/prototype/pat -q
# 189 passed, 187 subtests passed in 40.33s
ruff check --select F,I torchao/prototype/pat/optim/pruneopt.py torchao/prototype/pat/optim/prox_executor.py test/prototype/pat/test_pat_coupled_min_sparsity.py test/prototype/pat/test_pat_coupled_sharded_parity.py
git diff --check

The complete stack passed after rebasing the coupled core and sharded acceleration onto the revised PR #4984.

Tags: stack_created_with_devmate

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 18 Pending

As of commit 795ef85 with merge base 7b0382e ():
💚 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 2m 41s —— View job

---

Review: coupled shared-axis pruning

• Gather context
• Review source changes (pruneopt.py, prox_executor.py, min_sparsity.py)
• Review tests
• Post review

I didn't find any correctness bugs. The core algorithm is sound. Squared norms are summed per shared index, then one top-k runs and the same zero_idx is applied to every member. Zeros are counted after stack.close(), so reshaped groupers have already written back before counting. On DTensor, distribute_tensor(..., src_data_rank=None) reshards each rank's local copy, which avoids any global-rank-0 assumption. Accounting is also correct: a coupled param is skipped in both the per-param path and the global path (the change from global_result.zero_elts to per-parameter sums at pruneopt.py:520-523), so it is only counted once, from final literal zeros. The forced final hard mask at healing_start_step - 1 works for nested stages because cfg["prox_type"] is set to CoupledMinSparsityConstraint before should_prune reads it. The tests cover these paths well.

I couldn't run the tests: the sandbox requires approval for pytest invocations.

Suggestions below, most important first.

1. _coupled_clusters() runs on every step, including warmup and healing (pruneopt.py:427)

On every call it:

• rebuilds the cluster dict;
• builds a meta-tensor grouper for every clustered param (_validate_coupled_specs);
• instantiates a CoupledMinSparsityConstraint.

It's called at the very top of step(), before the warmup/healing early return. In those steps coupled_clusters and coupled_param_ids are never used. On large models that is wasted work on every step for the whole healing phase.

Moving the call below the if self.num_steps < self.warmup_steps or is_healing: early-return block keeps the behavior that test_pruning_step_revalidates_changed_cluster relies on. Another option is to cache the result when nothing has changed. Fix this →

2. Layout logic is duplicated between _validate_coupled_specs and grouped_view (prox_executor.py:189-196 vs :276-285)

The preflight copies the Elem/Layer/in_dims == 1 shape rules from grouped_view. If someone changes grouped_view later (say, for a new grouper), the preflight could quietly start accepting or rejecting different layouts from the ones that actually run.

The surrogate is a plain meta tensor, so grouped_view(grouper) already works on it. Calling it and reading view.shape would give one source of truth.

3. except RuntimeError is broader than it looks (prox_executor.py:207)

NotImplementedError is a subclass of RuntimeError. If a grouper uses an op with no meta kernel, the error is reported as "unsupported tensor layout", which is misleading. Consider narrowing the catch, or including str(error) in the message.

4. Redundant check in __init__ (pruneopt.py:64-68)

The early "coupled" in group and not group.get("prox_type") check duplicates the one in _coupled_config. That code path is unreachable for groups without a prox_type, though, because regularized_param_groups() filters them out first. So the early check is the one that actually fires. A short comment explaining that would stop someone from deleting it as dead code.

5. Minor / documentation

• Nested stage off-cadence (pruneopt.py:569-571). When the parent prox runs but the nested coupled stage doesn't, the parent's results are skipped for coupled params. Their state["sparsity_frac"] then stays stale until the next coupled step. That matches "retain cached totals", but the README could say explicitly that per-parameter sparsity_frac is also cached.
• reg_lambda in _COUPLED_INHERITED_KEYS has no effect, because hard constraints ignore it. Inheriting it is harmless, but the README lists it as one of the inherited settings, which suggests it matters.
• Partial validation before raising. __init__ runs group.setdefault("gamma"/"reg_lambda") before _coupled_clusters() validates, so a failed construction leaves those keys in the base optimizer's param groups. This is mostly cosmetic, since the weights themselves aren't changed.
• combine_scores with param_cost divides by group_size_sum (all members combined). The per-tensor score divides by a single tensor's group size. Both are fine because the denominator is the same for every channel. That also means the three score types (rms/l2/param_cost) all select the same mask in coupled mode. It would help to say so in the docstring or README, so users don't expect score_type to change which channels get pruned.

Overall it's well structured. Validation happens before any weights are changed, the docs clearly say what coupling does not do, and the tests cover composition, scheduling, and subset meshes. The only suggestion I'd treat as more than a nit is #1.

[z-a-f]

nit: Could we preserve float64 (also promote float16/bfloat16)? This currently downcasts all double precision