[pat] Add coupled shared-axis pruning

test/prototype/pat/test_pat_subset_mesh.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from datetime import timedelta
+from unittest import mock
 
 import torch
 import torch.distributed as dist
@@ -15,8 +16,13 @@
 from torch.testing._internal.common_utils import run_tests
 
 from test.prototype.pat.test_common import make_prox_kwargs
-from torchao.prototype.pat.group import Dim0Grouper
-from torchao.prototype.pat.optim import ProxGroupLasso, PruneOptimizer
+from torchao.prototype.pat.group import Dim0Grouper, Dim1Grouper
+from torchao.prototype.pat.optim import (
+    CoupledMinSparsityConstraint,
+    ProxGroupLasso,
+    PruneOptimizer,
+)
+from torchao.prototype.pat.optim import prox_executor as executor
 from torchao.prototype.pat.optim.prox_executor import apply_prox
 
 
@@ -48,6 +54,117 @@ def _step(optimizer, params):
             param.grad = torch.zeros_like(param)
         optimizer.step()
 
+    def test_coupled_subset_mesh_materialized_state(self):
+        self._init_process_group()
+        try:
+            mesh = DeviceMesh("cpu", [1, 2])
+            participant = mesh.get_coordinate() is not None
+            scores = torch.tensor(
+                [100.0, 1.0, 10000.0, 10.0, 1000.0, 2.0, 100000.0, 20.0]
+            )
+            originals = [scores.repeat(5, 1), scores[:, None].repeat(1, 3)]
+            originals[0][-1, 0] = 0.0
+            originals[1][0, -1] = 0.0
+            selected = torch.tensor([1, 3, 5, 7])
+            expected = [p.clone() for p in originals]
+            expected[0][:, selected] = 0.0
+            expected[1][selected] = 0.0
+            counts = [int(p.eq(0).sum()) for p in expected]
+            total = sum(p.numel() for p in expected)
+            groupers = (Dim1Grouper, Dim0Grouper)
+            direct = [
+                distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None)
+                for p in originals
+            ]
+            params = [
+                torch.nn.Parameter(
+                    distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None)
+                )
+                for p in originals
+            ]
+            groups = [
+                {
+                    "params": [p],
+                    "group_type": cls.__name__,
+                    "prox_type": "CoupledMinSparsityConstraint",
+                    "couple_key": "residual",
+                    "min_sparsity": 0.5,
+                }
+                for p, cls in zip(params, groupers)
+            ]
+            optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.0))
+            optimizer.relative_sparsity = 0.375
+            optimizer.relative_factored_frac = 0.25
+            with mock.patch.object(
+                executor, "distribute_tensor", wraps=executor.distribute_tensor
+            ) as redistributed:
+                result = executor.apply_coupled_prox(
+                    [(p, cls, {}) for p, cls in zip(direct, groupers)],
+                    CoupledMinSparsityConstraint(0.0, 0.5),
+                    0.5,
+                )
+                for _ in range(2):
+                    self._step(optimizer, params)
+                    if participant:
+                        self.assertEqual(
+                            optimizer.relative_sparsity, sum(counts) / total
+                        )
+                    else:
+                        self.assertEqual(optimizer.relative_sparsity, 0.375)
+                        self.assertEqual(optimizer.relative_factored_frac, 0.25)
+                        for p in params:
+                            self.assertNotIn("sparsity_frac", optimizer.state[p])
+            self.assertEqual(redistributed.call_count, 6 if participant else 0)
+            for call in redistributed.call_args_list:
+                self.assertIn("src_data_rank", call.kwargs)
+                self.assertIsNone(call.kwargs["src_data_rank"])
+            if participant:
+                self.assertEqual(result.zero_channels, 4)
+                self.assertEqual(result.channels, 8)
+                self.assertEqual(result.zero_elts, sum(counts))
+                self.assertEqual(result.numel, total)
+                self.assertEqual(len(result.parameters), 2)
+                for i, (p, q, dense) in enumerate(zip(direct, params, expected)):
+                    self.assertIs(result.parameters[i].parameter, p)
+                    self.assertEqual(result.parameters[i].zero_elts, counts[i])
+                    self.assertEqual(result.parameters[i].numel, dense.numel())
+                    self.assertEqual(p.full_tensor(), dense, rtol=0, atol=0)
+                    self.assertEqual(q.full_tensor(), dense, rtol=0, atol=0)
+                    self.assertEqual(
+                        optimizer.state[q]["sparsity_frac"], counts[i] / dense.numel()
+                    )
+                self.assertEqual(optimizer.relative_factored_frac, 0.0)
+                record = (
+                    True,
+                    result.zero_channels,
+                    result.zero_elts,
+                    optimizer.relative_sparsity,
+                )
+            else:
+                self.assertEqual(result.parameters, ())
+                self.assertEqual(
+                    (
+                        result.zero_elts,
+                        result.numel,
+                        result.zero_channels,
+                        result.channels,
+                    ),
+                    (0, 0, 0, 0),
+                )
+                record = (
+                    False,
+                    optimizer.relative_sparsity,
+                    optimizer.relative_factored_frac,
+                )
+            records = [None] * self.world_size
+            dist.all_gather_object(records, record)
+            self.assertEqual(records[0], (False, 0.375, 0.25))
+            self.assertEqual(records[1], records[2])
+            self.assertEqual(records[1], (True, 4, sum(counts), sum(counts) / total))
+        finally:
+            if dist.is_initialized():
+                dist.destroy_process_group()
+
     def test_subset_mesh_state_values_and_metrics(self):
         self._init_process_group()
         try:

torchao/prototype/pat/README.md

@@ -51,7 +51,7 @@ Pruning configs are dictionaries that define which parameter groups to prune and
 
 `PruneOptimizer(..., latent_weights=False)` and `build_prune_optimizer(..., latent_weights=False)` step directly from projected parameters instead of restoring PAT's dense latent weights. The default remains `True`. Direct projection removes PAT's latent copy but changes the optimization trajectory; it does not remove the base optimizer's momentum or AdamW moments. Gradients or momentum may revive zeros during pruning; the existing healing mask freezes them once healing starts.
 
-Direct mode supports only `MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint`, which ignore cumulative `gamma`. Other proximal maps are rejected at optimizer construction. Soft maps such as `ProxLasso`, `ProxGroupLasso`, and `ProxNuclearNorm` require `latent_weights=True` because applying cumulative shrinkage repeatedly to projected weights would over-regularize them.
+Direct mode supports only `MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `CoupledMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint`, which ignore cumulative `gamma`. Other proximal maps are rejected at optimizer construction. Soft maps such as `ProxLasso`, `ProxGroupLasso`, and `ProxNuclearNorm` require `latent_weights=True` because applying cumulative shrinkage repeatedly to projected weights would over-regularize them.
 
 The mode is a constructor setting, not part of the delegated optimizer state dict. When resuming, reconstruct the same mode before loading both model and optimizer checkpoints. Cross-mode resume is unsupported: direct checkpoints do not contain the latent weights needed by default-mode restoration.
 
@@ -84,6 +84,7 @@ Attention-head groupers accept exactly one positive integer `num_heads` or `head
 | `ProxNuclearNorm` | Soft-threshold singular values to shrink rank smoothly |
 | `MinSparsityConstraint(min_sparsity)` | Hard-zero the smallest-L2 `ceil(min_sparsity * n_groups)` groups in each tensor |
 | `GlobalMinSparsityConstraint(min_sparsity)` | Hard-zero one jointly scored group-count budget across all tensors in an optimizer parameter group |
+| `CoupledMinSparsityConstraint(min_sparsity)` | Hard-zero the same shared-channel indices across all tensors with one `couple_key` |
 | `MinRankConstraint(min_sparsity)` | Hard-zero the smallest `ceil(min_sparsity * k)` singular values in each matrix |
 | `NMSparseConstraint(n_nonzero)` | Keep the largest-magnitude `n_nonzero` elements per group |
 
@@ -95,18 +96,50 @@ 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` or `head_dim` for the selected projections |
+| **Attention head sparsity** | `AttentionHeadGrouperDim0` and/or `AttentionHeadGrouperDim1` | `MinSparsityConstraint` | Configure `num_heads` or fixed `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 |
+| **Shared residual-channel sparsity** | `Dim1Grouper` on readers and `Dim0Grouper` on writers | `CoupledMinSparsityConstraint` | One caller-defined cluster uses the same selected channel indices everywhere |
 | **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 |
 
-`MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint` are hard-zero maps: they ignore `reg_lambda` and `gamma` and are driven by their target argument. Set `min_sparsity_schedule: true` to ramp a minimum-sparsity or minimum-rank target cubically from the end of warmup to `healing_start_step`. Regardless of `prox_freq` alignment, PAT applies a hard constraint once at `healing_start_step - 1` so healing freezes the final target mask rather than an earlier mask.
+`MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `CoupledMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint` are hard-zero maps: they ignore `reg_lambda` and `gamma` and are driven by their target argument. Set `min_sparsity_schedule: true` to ramp a minimum-sparsity or minimum-rank target cubically from the end of warmup to `healing_start_step`. Regardless of `prox_freq` alignment, PAT applies a hard constraint once at `healing_start_step - 1` so healing freezes the final target mask rather than an earlier mask.
 
 `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.
+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, and CPU Gloo regression tests exercise actual multi-rank execution, including subset meshes.
+
+### Coupled shared-axis pruning
+
+Global pruning shares a budget but can remove different indices in each tensor. Coupled pruning sums squared norms at each shared channel index, scores the combined norm (`rms`, `l2`, or `param_cost`), and removes one index set of size `ceil(min_sparsity * channels)` from every tensor with the same `couple_key`, even across optimizer parameter groups.
+
+```python
+import torch
+from torchao.prototype.pat.optim import PruneOptimizer
+
+reader = torch.nn.Parameter(torch.randn(4, 8))  # input-channel axis
+writer = torch.nn.Parameter(torch.randn(8, 4))  # output-channel axis
+shared = {
+    "prox_type": "CoupledMinSparsityConstraint",
+    "couple_key": "residual",
+    "min_sparsity": 0.25,
+}
+groups = [
+    {**shared, "params": [reader], "group_type": "Dim1Grouper"},
+    {**shared, "params": [writer], "group_type": "Dim0Grouper"},
+]
+optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.01), healing_start_step=8)
+for p in (reader, writer):
+    p.grad = torch.zeros_like(p)
+optimizer.step()
+```
+
+An existing non-SVD prox can add a nested `coupled` dictionary containing `couple_key`, `group_type`, and `min_sparsity`. Its own per-parameter/global projection runs first, followed by coupling. The nested stage inherits only `min_sparsity_schedule`, `reg_lambda`, and `score_type`; specify its own target and optional `prox_freq` (default 1). All members must agree on effective target, schedule, score type, and coupled frequency. Scheduled coupling requires a finite `healing_start_step`; the final hard mask is applied at `healing_start_step - 1` even off cadence. Metrics count each participating parameter once using final literal zeros, and retain cached totals on skipped updates. SVD/coupled composition and padded layouts are unsupported.
+
+The caller must enumerate every relevant tensor and its channel axis. Coupling does not discover residual dependencies, handle normalization/bias dependencies automatically, resize modules, or produce a compact inference model.
+
+For DTensors, coupled selection materializes every clustered tensor on each mesh participant, selects one shared mask, and redistributes locally with `src_data_rank=None`. Peak dense memory is proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover the materializing reference on a subset excluding rank zero, including uneven row sharding and cached nonparticipant metrics. These correctness tests do not establish NCCL/FSDP2 throughput or model-training performance.
 
 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.
 

torchao/prototype/pat/__init__.py

@@ -16,6 +16,7 @@
     SVDGrouper,
 )
 from .optim import (  # noqa: F401
+    CoupledMinSparsityConstraint,
     ProxGroupLasso,
     ProxGroupLassoVectorized,
     ProxLasso,

torchao/prototype/pat/optim/__init__.py

@@ -11,6 +11,7 @@
 from .group_lasso import ProxGroupLasso, ProxGroupLassoVectorized  # noqa: F401
 from .lasso import ProxLasso  # noqa: F401
 from .min_sparsity import (  # noqa: F401
+    CoupledMinSparsityConstraint,
     GlobalMinSparsityConstraint,
     MinRankConstraint,
     MinSparsityConstraint,

torchao/prototype/pat/optim/min_sparsity.py

@@ -133,6 +133,22 @@ def zero_groups_(p: Tensor, zero_idx: Tensor) -> Tensor:
         return torch.tensor(zeros, device=p.device, dtype=torch.long)
 
 
+class CoupledMinSparsityConstraint(GlobalMinSparsityConstraint):
+    """Select one shared channel index set across a caller-defined cluster.
+
+    Scores combine squared norms from every participating tensor before one
+    top-k selection. ``PruneOptimizer`` owns clustering and scheduling.
+    """
+
+    def combine_scores(self, sq_norm_sum: Tensor, group_size_sum: int) -> Tensor:
+        norm = sq_norm_sum.sqrt()
+        if self.score_type == "rms":
+            return norm / math.sqrt(group_size_sum)
+        if self.score_type == "param_cost":
+            return norm / group_size_sum
+        return norm
+
+
 class MinRankConstraint(ProxMap, _TopKZeroMixin):
     """Zeros the smallest ``ceil(min_sparsity * k)`` singular values of an
     SVD-grouped tensor. Here the shared ``min_sparsity`` key is the fraction of