From 9d2a48dfff2c65618b0053897f13a0c1f1fbdeaf Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 10:49:29 -0700 Subject: [PATCH 1/4] [INITIAL] Update [ghstack-poisoned] --- .../pat/test_pat_coupled_min_sparsity.py | 707 ++++++++++++++++++ .../pat/test_pat_coupled_sharded_parity.py | 335 +++++++++ test/prototype/pat/test_pat_subset_mesh.py | 156 +++- torchao/prototype/pat/README.md | 39 +- torchao/prototype/pat/__init__.py | 1 + torchao/prototype/pat/optim/__init__.py | 1 + torchao/prototype/pat/optim/min_sparsity.py | 16 + torchao/prototype/pat/optim/prox_executor.py | 338 +++++++++ torchao/prototype/pat/optim/pruneopt.py | 149 +++- 9 files changed, 1734 insertions(+), 8 deletions(-) create mode 100644 test/prototype/pat/test_pat_coupled_min_sparsity.py create mode 100644 test/prototype/pat/test_pat_coupled_sharded_parity.py diff --git a/test/prototype/pat/test_pat_coupled_min_sparsity.py b/test/prototype/pat/test_pat_coupled_min_sparsity.py new file mode 100644 index 00000000000..4c44d663b6a --- /dev/null +++ b/test/prototype/pat/test_pat_coupled_min_sparsity.py @@ -0,0 +1,707 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 sys +import unittest +from unittest import mock + +import torch +from torch.distributed.tensor import distribute_tensor, init_device_mesh +from torch.distributed.tensor.placement_types import Shard +from torch.testing._internal import common_utils + +from test.prototype.pat.test_common import DistributedTestMixin +from torchao.prototype.pat.group import ( + Dim0Grouper, + Dim1Grouper, + ElemGrouper, + Grouper, + KElementGrouper, + LayerGrouper, +) +from torchao.prototype.pat.optim import CoupledMinSparsityConstraint, PruneOptimizer +from torchao.prototype.pat.optim.prox_executor import apply_coupled_prox, grouped_view + +D, INNER, HIDDEN, VOCAB = 4, 4, 6, 3 +EXECUTOR = "torchao.prototype.pat.optim.prox_executor" +OPTIMIZER = "torchao.prototype.pat.optim.pruneopt" + + +def _coupled_group(params, target, group_type, **overrides): + return dict( + params=list(params), + prox_type="CoupledMinSparsityConstraint", + group_type=group_type, + couple_key="resid", + min_sparsity=target, + **overrides, + ) + + +def _zero_channels(p, axis): + view = p.detach() if axis == 0 else p.detach().transpose(0, 1) + return set(view.eq(0).all(dim=1).nonzero(as_tuple=True)[0].tolist()) + + +def _step(optimizer, params, steps=1, dense_gradient=False): + for _ in range(steps): + optimizer.zero_grad() + for p in params: + p.grad = torch.ones_like(p) if dense_gradient else torch.zeros_like(p) + optimizer.step() + + +def _optimizer(groups, lr=0.0, healing=100, latent_weights=True, momentum=0.0): + return PruneOptimizer( + torch.optim.SGD(groups, lr=lr, momentum=momentum), + healing_start_step=healing, + latent_weights=latent_weights, + ) + + +class ToyBlockStack(torch.nn.Module): + """Small reader/writer stack sharing one residual channel axis.""" + + def __init__(self): + super().__init__() + self.tok_embeddings = torch.nn.Parameter(torch.randn(VOCAB, D)) + for name, shape in ( + ("wv", (INNER, D)), + ("wo", (D, INNER)), + ("w1", (HIDDEN, D)), + ("w2", (D, HIDDEN)), + ): + setattr( + self, + name, + torch.nn.ParameterList( + torch.nn.Parameter(torch.randn(*shape)) for _ in range(2) + ), + ) + + def readers(self): + return [self.tok_embeddings, *self.wv, *self.w1] + + def writers(self): + return [*self.wo, *self.w2] + + def specs(self): + return [(p, Dim1Grouper, {}) for p in self.readers()] + [ + (p, Dim0Grouper, {}) for p in self.writers() + ] + + def param_groups(self, target, **overrides): + return [ + _coupled_group(self.readers(), target, "Dim1Grouper", **overrides), + _coupled_group(self.writers(), target, "Dim0Grouper", **overrides), + ] + + def nested_groups(self, target, **overrides): + groups = self.param_groups(target) + for group in groups: + group["coupled"] = { + key: group.pop(key) for key in ("couple_key", "min_sparsity") + } + group["coupled"]["group_type"] = group["group_type"] + group.update( + prox_type="MinSparsityConstraint", min_sparsity=0.0, **overrides + ) + return groups + + +class _CoupledAssertions: + def setUp(self): + super().setUp() + torch.manual_seed(0) + + def assert_shared_mask(self, model, target): + masks = [_zero_channels(p, 1) for p in model.readers()] + masks += [_zero_channels(p, 0) for p in model.writers()] + self.assertEqual(masks, [masks[0]] * len(masks)) + self.assertEqual(len(masks[0]), math.ceil(target * D)) + return masks[0] + + def assert_literal_metrics(self, optimizer, params): + total = sum(p.numel() for p in params) + zeros = sum(p.eq(0).sum().item() for p in params) + self.assertAlmostEqual(optimizer.relative_sparsity, zeros / total) + for p in params: + self.assertAlmostEqual( + optimizer.state[p]["sparsity_frac"], p.eq(0).sum().item() / p.numel() + ) + + +class TestCoupledExecutor(_CoupledAssertions, common_utils.TestCase): + def test_combined_scores_match_concatenated_norm(self): + a, b = torch.randn(4, 3), torch.randn(4, 5) + joint = torch.cat((a, b), dim=1) + squared = a.square().sum(dim=1) + b.square().sum(dim=1) + for score_type, divisor in ( + ("rms", math.sqrt(8)), + ("l2", 1), + ("param_cost", 8), + ): + with self.subTest(score_type=score_type): + prox = CoupledMinSparsityConstraint(0.0, 0.5, score_type) + scores = prox.combine_scores(squared, 8) + self.assertEqual( + scores, torch.linalg.vector_norm(joint, dim=1) / divisor + ) + self.assertEqual(scores, prox.score(joint)) + + @torch.no_grad() + def test_shared_mask_and_result_contract(self): + for target in (0.0, 0.3, 1.0): + with self.subTest(target=target): + model = ToyBlockStack() + params = model.readers() + model.writers() + result = apply_coupled_prox( + model.specs(), CoupledMinSparsityConstraint(0.0, target), target + ) + self.assert_shared_mask(model, target) + self.assertEqual(result.channels, D) + self.assertEqual(result.zero_channels, math.ceil(target * D)) + self.assertEqual(result.numel, sum(p.numel() for p in params)) + self.assertEqual( + result.zero_elts, sum(p.eq(0).sum().item() for p in params) + ) + self.assertEqual(len(result.parameters), len(params)) + for record, p in zip(result.parameters, params): + self.assertIs(record.parameter, p) + self.assertEqual(record.numel, p.numel()) + self.assertEqual(record.zero_elts, p.eq(0).sum().item()) + + def test_float64_scoring_and_noncontiguous_flatten_writeback(self): + for dtype in (torch.float32, torch.float64): + for cls, target in ((ElemGrouper, 0.25), (LayerGrouper, 1.0)): + with self.subTest(dtype=dtype, grouper=cls.__name__): + p = torch.arange(1, 13, dtype=dtype).reshape(3, 4).t() + self.assertFalse(p.is_contiguous()) + expected = p.clone() + if cls is LayerGrouper: + expected.zero_() + else: + expected[expected <= 3] = 0 + result = apply_coupled_prox( + [(p, cls, {})], + CoupledMinSparsityConstraint(0.0, target), + target, + ) + self.assertEqual(p, expected) + self.assertEqual(result.zero_elts, int(expected.eq(0).sum())) + self.assertEqual(result.numel, p.numel()) + + def test_selects_lowest_joint_score_not_individual_scores(self): + reader = torch.tensor([[0.01, 10.0, 0.01]]) + writer = torch.tensor([[10.0], [0.01], [0.01]]) + apply_coupled_prox( + [(reader, Dim1Grouper, {}), (writer, Dim0Grouper, {})], + CoupledMinSparsityConstraint(0.0, 1 / 3), + 1 / 3, + ) + self.assertEqual(_zero_channels(reader, 1), {2}) + self.assertEqual(_zero_channels(writer, 0), {2}) + + def test_counts_existing_zeros_not_only_chosen_channels(self): + p = torch.tensor([[0.0, 0.0], [0.0, 0.0], [0.0, 2.0], [3.0, 4.0]]) + result = apply_coupled_prox( + [(p, Dim0Grouper, {})], CoupledMinSparsityConstraint(0.0, 0.25), 0.25 + ) + self.assertEqual(result.zero_channels, 1) + self.assertEqual(len(_zero_channels(p, 0)), 2) + self.assertEqual(result.zero_elts, 5) + self.assertEqual(result.parameters[0].zero_elts, 5) + self.assertEqual(result.numel, 8) + + def test_invalid_views_leave_originals_unchanged(self): + cases = [ + (torch.randn(5, 2), Dim0Grouper, {}, "same leading dimension"), + (torch.randn(2, 3, 4), Grouper, {}, "2-D"), + (torch.empty(4, 2, device="meta"), Dim0Grouper, {}, "same device"), + (torch.randn(1, 5), KElementGrouper, {"k": 4}, "padded"), + ] + for bad, grouper, kwargs, message in cases: + with self.subTest(message=message): + good = torch.randn(4, 2) + originals = [p.clone() for p in (good, bad)] + with self.assertRaisesRegex(ValueError, message): + apply_coupled_prox( + [(good, Dim0Grouper, {}), (bad, grouper, kwargs)], + CoupledMinSparsityConstraint(0.0, 0.5), + 0.5, + ) + self.assertEqual(good, originals[0]) + if bad.device.type != "meta": + self.assertEqual(bad, originals[1]) + + +class TestCoupledConfig(_CoupledAssertions, common_utils.TestCase): + def test_cluster_mismatches_fail_at_construction(self): + for nested in (False, True): + for field, value in ( + ("min_sparsity", 0.25), + ("min_sparsity_schedule", True), + ("score_type", "l2"), + ("prox_freq", 2), + ): + with self.subTest(nested=nested, field=field): + model = ToyBlockStack() + groups = ( + model.nested_groups(0.5) if nested else model.param_groups(0.5) + ) + cfg = groups[1]["coupled"] if nested else groups[1] + cfg[field] = value + with self.assertRaisesRegex(ValueError, field): + _optimizer(groups) + + def test_implicit_defaults_match_explicit_defaults(self): + for nested in (False, True): + model = ToyBlockStack() + groups = model.nested_groups(0.5) if nested else model.param_groups(0.5) + cfg = groups[1]["coupled"] if nested else groups[1] + cfg.update(min_sparsity_schedule=False, score_type="rms", prox_freq=1) + optimizer = _optimizer(groups) + _step(optimizer, list(model.parameters())) + self.assert_shared_mask(model, 0.5) + self.assert_literal_metrics(optimizer, list(model.parameters())) + + def test_malformed_nested_config_fails_at_construction(self): + for value in (None, False, [], 1, "bad", {}): + with self.subTest(value=value): + group = ToyBlockStack().nested_groups(0.5)[0] + group["coupled"] = value + with self.assertRaises(ValueError): + _optimizer([group]) + + def test_required_fields_including_nested_target(self): + for nested in (False, True): + for field in ("group_type", "couple_key", "min_sparsity"): + with self.subTest(nested=nested, field=field): + model = ToyBlockStack() + group = ( + model.nested_groups(0.5) if nested else model.param_groups(0.5) + )[0] + if nested: + group["min_sparsity"] = 0.75 + del (group["coupled"] if nested else group)[field] + with self.assertRaisesRegex(ValueError, field): + _optimizer([group]) + + def test_rejects_missing_parent_svd_and_ambiguous_stages(self): + for case in ( + "no_parent", + "svd_parent", + "svd_coupled", + "packed_coupled", + "ambiguous", + ): + with self.subTest(case=case): + group = ToyBlockStack().nested_groups(0.5)[0] + if case == "no_parent": + del group["prox_type"] + elif case == "svd_parent": + group.update( + group_type="SVDGrouper", + prox_type="MinRankConstraint", + min_sparsity=0.5, + ) + elif case in ("svd_coupled", "packed_coupled"): + group["coupled"]["group_type"] = ( + "SVDGrouper" if case == "svd_coupled" else "PackedSVDGrouper" + ) + group["coupled"]["npack"] = 2 + else: + group.update( + prox_type="CoupledMinSparsityConstraint", + couple_key="resid", + min_sparsity=0.5, + ) + with self.assertRaises(ValueError): + _optimizer([group]) + + def test_invalid_frequency_and_key(self): + for nested in (False, True): + for field, values in ( + ("prox_freq", (0, -1, 1.5, True, "2")), + ("couple_key", ("", None, 1, False)), + ): + for value in values: + with self.subTest(nested=nested, field=field, value=value): + model = ToyBlockStack() + group = ( + model.nested_groups(0.5) + if nested + else model.param_groups(0.5) + )[0] + (group["coupled"] if nested else group)[field] = value + with self.assertRaisesRegex(ValueError, field): + _optimizer([group]) + + def test_nested_schedule_independently_requires_finite_healing(self): + for healing in (sys.maxsize, float("inf")): + group = ToyBlockStack().nested_groups(0.5, min_sparsity_schedule=False)[0] + group["coupled"]["min_sparsity_schedule"] = True + with self.assertRaisesRegex(ValueError, "finite healing_start_step"): + _optimizer([group], healing=healing) + + def test_layout_validation_precedes_optimizer_mutation(self): + reader = torch.nn.Parameter(torch.ones(3, 4)) + writer = torch.nn.Parameter(torch.ones(5, 2)) + groups = [ + _coupled_group([reader], 0.5, "Dim1Grouper"), + _coupled_group([writer], 0.5, "Dim0Grouper"), + ] + with self.assertRaisesRegex(ValueError, "same leading dimension"): + _optimizer(groups, lr=0.1) + writer = torch.nn.Parameter(torch.ones(4, 2)) + groups[1]["params"] = [writer] + optimizer = _optimizer(groups, lr=0.1) + with torch.no_grad(): + writer.set_(torch.ones(5, 2)) + originals = [p.detach().clone() for p in (reader, writer)] + with self.assertRaisesRegex(ValueError, "same leading dimension"): + _step(optimizer, [reader, writer], dense_gradient=True) + self.assertEqual([reader, writer], originals) + self.assertEqual(optimizer.num_steps, 0) + self.assertEqual(len(optimizer.state), 0) + + def test_dense_attention_strides_fail_before_optimizer_mutation(self): + p = torch.nn.Parameter(torch.arange(1.0, 17.0).reshape(4, 4).t()) + group = _coupled_group([p], 0.5, "AttentionHeadGrouperDim0", num_heads=2) + original = p.detach().clone() + with self.assertRaisesRegex(ValueError, "unsupported tensor layout"): + _optimizer([group], lr=0.1) + self.assertEqual(p, original) + p = torch.nn.Parameter(original.contiguous()) + group["params"] = [p] + optimizer = _optimizer([group], lr=0.1) + with torch.no_grad(): + p.set_(p.t()) + before = p.detach().clone() + with self.assertRaisesRegex(ValueError, "unsupported tensor layout"): + _step(optimizer, [p], dense_gradient=True) + self.assertEqual(p, before) + self.assertEqual(optimizer.num_steps, 0) + self.assertEqual(len(optimizer.state), 0) + + def test_pruning_step_revalidates_changed_cluster(self): + for field, value in ( + ("min_sparsity", 0.25), + ("min_sparsity_schedule", True), + ("score_type", "l2"), + ("prox_freq", 2), + ): + with self.subTest(field=field): + model = ToyBlockStack() + optimizer = _optimizer(model.nested_groups(0.5)) + _step(optimizer, list(model.parameters())) + optimizer.param_groups[1]["coupled"][field] = value + with self.assertRaisesRegex(ValueError, field): + _step(optimizer, list(model.parameters())) + + +class TestCoupledOptimizer(_CoupledAssertions, common_utils.TestCase): + def test_exact_shared_budget_with_and_without_latent_weights(self): + for latent in (False, True): + model = ToyBlockStack() + optimizer = _optimizer(model.param_groups(0.5), latent_weights=latent) + _step(optimizer, list(model.parameters()), steps=3) + self.assert_shared_mask(model, 0.5) + self.assert_literal_metrics(optimizer, list(model.parameters())) + self.assertEqual(optimizer.relative_sparsity, 0.5) + self.assertEqual( + ["latent" in optimizer.state[p] for p in model.parameters()], + [latent] * len(list(model.parameters())), + ) + + def test_nested_inherits_schedule_score_and_regularization(self): + model = ToyBlockStack() + groups = model.nested_groups( + 0.75, min_sparsity_schedule=True, score_type="param_cost", reg_lambda=0.7 + ) + optimizer = _optimizer(groups, healing=8) + with mock.patch( + f"{OPTIMIZER}.apply_coupled_prox", wraps=apply_coupled_prox + ) as apply: + _step(optimizer, list(model.parameters()), steps=2) + self.assertEqual(apply.call_count, 2) + prox = apply.call_args.args[1] + self.assertEqual((prox.score_type, prox.reg_lambda), ("param_cost", 0.7)) + self.assertEqual(apply.call_args_list[0].args[2], 0.0) + self.assertAlmostEqual(apply.call_args.args[2], 0.75 * (1 - (1 - 1 / 8) ** 3)) + + def test_nested_overrides_inherited_values(self): + model = ToyBlockStack() + groups = model.nested_groups( + 0.5, min_sparsity_schedule=True, score_type="param_cost", reg_lambda=0.7 + ) + for group in groups: + group["coupled"].update( + min_sparsity_schedule=False, score_type="l2", reg_lambda=0.2 + ) + optimizer = _optimizer(groups) + with mock.patch( + f"{OPTIMIZER}.apply_coupled_prox", wraps=apply_coupled_prox + ) as apply: + _step(optimizer, list(model.parameters())) + prox = apply.call_args.args[1] + self.assertEqual((prox.score_type, prox.reg_lambda), ("l2", 0.2)) + self.assertEqual(apply.call_args.args[2], 0.5) + self.assert_shared_mask(model, 0.5) + + def test_schedule_ramps_to_final_mask_at_healing_boundary(self): + for nested in (False, True): + model = ToyBlockStack() + groups = ( + model.nested_groups(0.75, min_sparsity_schedule=True) + if nested + else model.param_groups(0.75, min_sparsity_schedule=True) + ) + optimizer = _optimizer(groups, healing=8) + counts = [] + for _ in range(8): + _step(optimizer, list(model.parameters())) + counts.append(len(_zero_channels(model.tok_embeddings, 1))) + self.assertEqual(counts[0], 0) + self.assertEqual(counts, sorted(counts)) + self.assertGreater(len(set(counts)), 2) + self.assert_shared_mask(model, 0.75) + + def test_coupled_frequency_is_independent_and_skips_cache_total_metrics(self): + for skipped in ("parent", "coupled"): + p = torch.nn.Parameter(torch.arange(1.0, 17.0).reshape(4, 4)) + group = dict( + params=[p], + group_type="Dim0Grouper", + prox_type="MinSparsityConstraint", + min_sparsity=0.25, + prox_freq=4 if skipped == "parent" else 1, + coupled=dict( + group_type="Dim1Grouper", couple_key="resid", min_sparsity=0.5 + ), + ) + if skipped == "coupled": + group["coupled"]["prox_freq"] = 4 + optimizer = _optimizer([group]) + with mock.patch( + f"{OPTIMIZER}.apply_coupled_prox", wraps=apply_coupled_prox + ) as apply: + _step(optimizer, [p]) + cached = (optimizer.relative_sparsity, optimizer.relative_factored_frac) + self.assert_literal_metrics(optimizer, [p]) + _step(optimizer, [p]) + self.assertEqual( + (optimizer.relative_sparsity, optimizer.relative_factored_frac), cached + ) + self.assertEqual(apply.call_count, 2 if skipped == "parent" else 1) + self.assertEqual(len(_zero_channels(p, 1)), 2 if skipped == "parent" else 0) + self.assertEqual(len(_zero_channels(p, 0)), 0 if skipped == "parent" else 1) + + def test_final_prune_forces_coupled_frequency_before_healing(self): + for nested in (False, True): + for scheduled in (False, True): + model = ToyBlockStack() + groups = model.nested_groups(0.5) if nested else model.param_groups(0.5) + for group in groups: + (group["coupled"] if nested else group).update( + prox_freq=4, min_sparsity_schedule=scheduled + ) + optimizer = _optimizer(groups, healing=3) + _step(optimizer, list(model.parameters()), steps=3) + self.assertEqual(optimizer.num_steps, 3) + self.assert_shared_mask(model, 0.5) + + def test_healing_dense_gradients_and_momentum_cannot_revive_mask(self): + for latent in (False, True): + model = ToyBlockStack() + params = list(model.parameters()) + optimizer = _optimizer( + model.nested_groups(0.5), + lr=0.1, + healing=2, + latent_weights=latent, + momentum=0.9, + ) + _step(optimizer, params, steps=2, dense_gradient=True) + frozen = self.assert_shared_mask(model, 0.5) + masks = [p.eq(0) for p in params] + self.assertTrue( + any( + optimizer.state[p]["momentum_buffer"][mask].ne(0).any() + for p, mask in zip(params, masks) + ) + ) + _step(optimizer, params, steps=4, dense_gradient=True) + self.assertEqual(self.assert_shared_mask(model, 0.5), frozen) + for p, mask in zip(params, masks): + self.assertTrue(p[mask].eq(0).all()) + + def test_composition_order_and_literal_metrics_with_existing_individual_zero(self): + # Removing row 0 changes the lowest channel from 0 to 1. + for prox_type in ( + "MinSparsityConstraint", + "GlobalMinSparsityConstraint", + "ProxGroupLasso", + ): + for latent in (False, True): + p = torch.nn.Parameter( + torch.tensor( + [ + [0.01, 3.0, 0.01, 0.01], + [1.0, 0.1, 0.0, 3.0], + [1.0, 0.1, 2.0, 3.0], + [1.0, 0.1, 2.0, 3.0], + ] + ) + ) + other = torch.nn.Parameter(torch.ones(2, 2)) + group = dict( + params=[p], + prox_type=prox_type, + group_type="Dim0Grouper", + min_sparsity=0.25, + reg_lambda=1.525 if prox_type == "ProxGroupLasso" else 0.0, + coupled=dict( + group_type="Dim1Grouper", couple_key="resid", min_sparsity=0.25 + ), + ) + untouched = dict( + params=[other], + prox_type="MinSparsityConstraint", + group_type="Dim0Grouper", + min_sparsity=0.0, + ) + optimizer = _optimizer( + [group, untouched], lr=1.0, latent_weights=latent + ) + _step(optimizer, [p, other]) + self.assertEqual(_zero_channels(p, 0), {0}) + self.assertEqual(_zero_channels(p, 1), {1}) + self.assertEqual(p[1, 2], 0.0) + self.assert_literal_metrics(optimizer, [p, other]) + self.assertAlmostEqual(optimizer.relative_sparsity, 8 / 20) + + def test_global_heads_with_head_dim_and_nested_residual(self): + for latent in (False, True): + model = ToyBlockStack() + groups = [] + for params, axis, coupled in ((model.wv, 0, 1), (model.wo, 1, 0)): + groups.append( + dict( + params=list(params), + group_type=f"AttentionHeadGrouperDim{axis}", + head_dim=2, + prox_type="GlobalMinSparsityConstraint", + min_sparsity=0.25, + coupled=dict( + group_type=f"Dim{coupled}Grouper", + couple_key="resid", + min_sparsity=0.5, + ), + ) + ) + groups += [ + _coupled_group([model.tok_embeddings, *model.w1], 0.5, "Dim1Grouper"), + _coupled_group(model.w2, 0.5, "Dim0Grouper"), + ] + optimizer = _optimizer(groups, latent_weights=latent) + _step(optimizer, list(model.parameters()), steps=2) + self.assert_shared_mask(model, 0.5) + for params, axis in ((model.wv, 0), (model.wo, 1)): + zero_heads = sum( + len( + _zero_channels( + (p if axis == 0 else p.transpose(0, 1)).reshape( + INNER // 2, -1 + ), + 0, + ) + ) + for p in params + ) + self.assertGreaterEqual(zero_heads, 1) + self.assert_literal_metrics(optimizer, list(model.parameters())) + + +class TestCoupledDTensor( + DistributedTestMixin, _CoupledAssertions, common_utils.TestCase +): + @torch.no_grad() + def test_world_size_one_multishard_fallback_matches_dense(self): + model = ToyBlockStack() + dense_specs = [ + (p.detach().clone(), cls, kwargs) for p, cls, kwargs in model.specs() + ] + specs = [ + (distribute_tensor(p.clone(), self.mesh, [Shard(0), Shard(0)]), cls, kwargs) + for p, cls, kwargs in dense_specs + ] + prox = CoupledMinSparsityConstraint(0.0, 0.5) + dense = apply_coupled_prox(dense_specs, prox, 0.5) + with mock.patch(f"{EXECUTOR}.grouped_view", wraps=grouped_view) as view: + result = apply_coupled_prox(specs, prox, 0.5) + self.assertEqual(view.call_count, len(specs)) + self.assertEqual( + (result.zero_elts, result.numel, result.zero_channels, result.channels), + (dense.zero_elts, dense.numel, dense.zero_channels, dense.channels), + ) + for record, (p, _, _), (want, _, _) in zip( + result.parameters, specs, dense_specs + ): + self.assertIs(record.parameter, p) + self.assertEqual(p.full_tensor(), want) + self.assertEqual(record.zero_elts, want.eq(0).sum().item()) + + def test_nonparticipant_skips_before_grouping(self): + p = distribute_tensor(torch.randn(4, 2), self.mesh, [Shard(0), Shard(0)]) + original = p.to_local().clone() + with ( + mock.patch.object(p.device_mesh, "get_coordinate", return_value=None), + mock.patch( + f"{EXECUTOR}.grouped_view", + side_effect=AssertionError("unexpected grouping"), + ), + ): + result = apply_coupled_prox( + [(p, Dim0Grouper, {})], CoupledMinSparsityConstraint(0.0, 0.5), 0.5 + ) + self.assertEqual(result.parameters, ()) + self.assertEqual( + (result.zero_elts, result.numel, result.zero_channels, result.channels), + (0, 0, 0, 0), + ) + self.assertEqual(p.to_local(), original) + + def test_mixed_dense_or_mesh_rejection_leaves_originals_unchanged(self): + p = distribute_tensor(torch.randn(4, 2), self.mesh, [Shard(0), Shard(0)]) + different_mesh = init_device_mesh("cpu", (1,)) + alternatives = [ + torch.randn(4, 2), + distribute_tensor(torch.randn(4, 2), different_mesh, [Shard(0)]), + ] + for other, message in zip( + alternatives, ("cannot mix dense", "same DeviceMesh") + ): + original = p.to_local().clone() + other_original = ( + other.to_local() if hasattr(other, "to_local") else other + ).clone() + with self.assertRaisesRegex(ValueError, message): + apply_coupled_prox( + [(p, Dim0Grouper, {}), (other, Dim0Grouper, {})], + CoupledMinSparsityConstraint(0.0, 0.5), + 0.5, + ) + self.assertEqual(p.to_local(), original) + self.assertEqual( + other.to_local() if hasattr(other, "to_local") else other, + other_original, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/prototype/pat/test_pat_coupled_sharded_parity.py b/test/prototype/pat/test_pat_coupled_sharded_parity.py new file mode 100644 index 00000000000..d7dbed011ed --- /dev/null +++ b/test/prototype/pat/test_pat_coupled_sharded_parity.py @@ -0,0 +1,335 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 +from contextlib import ExitStack +from unittest import mock + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, distribute_tensor +from torch.distributed.tensor.placement_types import Replicate, Shard +from torch.testing._internal.common_utils import run_tests + +from test.prototype.pat.test_pat_subset_mesh import _GlooMultiProcessTestCase +from torchao.prototype.pat.group import ( + AttentionHeadGrouperDim0, + AttentionHeadGrouperDim1, + Dim0Grouper, + Dim1Grouper, +) +from torchao.prototype.pat.optim import CoupledMinSparsityConstraint +from torchao.prototype.pat.optim import prox_executor as executor + + +def _channel_view(p, grouper_cls, kwargs): + if grouper_cls in (Dim1Grouper, AttentionHeadGrouperDim1): + channels = kwargs.get("num_heads", p.size(1)) + return p.reshape(p.size(0), channels, -1).permute(1, 0, 2).reshape(channels, -1) + return p.reshape(kwargs.get("num_heads", p.size(0)), -1) + + +def _fixture(groupers): + # Different tensors dominate different channels, with large score gaps. + amplitudes = 10.0 ** torch.tensor([3, 0, 6, 1, 4, 2, 7, 5]) + specs = [] + for i, cls in enumerate(groupers): + width = 5 + 2 * i + scale = torch.where(torch.arange(8) % len(groupers) == i, 1.0, 0.001) + view = (amplitudes * scale)[:, None] * torch.arange(1.0, width + 1)[None, :] + view[-1, -1] = 0.0 # An isolated zero outside the usual selected channels. + p = view.t().contiguous() if cls is Dim1Grouper else view.clone() + specs.append((p, cls, {})) + return specs + + +def _dense_reference(specs, budget): + views = [_channel_view(p, cls, kwargs) for p, cls, kwargs in specs] + squared_scores = sum(view.float().square().sum(1) for view in views) + chosen = math.ceil(budget * views[0].size(0)) + indices = squared_scores.argsort()[:chosen] + expected = [(p.clone(), cls, kwargs) for p, cls, kwargs in specs] + for p, cls, kwargs in expected: + axis = int(cls in (Dim1Grouper, AttentionHeadGrouperDim1)) + width = p.size(axis) // views[0].size(0) + packed_indices = (indices[:, None] * width + torch.arange(width)).flatten() + p.index_fill_(axis, packed_indices, 0.0) + return expected, chosen + + +class _CoupledParityMixin: + def _distribute(self, originals, mesh, placements): + return [ + ( + distribute_tensor(p.clone(), mesh, placements, src_data_rank=None), + cls, + kwargs, + ) + for p, cls, kwargs in originals + ] + + def _apply(self, specs, budget, score_type, route): + prox = CoupledMinSparsityConstraint(0.0, budget, score_type) + with ExitStack() as stack: + if route == "fast": + self.assertIsNotNone(executor._coupled_fast_path(specs)) + sharded = stack.enter_context( + mock.patch.object( + executor, + "_coupled_prox_sharded", + wraps=executor._coupled_prox_sharded, + ) + ) + stack.enter_context( + mock.patch.object( + executor, + "grouped_view", + side_effect=AssertionError("materialized grouped view"), + ) + ) + stack.enter_context( + mock.patch.object( + DTensor, + "full_tensor", + side_effect=AssertionError("materialized DTensor"), + ) + ) + collectives = [] + for name in ( + "all_reduce", + "all_gather_into_tensor", + "all_gather", + "broadcast", + "reduce", + "reduce_scatter_tensor", + "all_to_all_single", + ): + collectives.append( + stack.enter_context( + mock.patch.object(dist, name, wraps=getattr(dist, name)) + ) + ) + else: + if route == "forced": + stack.enter_context( + mock.patch.object( + executor, "_coupled_fast_path", return_value=None + ) + ) + else: + self.assertIsNone(executor._coupled_fast_path(specs)) + stack.enter_context( + mock.patch.object( + executor, + "_coupled_prox_sharded", + side_effect=AssertionError("unexpected fast path"), + ) + ) + grouped = stack.enter_context( + mock.patch.object( + executor, "grouped_view", wraps=executor.grouped_view + ) + ) + redistributed = stack.enter_context( + mock.patch.object( + executor, "distribute_tensor", wraps=executor.distribute_tensor + ) + ) + result = executor.apply_coupled_prox(specs, prox, budget) + if route == "fast": + self.assertEqual(sharded.call_count, 1) + self.assertGreater(sum(call.call_count for call in collectives), 0) + self.assertLessEqual(sum(call.call_count for call in collectives), 3) + shard_group = executor._coupled_fast_path(specs)[0] + for collective in collectives: + for call in collective.call_args_list: + self.assertIs(call.kwargs["group"], shard_group) + else: + self.assertEqual(grouped.call_count, len(specs)) + self.assertEqual(redistributed.call_count, len(specs)) + for call in redistributed.call_args_list: + self.assertIn("src_data_rank", call.kwargs) + self.assertIsNone(call.kwargs["src_data_rank"]) + return result + + def _check_result(self, result, specs, expected, chosen): + self.assertEqual(result.zero_channels, chosen) + self.assertEqual(result.channels, 8) + self.assertEqual(len(result.parameters), len(specs)) + zeros = [] + sizes = [] + for entry, (p, _, _), (dense, cls, kwargs) in zip( + result.parameters, specs, expected + ): + full = p.full_tensor() + self.assertEqual(full.eq(0), dense.eq(0)) + self.assertEqual(full, dense, rtol=0, atol=0) + self.assertIs(entry.parameter, p) + zeros.append(int(dense.eq(0).sum())) + sizes.append(dense.numel()) + self.assertEqual(entry.zero_elts, zeros[-1]) + self.assertEqual(entry.numel, sizes[-1]) + self.assertEqual( + int(_channel_view(full, cls, kwargs).eq(0).all(1).sum()), chosen + ) + self.assertEqual(result.zero_elts, sum(zeros)) + self.assertEqual(result.numel, sum(sizes)) + + def _parity(self, originals, mesh, placements, route="fast", budgets=(0.26,)): + for score_type in ("rms", "l2", "param_cost"): + for budget in budgets: + with self.subTest( + score_type=score_type, budget=budget, placements=placements + ): + expected, chosen = _dense_reference(originals, budget) + actual = self._distribute(originals, mesh, placements) + forced = self._distribute(originals, mesh, placements) + a = self._apply(actual, budget, score_type, route) + b = self._apply(forced, budget, score_type, "forced") + self._check_result(a, actual, expected, chosen) + self._check_result(b, forced, expected, chosen) + self.assertEqual( + (a.zero_elts, a.numel, a.zero_channels, a.channels), + (b.zero_elts, b.numel, b.zero_channels, b.channels), + ) + + +class TestPATCoupledShardedParity(_CoupledParityMixin, _GlooMultiProcessTestCase): + @property + def world_size(self): + return 4 + + def test_mixed_readers_writers_and_replicated_axes(self): + self._init_process_group() + try: + originals = _fixture( + (Dim1Grouper, Dim0Grouper, Dim1Grouper, Dim0Grouper, Dim0Grouper) + ) + mesh = DeviceMesh("cpu", [0, 1, 2, 3]) + self._parity(originals, mesh, (Shard(0),), budgets=(0.0, 0.26, 1.0)) + double_specs = [(p.double(), cls, kwargs) for p, cls, kwargs in originals] + self._parity(double_specs, mesh, (Shard(0),)) + mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) + for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))): + self._parity(originals, mesh, placements) + finally: + dist.destroy_process_group() + + def test_readers_only_uneven_rows_and_writers_only(self): + self._init_process_group() + try: + mesh = DeviceMesh("cpu", [0, 1, 2, 3]) + for groupers in ( + (Dim1Grouper,) * 3, + (Dim0Grouper,) * 3, + (Dim1Grouper, Dim0Grouper), + ): + self._parity(_fixture(groupers), mesh, (Shard(0),)) + finally: + dist.destroy_process_group() + + def test_unsupported_layouts_and_groupers_materialize(self): + self._init_process_group() + try: + originals = _fixture((Dim1Grouper, Dim0Grouper)) + mesh = DeviceMesh("cpu", [0, 1, 2, 3]) + self._parity(originals, mesh, (Shard(1),), "fallback") + kwargs_specs = [ + (p, cls, {"start_dim": 0, "end_dim": 1}) for p, cls, _ in originals + ] + self._parity(kwargs_specs, mesh, (Shard(0),), "fallback") + head_specs = [ + ( + p.repeat_interleave(2, dim=1 if cls is Dim1Grouper else 0), + AttentionHeadGrouperDim1 + if cls is Dim1Grouper + else AttentionHeadGrouperDim0, + {"num_heads": 8}, + ) + for p, cls, _ in originals + ] + self._parity(head_specs, mesh, (Shard(0),), "fallback") + mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) + self._parity(originals, mesh, (Shard(0), Shard(0)), "fallback") + finally: + dist.destroy_process_group() + + def test_tied_and_near_tied_selection_agrees_on_all_ranks(self): + self._init_process_group() + try: + mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) + for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))): + for near_tied in (False, True): + for route in ("fast", "forced"): + with self.subTest( + placements=placements, near_tied=near_tied, route=route + ): + scales = torch.ones(8) + if near_tied: + scales += ( + torch.arange(8) * torch.finfo(torch.float32).eps + ) + originals = [ + (scales.repeat(8, 1), Dim1Grouper, {}), + (scales[:, None].repeat(1, 4), Dim0Grouper, {}), + ] + specs = self._distribute(originals, mesh, placements) + result = self._apply(specs, 0.5, "rms", route) + mask = specs[0][0].to_local().eq(0).all(0) + self.assertEqual(int(mask.sum()), 4) + shard_dim = next( + i + for i, p in enumerate(placements) + if isinstance(p, Shard) + ) + shard_rank = mesh.get_coordinate()[shard_dim] + self.assertEqual( + specs[1][0].to_local().eq(0).all(1), + mask.chunk(2)[shard_rank], + ) + record = ( + mask.tolist(), + result.zero_channels, + result.zero_elts, + result.numel, + ) + records = [None] * self.world_size + dist.all_gather_object(records, record) + self.assertEqual(records, [record] * self.world_size) + # Dense tie-breaking need not match; shared masks must. + expected = [ + (p.clone(), cls, kwargs) for p, cls, kwargs in originals + ] + expected[0][0][:, mask] = 0.0 + expected[1][0][mask] = 0.0 + self._check_result(result, specs, expected, 4) + finally: + dist.destroy_process_group() + + +class TestPATCoupledUnevenShards(_CoupledParityMixin, _GlooMultiProcessTestCase): + @property + def world_size(self): + return 3 + + def test_uneven_writers_fall_back_but_uneven_readers_are_fast(self): + self._init_process_group() + try: + mesh = DeviceMesh("cpu", [0, 1, 2]) + self._parity( + _fixture((Dim1Grouper, Dim0Grouper, Dim0Grouper)), + mesh, + (Shard(0),), + "fallback", + ) + self._parity(_fixture((Dim1Grouper,) * 3), mesh, (Shard(0),)) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + run_tests() diff --git a/test/prototype/pat/test_pat_subset_mesh.py b/test/prototype/pat/test_pat_subset_mesh.py index f3a8cd45ab1..bef79fbb012 100644 --- a/test/prototype/pat/test_pat_subset_mesh.py +++ b/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,152 @@ def _step(optimizer, params): param.grad = torch.zeros_like(param) optimizer.step() + def test_coupled_subset_mesh_fast_and_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) + for forced in (False, True): + with self.subTest(forced=forced): + 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 + route_patch = ( + mock.patch.object( + executor, "_coupled_fast_path", return_value=None + ) + if forced + else mock.patch.object( + executor, + "_coupled_prox_sharded", + wraps=executor._coupled_prox_sharded, + ) + ) + with ( + route_patch as route, + 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] + ) + if forced: + 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"]) + else: + self.assertEqual(route.call_count, 3 if participant else 0) + self.assertEqual(redistributed.call_count, 0) + 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: diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 240c9815e5d..589331c34d0 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -82,6 +82,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 | @@ -93,18 +94,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` 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 2-D DTensors with exact default `Dim0Grouper`/`Dim1Grouper` layouts, one `Shard(0)` mesh axis, replication elsewhere, and evenly split writers, coupling uses local masking and at most three mesh-scoped collectives per cluster, independent of tensor count. Other layouts fall back to full materialization; peak dense memory is then proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover four-rank 1-D and 2-D meshes, both replicated/sharded axis orders, three-rank uneven fallback, and a subset excluding rank zero. Reduction order can perturb tied or near-tied selection across layouts; exact cross-layout mask parity is only asserted for well-separated scores. 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. diff --git a/torchao/prototype/pat/__init__.py b/torchao/prototype/pat/__init__.py index 206be4fbb5e..ebc6fe6752c 100644 --- a/torchao/prototype/pat/__init__.py +++ b/torchao/prototype/pat/__init__.py @@ -16,6 +16,7 @@ SVDGrouper, ) from .optim import ( # noqa: F401 + CoupledMinSparsityConstraint, ProxGroupLasso, ProxGroupLassoVectorized, ProxLasso, diff --git a/torchao/prototype/pat/optim/__init__.py b/torchao/prototype/pat/optim/__init__.py index c96688cd640..7801ee00205 100644 --- a/torchao/prototype/pat/optim/__init__.py +++ b/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, diff --git a/torchao/prototype/pat/optim/min_sparsity.py b/torchao/prototype/pat/optim/min_sparsity.py index 01deeb80624..b4a0ef6965f 100644 --- a/torchao/prototype/pat/optim/min_sparsity.py +++ b/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 diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 43fe8bc1e38..76e9ae7a841 100644 --- a/torchao/prototype/pat/optim/prox_executor.py +++ b/torchao/prototype/pat/optim/prox_executor.py @@ -11,12 +11,14 @@ from typing import Any, Sequence import torch +import torch.distributed as dist from torch import Tensor from torch.distributed.tensor import distribute_tensor from torch.distributed.tensor.experimental import local_map from torch.distributed.tensor.placement_types import Partial, Shard from ..distributed_utils import _is_dtensor +from ..group.dim import Dim0Grouper, Dim1Grouper from ..group.grouper import ElemGrouper, LayerGrouper from ..utils import get_index_linspace from .group_lasso import ProxGroupLasso, ProxGroupLassoVectorized @@ -53,6 +55,17 @@ class GlobalProxResult: numel: int +@dataclass(frozen=True) +class CoupledProxResult: + """Final literal zeros and shared channel selection for a coupled cluster.""" + + parameters: tuple[ParameterSparsity, ...] + zero_elts: int + numel: int + zero_channels: int + channels: int + + class GlobalDTensorMaterializationPolicy(Enum): """Internal policy for full DTensor views used by global selection.""" @@ -276,6 +289,331 @@ def grouped_view(grouper): return view, full +def _check_mesh_compatibility(params: Sequence[Tensor], what: str) -> bool: + """Return False when this rank sits outside the shared DeviceMesh.""" + dtensor_params = [p for p in params if _is_dtensor(p)] + if not dtensor_params: + return True + if len(dtensor_params) != len(params): + raise ValueError(f"{what} cannot mix dense tensors and DTensors in one group.") + first_mesh = dtensor_params[0].device_mesh + if any(p.device_mesh != first_mesh for p in dtensor_params[1:]): + raise ValueError(f"{what} requires all DTensors to use the same DeviceMesh.") + return first_mesh.get_coordinate() is not None + + +def _validate_coupled_specs(specs) -> bool: + """Preflight layouts using meta tensors, without gathering or changing weights.""" + params = [p for p, _, _ in specs] + if not params: + return True + if len({id(p) for p in params}) != len(params): + raise ValueError("CoupledMinSparsityConstraint requires distinct parameters") + if not _check_mesh_compatibility(params, "CoupledMinSparsityConstraint"): + return False + if any(p.device != params[0].device for p in params): + raise ValueError( + "CoupledMinSparsityConstraint requires every parameter to live on the same device." + ) + if any(not p.is_floating_point() for p in params): + raise ValueError( + "CoupledMinSparsityConstraint requires real floating-point tensors." + ) + n_channels = None + for p, cls, kwargs in specs: + if cls.__name__.endswith("SVDGrouper"): + raise ValueError("SVD groupers cannot be composed with coupled pruning") + surrogate = ( + torch.empty(p.shape, device="meta", dtype=p.dtype) + if _is_dtensor(p) + else torch.empty_strided(p.shape, p.stride(), device="meta", dtype=p.dtype) + ) + try: + with cls(surrogate, **kwargs) as grouper: + if hasattr(grouper, "_pad_size"): + raise ValueError( + "CoupledMinSparsityConstraint does not support padded KElementGrouper groups." + ) + if isinstance(grouper, ElemGrouper): + shape = (grouper.p.numel(), 1) + elif isinstance(grouper, LayerGrouper): + shape = (1, grouper.p.numel()) + elif getattr(grouper, "in_dims", 0) == 1 and grouper.p.dim() == 2: + shape = tuple(reversed(grouper.p.shape)) + else: + shape = grouper.p.shape + if len(shape) != 2: + raise ValueError( + "CoupledMinSparsityConstraint requires a grouper that produces a 2-D (n_channels, group_size) view." + ) + if n_channels is None: + n_channels = shape[0] + elif shape[0] != n_channels: + raise ValueError( + "CoupledMinSparsityConstraint requires every parameter to expose the same leading dimension." + ) + except RuntimeError as error: + raise ValueError( + "CoupledMinSparsityConstraint encountered an unsupported tensor layout" + ) from error + return True + + +def _shard_mesh_dim(p: Tensor) -> int | None: + """Find a single plain Shard(0) mesh axis, with replication elsewhere.""" + found = None + for mesh_dim, placement in enumerate(p.placements): + if type(placement) is Shard: + if placement.dim != 0 or found is not None: + return None + found = mesh_dim + elif not placement.is_replicate(): + return None + return found + + +def _coupled_fast_path(specs) -> tuple[Any, int, int] | None: + """Find eligible mesh-local reader reductions and writer gathers. + + Readers (Dim1) are row-sharded and hold partial norms for every channel. + Writers (Dim0) own complete channels, requiring evenly sized row blocks. + Reshaped groupers and other layouts use the materializing reference path. + """ + if not specs: + return None + mesh = None + mesh_dim = None + n_channels = None + defaults = {"start_dim": 1, "end_dim": -1} + for p, grouper_cls, grouper_kwargs in specs: + if grouper_cls not in (Dim0Grouper, Dim1Grouper): + return None + if any( + key not in defaults or value != defaults[key] + for key, value in grouper_kwargs.items() + ): + return None + if not _is_dtensor(p) or p.dim() != 2: + return None + p_mesh_dim = _shard_mesh_dim(p) + if p_mesh_dim is None: + return None + if mesh is None: + mesh, mesh_dim = p.device_mesh, p_mesh_dim + elif p.device_mesh != mesh or p_mesh_dim != mesh_dim: + return None + channels = p.size(1) if grouper_cls is Dim1Grouper else p.size(0) + if n_channels is None: + n_channels = channels + elif channels != n_channels: + return None + n_shard = mesh.size(mesh_dim) + coordinate = mesh.get_coordinate() + if coordinate is None: + return None + if any(cls is Dim0Grouper and p.size(0) % n_shard for p, cls, _ in specs): + return None + return mesh.get_group(mesh_dim), coordinate[mesh_dim], n_shard + + +def _coupled_prox_sharded( + specs, prox_map, min_sparsity, shard_group, shard_rank, n_shard +) -> CoupledProxResult: + """Apply a coupled mask with at most three mesh-scoped collectives. + + Aggregate all readers' partials together, gather all writers' channel + blocks together, then reduce one vector of literal-zero counts. Replica + mesh dimensions never participate in these collectives. + """ + reader_sq = None + writer_blocks = [] + group_size_sum = 0 + for p, cls, _ in specs: + local = p.to_local() + if cls is Dim1Grouper: + part = local.float().square().sum(dim=0) + reader_sq = part if reader_sq is None else reader_sq + part + group_size_sum += p.size(0) + else: + writer_blocks.append(local.float().square().sum(dim=1)) + group_size_sum += p.size(1) + + if reader_sq is not None: + dist.all_reduce(reader_sq, group=shard_group) + writer_sq = None + if writer_blocks: + stacked = torch.stack(writer_blocks).contiguous() + n_writers, local_channels = stacked.shape + gathered = torch.empty( + n_shard * n_writers, + local_channels, + dtype=stacked.dtype, + device=stacked.device, + ) + dist.all_gather_into_tensor(gathered, stacked, group=shard_group) + writer_sq = ( + gathered.view(n_shard, n_writers, local_channels) + .permute(1, 0, 2) + .reshape(n_writers, n_shard * local_channels) + .sum(dim=0) + ) + sq_norm_sum = ( + reader_sq + if writer_sq is None + else (writer_sq if reader_sq is None else reader_sq + writer_sq) + ) + scores = prox_map.combine_scores(sq_norm_sum, group_size_sum) + n_channels = scores.numel() + n_zero = math.ceil(min_sparsity * n_channels) + if n_zero <= 0: + zero_idx = torch.empty(0, dtype=torch.long, device=scores.device) + elif n_zero >= n_channels: + zero_idx = torch.arange(n_channels, device=scores.device) + else: + _, zero_idx = torch.topk(scores, k=n_zero, largest=False, sorted=False) + + zero_counts = [] + for p, cls, _ in specs: + local = p.to_local() + if cls is Dim1Grouper: + local[:, zero_idx] = 0.0 + else: + offset = shard_rank * local.size(0) + local_idx = ( + zero_idx[(zero_idx >= offset) & (zero_idx < offset + local.size(0))] + - offset + ) + local[local_idx] = 0.0 + zero_counts.append((local == 0).sum()) + counts = torch.stack(zero_counts).to(torch.long) + dist.all_reduce(counts, group=shard_group) + parameters = tuple( + ParameterSparsity(parameter=p, zero_elts=zeros, numel=p.numel()) + for (p, _, _), zeros in zip(specs, counts.tolist()) + ) + return CoupledProxResult( + parameters=parameters, + zero_elts=sum(p.zero_elts for p in parameters), + numel=sum(p.numel for p in parameters), + zero_channels=zero_idx.numel(), + channels=n_channels, + ) + + +def apply_coupled_prox( + specs: Sequence[tuple[Tensor, Any, dict[str, Any]]], + prox_map, + min_sparsity: float, + *, + materialization_policy: GlobalDTensorMaterializationPolicy = ( + GlobalDTensorMaterializationPolicy.CACHE_FULL_TENSORS + ), +) -> CoupledProxResult: + """Zero one shared set of channel indices across ``specs`` in place. + + Each spec holds (parameter, grouper class, grouper kwargs). Groupers must + expose 2-D (channels, group_size) views of the same leading dimension. + Reported element counts include all final literal zeros, including zeros + from a prior prox; callers must account each participating parameter once. + + Eligible Shard(0) layouts avoid full tensors. The fallback caches every + grouped full tensor, so peak dense memory is the sum of clustered tensors. + Floating-point reduction order can affect selection at tied or near-tied + score boundaries across layouts; selection is consistent within a mesh. + """ + if ( + materialization_policy + is not GlobalDTensorMaterializationPolicy.CACHE_FULL_TENSORS + ): + raise NotImplementedError( + f"Unsupported materialization policy: {materialization_policy}" + ) + if not 0.0 <= min_sparsity <= 1.0: + raise ValueError("min_sparsity must be in [0, 1]") + empty = CoupledProxResult( + parameters=(), zero_elts=0, numel=0, zero_channels=0, channels=0 + ) + if not specs: + return empty + if not _validate_coupled_specs(specs): + return empty + fast_path = _coupled_fast_path(specs) if dist.is_initialized() else None + if fast_path is not None: + return _coupled_prox_sharded(specs, prox_map, min_sparsity, *fast_path) + + with contextlib.ExitStack() as stack: + entries = [] + n_channels = None + sq_norm_sum = None + group_size_sum = 0 + for p, cls, kwargs in specs: + # Group the dense view so uneven DTensor reshapes cannot fail before + # the fallback gets a chance to materialize the parameter. + dense = p.full_tensor() if _is_dtensor(p) else p + grouper = stack.enter_context(cls(dense, **kwargs)) + if hasattr(grouper, "_pad_size"): + raise ValueError( + "CoupledMinSparsityConstraint does not support padded KElementGrouper groups." + ) + view, full = grouped_view(grouper) + if view.dim() != 2: + raise ValueError( + "CoupledMinSparsityConstraint requires a grouper that produces a 2-D " + f"(n_channels, group_size) view; {cls.__name__} produced shape {tuple(view.shape)}." + ) + if n_channels is None: + n_channels = view.size(0) + elif view.size(0) != n_channels: + raise ValueError( + "CoupledMinSparsityConstraint requires every parameter to expose the same leading dimension." + ) + sq_norm = torch.linalg.vector_norm(view.float(), dim=1).square() + sq_norm_sum = sq_norm if sq_norm_sum is None else sq_norm_sum + sq_norm + group_size_sum += view.size(1) + entries.append((p, grouper, view, full, dense)) + + scores = prox_map.combine_scores(sq_norm_sum, group_size_sum) + n_zero = math.ceil(min_sparsity * n_channels) + if n_zero <= 0: + zero_idx = torch.empty(0, dtype=torch.long, device=scores.device) + elif n_zero >= n_channels: + zero_idx = torch.arange(n_channels, device=scores.device) + else: + _, zero_idx = torch.topk(scores, k=n_zero, largest=False, sorted=False) + for p, grouper, view, full, dense in entries: + prox_map.zero_groups_(view, zero_idx) + if isinstance(grouper, (ElemGrouper, LayerGrouper)): + # Flattening a noncontiguous tensor may have allocated a copy. + full.copy_(view.reshape(full.shape)) + # Flush reshaped groupers' write-back before redistributing/counting. + stack.close() + zero_counts = [] + for p, _, _, _, dense in entries: + if _is_dtensor(p): + p.copy_( + distribute_tensor( + dense, + device_mesh=p.device_mesh, + placements=p.placements, + src_data_rank=None, + ) + ) + zero_counts.append((dense == 0).sum()) + parameters = tuple( + ParameterSparsity(parameter=p, zero_elts=zeros, numel=dense.numel()) + for (p, _, _, _, dense), zeros in zip( + entries, torch.stack(zero_counts).to(torch.long).tolist() + ) + ) + return CoupledProxResult( + parameters=parameters, + zero_elts=sum(p.zero_elts for p in parameters), + numel=sum(p.numel for p in parameters), + zero_channels=zero_idx.numel(), + channels=n_channels, + ) + + def apply_global_prox( params: Sequence[Tensor], prox_map, diff --git a/torchao/prototype/pat/optim/pruneopt.py b/torchao/prototype/pat/optim/pruneopt.py index c50fcc6b303..603eb83a7d6 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.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 sys from collections import defaultdict from collections.abc import Callable @@ -21,7 +22,15 @@ _sum_async_streams, ) from ..utils import instantiate_module -from .prox_executor import apply_global_prox, apply_prox_to_param +from .prox_executor import ( + _validate_coupled_specs, + apply_coupled_prox, + apply_global_prox, + apply_prox_to_param, +) + +_COUPLED_INHERITED_KEYS = ("min_sparsity_schedule", "reg_lambda", "score_type") +_GROUP_STATE_KEYS = ("params", "gamma", "num_steps", "factored_frac", "initial_lr") class PruneOptimizer(Optimizer): @@ -52,6 +61,11 @@ def __init__( ) -> None: # need to reconstruct these objects if loading checkpoint self.base_optimizer = base_optimizer + for group in self.param_groups: + if "coupled" in group and not group.get("prox_type"): + raise ValueError( + "nested coupled requires a regularized parent prox_type" + ) # need to store these attributes in state_dict for checkpoint assert warmup_steps < healing_start_step, ( @@ -72,6 +86,7 @@ def __init__( "the ramp ends when the mask freezes." ) + self._coupled_clusters() self.relative_sparsity = 0 self.relative_factored_frac = 0 @@ -227,6 +242,98 @@ def _build_group_artifacts(self, group: dict[str, Any]): ) return grouper_cls, self._get_grouper_kwargs(group) + def _coupled_config(self, group: dict[str, Any]) -> dict[str, Any] | None: + """Resolve top-level or nested coupling without inheriting target/frequency.""" + if group.get("prox_type") == "CoupledMinSparsityConstraint": + if "coupled" in group: + raise ValueError( + "top-level coupled prox cannot also carry a nested coupled stage" + ) + cfg = { + key: value + for key, value in group.items() + if key not in self.base_optimizer.defaults + and key not in _GROUP_STATE_KEYS + } + else: + if "coupled" not in group: + return None + if not group.get("prox_type"): + raise ValueError( + "nested coupled requires a regularized parent prox_type" + ) + nested = group["coupled"] + if not isinstance(nested, dict): + raise ValueError("nested coupled must be a dictionary") + cfg = {key: group[key] for key in _COUPLED_INHERITED_KEYS if key in group} + cfg.update(nested) + for key in ("couple_key", "group_type", "min_sparsity"): + if key not in cfg: + raise ValueError(f"coupled prune group requires {key!r}") + if not isinstance(cfg["couple_key"], str) or not cfg["couple_key"]: + raise ValueError("couple_key must be a nonempty string") + if any( + str(config.get("group_type", "")).endswith("SVDGrouper") + for config in (group, cfg) + ): + raise ValueError("SVD groupers cannot be composed with coupled pruning") + cfg["prox_type"] = "CoupledMinSparsityConstraint" + cfg.setdefault("min_sparsity_schedule", False) + cfg.setdefault("score_type", "rms") + cfg.setdefault("prox_freq", 1) + freq = cfg["prox_freq"] + if not isinstance(freq, int) or isinstance(freq, bool) or freq <= 0: + raise ValueError("coupled prox_freq must be a positive integer") + if cfg["min_sparsity_schedule"] and ( + self.healing_start_step == sys.maxsize + or not math.isfinite(self.healing_start_step) + ): + raise ValueError( + "coupled min_sparsity_schedule requires a finite healing_start_step" + ) + return cfg + + def _coupled_clusters(self) -> dict[str, tuple[dict[str, Any], list]]: + """Collect tensor-axis specs and validate effective cluster settings.""" + clusters = {} + seen = set() + for group in self.param_groups: + cfg = self._coupled_config(group) + if cfg is None: + continue + # Validate target and score type before any optimizer mutation. + instantiate_module( + "torchao.prototype.pat.optim.CoupledMinSparsityConstraint" + )(cfg.get("reg_lambda", 0.0), cfg["min_sparsity"], cfg["score_type"]) + key = cfg["couple_key"] + cls, kwargs = self._build_group_artifacts(cfg) + if key not in clusters: + clusters[key] = (cfg, []) + else: + first = clusters[key][0] + for field in ( + "min_sparsity", + "min_sparsity_schedule", + "score_type", + "prox_freq", + ): + if first[field] != cfg[field]: + raise ValueError( + f"coupled group {key!r} has conflicting {field!r}: {first[field]!r} vs {cfg[field]!r}" + ) + for p in group["params"]: + if not p.requires_grad: + continue + if id(p) in seen: + raise ValueError( + "each parameter may participate in only one coupled cluster" + ) + seen.add(id(p)) + clusters[key][1].append((p, cls, kwargs)) + for _, specs in clusters.values(): + _validate_coupled_specs(specs) + return clusters + def _build_global_prox_artifacts(self, group: dict[str, Any]): """Build global prox artifacts without resolving the scheduled budget.""" assert "min_sparsity" in group, ( @@ -263,6 +370,7 @@ def _build_prox_artifacts(self, group: dict[str, Any]): def should_prune(self, group: dict[str, Any], step: int) -> bool: """Run the group's prox map every ``prox_freq`` steps after warmup.""" hard_constraints = { + "CoupledMinSparsityConstraint", "GlobalMinSparsityConstraint", "MinRankConstraint", "MinSparsityConstraint", @@ -318,6 +426,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: # opted into through-heal instead reapply their prox map because dense # low-rank weights have no literal zeros to mask. Soft SVD maps default # to unconstrained healing unless they explicitly opt in. + coupled_clusters = self._coupled_clusters() + coupled_param_ids = { + id(p) for _, specs in coupled_clusters.values() for p, _, _ in specs + } healing_masks = {} is_healing = self.num_steps >= self.healing_start_step if is_healing: @@ -397,6 +509,13 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: self.state[p]["latent"].copy_(p) continue + if group["prox_type"] == "CoupledMinSparsityConstraint": + if self.latent_weights: + for p in group["params"]: + if p.requires_grad: + self.state[p]["latent"].copy_(p) + continue + if group["prox_type"] == "GlobalMinSparsityConstraint": prox_map, grouper_cls, grouper_kwargs = ( self._build_global_prox_artifacts(group) @@ -414,14 +533,16 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: score_group_count_ref=group.get("score_group_count_ref"), ) for param_result in global_result.parameters: + if id(param_result.parameter) in coupled_param_ids: + continue state = self.state[param_result.parameter] state["sparsity_frac"] = ( param_result.zero_elts / param_result.numel if param_result.numel else 0.0 ) - regularized_zeros += global_result.zero_elts - regularized_params += global_result.numel + regularized_zeros += param_result.zero_elts + regularized_params += param_result.numel continue prox_map, grouper_cls, grouper_kwargs, prox_kwargs = ( @@ -445,6 +566,9 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: if result is None: continue + if id(p) in coupled_param_ids: + # Account coupled tensors once, from final surviving zeros. + continue zero_elts = result.zero_elts zeros_are_summed = result.zeros_are_summed numel = result.numel @@ -484,6 +608,25 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: regularized_zeros += zero_elts regularized_params += numel + for cfg, specs in coupled_clusters.values(): + if not self.should_prune(cfg, self.num_steps): + all_groups_ran = False + continue + prox_map = instantiate_module( + "torchao.prototype.pat.optim.CoupledMinSparsityConstraint" + )(cfg.get("reg_lambda", 0.0), cfg["min_sparsity"], cfg["score_type"]) + result = apply_coupled_prox( + specs, prox_map, self._effective_min_sparsity(cfg) + ) + for param_result in result.parameters: + self.state[param_result.parameter]["sparsity_frac"] = ( + param_result.zero_elts / param_result.numel + if param_result.numel + else 0.0 + ) + regularized_zeros += param_result.zero_elts + regularized_params += param_result.numel + self.num_steps += 1 if torch.distributed.is_initialized() and _is_main_process(): From 87e4944b1c9d51ade8734386b25ab9d92ef89123 Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 10:49:29 -0700 Subject: [PATCH 2/4] Update (base update) [ghstack-poisoned] --- .../pat/test_pat_direct_projection.py | 562 ++++++++++++++++++ .../pat/test_pat_global_min_sparsity.py | 74 +++ test/prototype/pat/test_pat_grouper.py | 68 +++ torchao/prototype/pat/README.md | 14 +- torchao/prototype/pat/group/attention.py | 50 +- torchao/prototype/pat/optim/__init__.py | 2 + torchao/prototype/pat/optim/prox_executor.py | 27 +- torchao/prototype/pat/optim/pruneopt.py | 40 +- 8 files changed, 805 insertions(+), 32 deletions(-) create mode 100644 test/prototype/pat/test_pat_direct_projection.py diff --git a/test/prototype/pat/test_pat_direct_projection.py b/test/prototype/pat/test_pat_direct_projection.py new file mode 100644 index 00000000000..f03c6ae1446 --- /dev/null +++ b/test/prototype/pat/test_pat_direct_projection.py @@ -0,0 +1,562 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +"""Regressions for latent and direct projected PAT optimizer updates.""" + +import copy +import math +from unittest import mock + +import torch +from torch.testing._internal import common_utils + +from test.prototype.pat.test_common import TwoLayerMLP, optim_step +from torchao.prototype.pat.optim import ( + MinSparsityConstraint, + PruneOptimizer, + build_prune_optimizer, +) +from torchao.prototype.pat.utils import get_param_groups + + +def _param_groups(model, min_sparsity=0.5, prox_freq=1, scheduled=False): + config = { + (torch.nn.Linear, "weight"): { + "group_type": "Dim0Grouper", + "prox_type": "MinSparsityConstraint", + "min_sparsity": min_sparsity, + "min_sparsity_schedule": scheduled, + "prox_freq": prox_freq, + } + } + return get_param_groups(model, config, verbose=False) + + +def _base_optimizer(groups, name, lr=0.1): + if name == "sgd": + return torch.optim.SGD(groups, lr=lr, momentum=0.9, foreach=False) + return torch.optim.AdamW( + groups, lr=lr, betas=(0.8, 0.9), weight_decay=0.01, foreach=False + ) + + +def _optimizer( + model, + name="sgd", + latent_weights=None, + warmup_steps=2, + healing_start_step=8, + prox_freq=1, + scheduled=False, + min_sparsity=0.5, + lr=0.1, + reg_lambda=0.0, + builder=False, +): + base = _base_optimizer( + _param_groups(model, min_sparsity, prox_freq, scheduled), name, lr + ) + kwargs = {} if latent_weights is None else {"latent_weights": latent_weights} + if builder: + return build_prune_optimizer( + base, + prune_reg_lambda=reg_lambda, + prune_warmup_steps=warmup_steps, + prune_healing_start_step=healing_start_step, + **kwargs, + ) + return PruneOptimizer( + base, + warmup_steps=warmup_steps, + healing_start_step=healing_start_step, + reg_lambda=reg_lambda, + **kwargs, + ) + + +def _linear_model(): + with torch.random.fork_rng(devices=[]): + model = torch.nn.Linear(3, 6, dtype=torch.float32) + with torch.no_grad(): + model.weight.copy_(torch.arange(1, 19, dtype=torch.float32).reshape(6, 3) / 4) + model.bias.copy_(torch.arange(1, 7, dtype=torch.float32) / 10) + return model + + +def _mlp_and_data(steps): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + model = TwoLayerMLP(input_size=4, output_size=2).float() + generator = torch.Generator().manual_seed(1234) + inputs = torch.randn(steps, 4, dtype=torch.float32, generator=generator) + labels = torch.randint(0, 2, (steps,), generator=generator) + return model, inputs, labels + + +def _set_dense_gradients(model, step): + for p in model.parameters(): + gradient = torch.linspace(0.2, 1.1, p.numel(), dtype=p.dtype, device=p.device) + p.grad = gradient.reshape_as(p) * (1 + step / 10) + + +def _dense_step(model, optimizer, step): + optimizer.zero_grad() + _set_dense_gradients(model, step) + optimizer.step() + + +def _moment_keys(name): + return ("momentum_buffer",) if name == "sgd" else ("step", "exp_avg", "exp_avg_sq") + + +class TestDirectProjection(common_utils.TestCase): + def _assert_no_latent(self, optimizer): + self.assertNotIn("_state", optimizer.__dict__) + for state in optimizer.state.values(): + self.assertNotIn("latent", state) + for state in optimizer.state_dict()["state"].values(): + self.assertNotIn("latent", state) + + def _assert_pair_equal(self, model_a, optimizer_a, model_b, optimizer_b): + self.assertEqual(model_a.state_dict(), model_b.state_dict(), atol=0, rtol=0) + self.assertEqual( + optimizer_a.state_dict(), optimizer_b.state_dict(), atol=0, rtol=0 + ) + self.assertEqual(optimizer_a.num_steps, optimizer_b.num_steps) + + def test_omitted_default_matches_explicit_latent_trajectory_and_state(self): + for name in ("sgd", "adamw"): + for warmup in (0, 2): + for prox_freq in (1, 3): + with self.subTest(name=name, warmup=warmup, prox_freq=prox_freq): + model_a, inputs, labels = _mlp_and_data(12) + model_b = copy.deepcopy(model_a) + kwargs = { + "name": name, + "warmup_steps": warmup, + "prox_freq": prox_freq, + "scheduled": True, + } + default = _optimizer(model_a, **kwargs) + explicit = _optimizer(model_b, latent_weights=True, **kwargs) + self.assertTrue(default.latent_weights) + self.assertTrue(explicit.latent_weights) + for step in range(12): + optim_step(model_a, default, inputs, labels, step) + optim_step(model_b, explicit, inputs, labels, step) + self._assert_pair_equal(model_a, default, model_b, explicit) + for group in default.regularized_param_groups(): + for p in group["params"]: + self.assertIn("latent", default.state[p]) + zeroed = model_a.fc1.weight.detach().eq(0).all(dim=1) + self.assertTrue(zeroed.any()) + self.assertGreater( + default.state[model_a.fc1.weight]["latent"][zeroed] + .abs() + .sum() + .item(), + 0.0, + ) + + def test_direct_mode_never_uses_latent_and_keeps_base_moments(self): + for name in ("sgd", "adamw"): + for warmup in (0, 2): + for prox_freq in (1, 3): + with self.subTest(name=name, warmup=warmup, prox_freq=prox_freq): + model = _linear_model() + optimizer = _optimizer( + model, + name=name, + latent_weights=False, + warmup_steps=warmup, + prox_freq=prox_freq, + ) + self.assertEqual(len(optimizer.state), 0) + self._assert_no_latent(optimizer) + with ( + mock.patch.object( + optimizer, + "_init_latent_state", + side_effect=AssertionError, + ), + mock.patch.object( + optimizer, + "save_latent_params", + side_effect=AssertionError, + ), + mock.patch.object( + optimizer, + "restore_latent_params", + side_effect=AssertionError, + ), + ): + for step in range(12): + _dense_step(model, optimizer, step) + self._assert_no_latent(optimizer) + for p in model.parameters(): + state = optimizer.state[p] + for key in _moment_keys(name): + self.assertIn(key, state) + self.assertTrue( + torch.isfinite(state[key]).all() + ) + moment = state[_moment_keys(name)[-1]] + self.assertEqual(moment.shape, p.shape) + self.assertGreater(moment.abs().sum().item(), 0.0) + + def test_direct_updates_match_manual_base_step_and_projection(self): + for name in ("sgd", "adamw"): + for warmup in (0, 2): + for prox_freq in (1, 3): + with self.subTest(name=name, warmup=warmup, prox_freq=prox_freq): + model = _linear_model() + reference = copy.deepcopy(model) + optimizer = _optimizer( + model, + name=name, + latent_weights=False, + warmup_steps=warmup, + healing_start_step=100, + prox_freq=prox_freq, + ) + base = _base_optimizer(_param_groups(reference), name) + prox = MinSparsityConstraint(reg_lambda=0.0, min_sparsity=0.5) + gamma = 0.0 + for step in range(10): + previous_zeros = reference.weight.detach().eq(0).all(dim=1) + _dense_step(model, optimizer, step) + _dense_step(reference, base, step) + # The oracle never calls PruneOptimizer or its helpers. + if step >= warmup: + gamma += 0.1 + if (step - warmup) % prox_freq == 0: + with torch.no_grad(): + prox.apply_(reference.weight, gamma=gamma) + elif previous_zeros.any(): + self.assertTrue( + model.weight[previous_zeros].ne(0).any() + ) + self.assertEqual( + model.state_dict(), + reference.state_dict(), + atol=0, + rtol=0, + ) + for p, q in zip( + model.parameters(), reference.parameters(), strict=True + ): + for key in _moment_keys(name): + self.assertEqual( + optimizer.state[p][key], + base.state[q][key], + atol=0, + rtol=0, + ) + self.assertEqual(optimizer.param_groups[0]["gamma"], gamma) + self.assertEqual(optimizer.num_steps, step + 1) + self._assert_no_latent(optimizer) + + def test_zero_gradient_preserves_projected_point_on_skipped_steps(self): + for warmup in (0, 2): + with self.subTest(warmup=warmup): + model = _linear_model() + optimizer = PruneOptimizer( + torch.optim.SGD(_param_groups(model, prox_freq=3), lr=0.1), + warmup_steps=warmup, + healing_start_step=100, + latent_weights=False, + ) + for _ in range(warmup + 1): + for p in model.parameters(): + p.grad = torch.zeros_like(p) + optimizer.step() + projected = copy.deepcopy(model.state_dict()) + zeroed = model.weight.detach().eq(0).all(dim=1) + self.assertEqual(zeroed.sum().item(), 3) + gamma = optimizer.param_groups[0]["gamma"] + # Includes two skipped steps, a projection, then two more skips. + for _ in range(5): + for p in model.parameters(): + p.grad = torch.zeros_like(p) + optimizer.step() + gamma += 0.1 + self.assertEqual(model.state_dict(), projected, atol=0, rtol=0) + self.assertTrue(model.weight[zeroed].eq(0).all()) + self.assertEqual(optimizer.param_groups[0]["gamma"], gamma) + self._assert_no_latent(optimizer) + + def test_healing_freezes_actual_zeros_with_dense_gradients_and_momentum(self): + for name in ("sgd", "adamw"): + for latent_weights in (False, True): + for warmup in (0, 2): + with self.subTest( + name=name, latent_weights=latent_weights, warmup=warmup + ): + model = _linear_model() + optimizer = _optimizer( + model, + name=name, + latent_weights=latent_weights, + warmup_steps=warmup, + healing_start_step=5, + prox_freq=3, + scheduled=True, + ) + for step in range(5): + _dense_step(model, optimizer, step) + frozen = model.weight.detach().eq(0) + self.assertEqual(frozen.all(dim=1).sum().item(), 3) + self.assertTrue((~frozen).any()) + key = "momentum_buffer" if name == "sgd" else "exp_avg" + self.assertTrue( + optimizer.state[model.weight][key][frozen].ne(0).all() + ) + gamma = optimizer.param_groups[0]["gamma"] + for step in range(5, 9): + before = model.weight.detach().clone() + _set_dense_gradients(model, step) + self.assertTrue(model.weight.grad[frozen].ne(0).all()) + optimizer.step() + self.assertTrue(model.weight.grad[frozen].eq(0).all()) + self.assertTrue(model.weight.grad[~frozen].ne(0).all()) + self.assertTrue(model.weight[frozen].eq(0).all()) + self.assertFalse( + torch.equal(model.weight[~frozen], before[~frozen]) + ) + self.assertTrue( + optimizer.state[model.weight][key][frozen].ne(0).all() + ) + self.assertEqual(optimizer.param_groups[0]["gamma"], gamma) + if not latent_weights: + self._assert_no_latent(optimizer) + + def test_builder_forwards_mode_and_schedule_arguments(self): + for name in ("sgd", "adamw"): + for latent_weights in (None, True, False): + with self.subTest(name=name, latent_weights=latent_weights): + model_a = _linear_model() + model_b = copy.deepcopy(model_a) + kwargs = { + "name": name, + "latent_weights": latent_weights, + "warmup_steps": 2, + "healing_start_step": 6, + "prox_freq": 2, + "reg_lambda": 0.7, + } + built = _optimizer(model_a, builder=True, **kwargs) + direct = _optimizer(model_b, **kwargs) + self.assertIsInstance(built, PruneOptimizer) + self.assertEqual(built.latent_weights, latent_weights is not False) + self.assertEqual(built.warmup_steps, 2) + self.assertEqual(built.healing_start_step, 6) + self.assertEqual(built.param_groups[0]["reg_lambda"], 0.7) + for step in range(9): + _dense_step(model_a, built, step) + _dense_step(model_b, direct, step) + self._assert_pair_equal(model_a, built, model_b, direct) + if latent_weights is False: + self._assert_no_latent(built) + + def test_both_modes_reach_target_sparsity_before_healing(self): + for latent_weights in (False, True): + for scheduled in (False, True): + for target in (0.0, 0.5, 1.0): + with self.subTest( + latent_weights=latent_weights, + scheduled=scheduled, + target=target, + ): + model = _linear_model() + optimizer = _optimizer( + model, + latent_weights=latent_weights, + warmup_steps=2, + healing_start_step=7, + prox_freq=3, + scheduled=scheduled, + min_sparsity=target, + lr=0.01, + ) + n_rows = model.weight.shape[0] + for step in range(10): + _dense_step(model, optimizer, step) + if step in (2, 5) or step >= 6: + effective_target = target + if scheduled and step < 6: + progress = (step - 2) / (7 - 2) + effective_target *= 1 - (1 - progress) ** 3 + expected_rows = math.ceil(effective_target * n_rows) + # Step 6 is off cadence but must set the final mask. + rows = ( + model.weight.detach().eq(0).all(dim=1).sum().item() + ) + self.assertEqual(rows, expected_rows) + self.assertEqual( + optimizer.relative_sparsity, expected_rows / n_rows + ) + self.assertEqual( + optimizer.state[model.weight][ + "sparsity_frac" + ].item(), + expected_rows / n_rows, + ) + if not latent_weights: + self._assert_no_latent(optimizer) + + def test_state_dict_is_delegated_and_does_not_serialize_mode(self): + for name in ("sgd", "adamw"): + for latent_weights in (False, True): + with self.subTest(name=name, latent_weights=latent_weights): + model = _linear_model() + optimizer = _optimizer( + model, name=name, latent_weights=latent_weights + ) + for step in range(5): + _dense_step(model, optimizer, step) + base = optimizer.base_optimizer + with mock.patch.object( + base, "state_dict", wraps=base.state_dict + ) as delegated: + checkpoint = optimizer.state_dict() + delegated.assert_called_once_with() + self.assertEqual(checkpoint, base.state_dict(), atol=0, rtol=0) + self.assertEqual(set(checkpoint), {"state", "param_groups"}) + self.assertNotIn("latent_weights", checkpoint) + for group in checkpoint["param_groups"]: + self.assertNotIn("latent_weights", group) + for state in checkpoint["state"].values(): + self.assertNotIn("latent_weights", state) + resumed = _optimizer( + copy.deepcopy(model), name=name, latent_weights=latent_weights + ) + resumed.load_state_dict(copy.deepcopy(checkpoint)) + self.assertEqual(resumed.latent_weights, latent_weights) + self.assertEqual(resumed.state_dict(), checkpoint, atol=0, rtol=0) + + def test_reconstruct_load_and_continue_in_each_training_phase(self): + for name in ("sgd", "adamw"): + for latent_weights in (False, True): + for checkpoint_step in (1, 6, 10): + with self.subTest( + name=name, + latent_weights=latent_weights, + checkpoint_step=checkpoint_step, + ): + model_a, inputs, labels = _mlp_and_data(12) + kwargs = { + "name": name, + "latent_weights": latent_weights, + "warmup_steps": 2, + "healing_start_step": 8, + "prox_freq": 3, + "scheduled": True, + } + optimizer_a = _optimizer(model_a, **kwargs) + for step in range(checkpoint_step): + optim_step(model_a, optimizer_a, inputs, labels, step) + checkpoint = copy.deepcopy(optimizer_a.state_dict()) + model_b = copy.deepcopy(model_a) + optimizer_b = _optimizer(model_b, **kwargs) + optimizer_b.load_state_dict(copy.deepcopy(checkpoint)) + self.assertEqual(optimizer_b.latent_weights, latent_weights) + self._assert_pair_equal( + model_a, optimizer_a, model_b, optimizer_b + ) + for p, q in zip( + model_a.parameters(), model_b.parameters(), strict=True + ): + self.assertNotEqual(p.data_ptr(), q.data_ptr()) + for key, value in optimizer_a.state[p].items(): + if isinstance(value, torch.Tensor): + self.assertNotEqual( + value.data_ptr(), + optimizer_b.state[q][key].data_ptr(), + ) + for step in range(checkpoint_step, 12): + optim_step(model_a, optimizer_a, inputs, labels, step) + optim_step(model_b, optimizer_b, inputs, labels, step) + self._assert_pair_equal( + model_a, optimizer_a, model_b, optimizer_b + ) + if not latent_weights: + self._assert_no_latent(optimizer_a) + self._assert_no_latent(optimizer_b) + + def test_legacy_latent_checkpoint_loads_into_default_mode(self): + for name in ("sgd", "adamw"): + for prox_freq in (1, 3): + with self.subTest(name=name, prox_freq=prox_freq): + legacy_model = _linear_model() + base = _base_optimizer( + _param_groups(legacy_model, prox_freq=prox_freq), name + ) + group = base.param_groups[0] + group.update(gamma=0.0, reg_lambda=0.0, num_steps=0) + prox = MinSparsityConstraint(reg_lambda=0.0, min_sparsity=0.5) + warmup, checkpoint_step = 2, 5 + # Build legacy state independently, including skipped prox steps. + for step in range(9): + if step == checkpoint_step: + checkpoint = copy.deepcopy(base.state_dict()) + resumed_model = copy.deepcopy(legacy_model) + resumed = _optimizer( + resumed_model, + name=name, + warmup_steps=warmup, + healing_start_step=100, + prox_freq=prox_freq, + ) + self.assertTrue(resumed.latent_weights) + resumed.load_state_dict(copy.deepcopy(checkpoint)) + self.assertEqual( + resumed.state_dict(), checkpoint, atol=0, rtol=0 + ) + self.assertIn("latent", resumed.state[resumed_model.weight]) + _set_dense_gradients(legacy_model, step) + with torch.no_grad(): + if step == warmup: + base.state[legacy_model.weight]["latent"].copy_( + legacy_model.weight + ) + elif step > warmup: + legacy_model.weight.copy_( + base.state[legacy_model.weight]["latent"] + ) + base.step() + state = base.state[legacy_model.weight] + if step < warmup: + if "latent" not in state: + state["latent"] = ( + legacy_model.weight.detach().clone() + ) + else: + state["latent"].copy_(legacy_model.weight) + group["gamma"] += 0.1 + if (step - warmup) % prox_freq == 0: + zeros, _ = prox.apply_( + legacy_model.weight, gamma=group["gamma"] + ) + state["sparsity_frac"] = ( + zeros / legacy_model.weight.numel() + ) + group["num_steps"] = step + 1 + if step >= checkpoint_step: + _dense_step(resumed_model, resumed, step) + self.assertEqual( + legacy_model.state_dict(), + resumed_model.state_dict(), + atol=0, + rtol=0, + ) + self.assertEqual( + base.state_dict(), resumed.state_dict(), atol=0, rtol=0 + ) + self.assertEqual(resumed.num_steps, step + 1) + + +if __name__ == "__main__": + common_utils.run_tests() diff --git a/test/prototype/pat/test_pat_global_min_sparsity.py b/test/prototype/pat/test_pat_global_min_sparsity.py index eb742a05f11..881b252152a 100644 --- a/test/prototype/pat/test_pat_global_min_sparsity.py +++ b/test/prototype/pat/test_pat_global_min_sparsity.py @@ -95,6 +95,80 @@ 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_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) + + 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..3ce6843fc98 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,72 @@ 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() + self.assertEqual(int(params[0].eq(0).sum()), 12) + for kwargs in ({}, {"num_heads": 2, "head_dim": 3}): + invalid = {"group_type": group["group_type"], **kwargs} + with self.assertRaisesRegex(ValueError, "exactly one"): + opt._get_grouper_kwargs(invalid) + + 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..240c9815e5d 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -47,6 +47,12 @@ Pruning configs are dictionaries that define which parameter groups to prune and - regex pattern (string): for example, `:.*attn\.qkv\.weight` - module type and parameter name suffix (`(class, string)` tuple): for example, `(torch.nn.Linear, "weight")` +## Direct projection + +`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. + +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. + ## Groupers and proximal maps A pruning entry pairs a **grouper** with a **proximal map**. The grouper reshapes a tensor into `(n_groups, group_size)`, or exposes singular values for an SVD grouper, and the proximal map is then applied to that view. @@ -60,11 +66,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 | @@ -94,6 +102,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..cea1b667f76 100644 --- a/torchao/prototype/pat/group/attention.py +++ b/torchao/prototype/pat/group/attention.py @@ -4,25 +4,42 @@ # 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") + name, value = ( + ("head_dim", head_dim) if head_dim is not None else ("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 +47,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/__init__.py b/torchao/prototype/pat/optim/__init__.py index a5c8440d7ba..c96688cd640 100644 --- a/torchao/prototype/pat/optim/__init__.py +++ b/torchao/prototype/pat/optim/__init__.py @@ -26,6 +26,7 @@ def build_prune_optimizer( prune_reg_lambda: float, prune_warmup_steps: int = 0, prune_healing_start_step: int = sys.maxsize, + latent_weights: bool = True, ) -> PruneOptimizer: prune_opt_cls = PruneOptimizer # TODO: support other prune optimizers return prune_opt_cls( @@ -33,4 +34,5 @@ def build_prune_optimizer( warmup_steps=prune_warmup_steps, healing_start_step=prune_healing_start_step, reg_lambda=prune_reg_lambda, + latent_weights=latent_weights, ) diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 3e43639c2b7..43fe8bc1e38 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,8 @@ 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: + 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 +383,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 +401,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..c50fcc6b303 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.py @@ -37,6 +37,9 @@ class PruneOptimizer(Optimizer): pruned parameters are frozen. Must be greater than warmup_steps. reg_lambda: Regularization strength for the proximal updates. Can be overridden per parameter group. + latent_weights: Keep dense latent weights for AProx updates by default. + If False, step directly from projected weights without latent storage. + Reconstruct the same mode when loading an optimizer checkpoint. """ def __init__( @@ -45,6 +48,7 @@ def __init__( warmup_steps: int = 0, healing_start_step: int = sys.maxsize, reg_lambda: float = 0.0, + latent_weights: bool = True, ) -> None: # need to reconstruct these objects if loading checkpoint self.base_optimizer = base_optimizer @@ -56,6 +60,7 @@ def __init__( self.num_steps = 0 self.warmup_steps = warmup_steps self.healing_start_step = healing_start_step + self.latent_weights = latent_weights for group in self.regularized_param_groups(): group.setdefault("gamma", 0.0) @@ -186,7 +191,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": @@ -338,16 +349,18 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: del healing_masks if is_healing: self._apply_prox_to_through_heal_groups() - self._init_latent_state() + if self.latent_weights: + self._init_latent_state() self.num_steps += 1 return loss - if self.num_steps == self.warmup_steps: - # first PAT step: save latent params - self.save_latent_params() - else: - # restore latent params for base optimizer update - self.restore_latent_params() + if self.latent_weights: + if self.num_steps == self.warmup_steps: + # first PAT step: save latent params + self.save_latent_params() + else: + # restore latent params for base optimizer update + self.restore_latent_params() # call base optimizer step() method to update latent parameters loss = self.base_optimizer.step(closure=closure) # pyre-ignore[6] @@ -379,7 +392,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: # Keep latent parameters aligned with the base optimizer while # retaining cached sparsity and factorization metrics. for p in group["params"]: - if not p.requires_grad: + if not p.requires_grad or not self.latent_weights: continue self.state[p]["latent"].copy_(p) continue @@ -389,14 +402,16 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: self._build_global_prox_artifacts(group) ) params = [p for p in group["params"] if p.requires_grad] - for p in params: - self.state[p]["latent"].copy_(p) + if self.latent_weights: + for p in params: + self.state[p]["latent"].copy_(p) global_result = apply_global_prox( params, prox_map, 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] @@ -417,7 +432,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None: continue state = self.state[p] - state["latent"].copy_(p) + if self.latent_weights: + state["latent"].copy_(p) result = apply_prox_to_param( p, prox_map, From 56b9deede9efb8d780661afeef57935b95988811 Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 8 Oct 2026 11:02:53 -0700 Subject: [PATCH 3/4] [UPDATE] Split coupled pruning core from sharded acceleration [ghstack-poisoned] --- .../pat/test_pat_coupled_sharded_parity.py | 335 ------------------ test/prototype/pat/test_pat_subset_mesh.py | 205 +++++------ torchao/prototype/pat/README.md | 2 +- torchao/prototype/pat/optim/prox_executor.py | 153 +------- 4 files changed, 89 insertions(+), 606 deletions(-) delete mode 100644 test/prototype/pat/test_pat_coupled_sharded_parity.py diff --git a/test/prototype/pat/test_pat_coupled_sharded_parity.py b/test/prototype/pat/test_pat_coupled_sharded_parity.py deleted file mode 100644 index d7dbed011ed..00000000000 --- a/test/prototype/pat/test_pat_coupled_sharded_parity.py +++ /dev/null @@ -1,335 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# 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 -from contextlib import ExitStack -from unittest import mock - -import torch -import torch.distributed as dist -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor, distribute_tensor -from torch.distributed.tensor.placement_types import Replicate, Shard -from torch.testing._internal.common_utils import run_tests - -from test.prototype.pat.test_pat_subset_mesh import _GlooMultiProcessTestCase -from torchao.prototype.pat.group import ( - AttentionHeadGrouperDim0, - AttentionHeadGrouperDim1, - Dim0Grouper, - Dim1Grouper, -) -from torchao.prototype.pat.optim import CoupledMinSparsityConstraint -from torchao.prototype.pat.optim import prox_executor as executor - - -def _channel_view(p, grouper_cls, kwargs): - if grouper_cls in (Dim1Grouper, AttentionHeadGrouperDim1): - channels = kwargs.get("num_heads", p.size(1)) - return p.reshape(p.size(0), channels, -1).permute(1, 0, 2).reshape(channels, -1) - return p.reshape(kwargs.get("num_heads", p.size(0)), -1) - - -def _fixture(groupers): - # Different tensors dominate different channels, with large score gaps. - amplitudes = 10.0 ** torch.tensor([3, 0, 6, 1, 4, 2, 7, 5]) - specs = [] - for i, cls in enumerate(groupers): - width = 5 + 2 * i - scale = torch.where(torch.arange(8) % len(groupers) == i, 1.0, 0.001) - view = (amplitudes * scale)[:, None] * torch.arange(1.0, width + 1)[None, :] - view[-1, -1] = 0.0 # An isolated zero outside the usual selected channels. - p = view.t().contiguous() if cls is Dim1Grouper else view.clone() - specs.append((p, cls, {})) - return specs - - -def _dense_reference(specs, budget): - views = [_channel_view(p, cls, kwargs) for p, cls, kwargs in specs] - squared_scores = sum(view.float().square().sum(1) for view in views) - chosen = math.ceil(budget * views[0].size(0)) - indices = squared_scores.argsort()[:chosen] - expected = [(p.clone(), cls, kwargs) for p, cls, kwargs in specs] - for p, cls, kwargs in expected: - axis = int(cls in (Dim1Grouper, AttentionHeadGrouperDim1)) - width = p.size(axis) // views[0].size(0) - packed_indices = (indices[:, None] * width + torch.arange(width)).flatten() - p.index_fill_(axis, packed_indices, 0.0) - return expected, chosen - - -class _CoupledParityMixin: - def _distribute(self, originals, mesh, placements): - return [ - ( - distribute_tensor(p.clone(), mesh, placements, src_data_rank=None), - cls, - kwargs, - ) - for p, cls, kwargs in originals - ] - - def _apply(self, specs, budget, score_type, route): - prox = CoupledMinSparsityConstraint(0.0, budget, score_type) - with ExitStack() as stack: - if route == "fast": - self.assertIsNotNone(executor._coupled_fast_path(specs)) - sharded = stack.enter_context( - mock.patch.object( - executor, - "_coupled_prox_sharded", - wraps=executor._coupled_prox_sharded, - ) - ) - stack.enter_context( - mock.patch.object( - executor, - "grouped_view", - side_effect=AssertionError("materialized grouped view"), - ) - ) - stack.enter_context( - mock.patch.object( - DTensor, - "full_tensor", - side_effect=AssertionError("materialized DTensor"), - ) - ) - collectives = [] - for name in ( - "all_reduce", - "all_gather_into_tensor", - "all_gather", - "broadcast", - "reduce", - "reduce_scatter_tensor", - "all_to_all_single", - ): - collectives.append( - stack.enter_context( - mock.patch.object(dist, name, wraps=getattr(dist, name)) - ) - ) - else: - if route == "forced": - stack.enter_context( - mock.patch.object( - executor, "_coupled_fast_path", return_value=None - ) - ) - else: - self.assertIsNone(executor._coupled_fast_path(specs)) - stack.enter_context( - mock.patch.object( - executor, - "_coupled_prox_sharded", - side_effect=AssertionError("unexpected fast path"), - ) - ) - grouped = stack.enter_context( - mock.patch.object( - executor, "grouped_view", wraps=executor.grouped_view - ) - ) - redistributed = stack.enter_context( - mock.patch.object( - executor, "distribute_tensor", wraps=executor.distribute_tensor - ) - ) - result = executor.apply_coupled_prox(specs, prox, budget) - if route == "fast": - self.assertEqual(sharded.call_count, 1) - self.assertGreater(sum(call.call_count for call in collectives), 0) - self.assertLessEqual(sum(call.call_count for call in collectives), 3) - shard_group = executor._coupled_fast_path(specs)[0] - for collective in collectives: - for call in collective.call_args_list: - self.assertIs(call.kwargs["group"], shard_group) - else: - self.assertEqual(grouped.call_count, len(specs)) - self.assertEqual(redistributed.call_count, len(specs)) - for call in redistributed.call_args_list: - self.assertIn("src_data_rank", call.kwargs) - self.assertIsNone(call.kwargs["src_data_rank"]) - return result - - def _check_result(self, result, specs, expected, chosen): - self.assertEqual(result.zero_channels, chosen) - self.assertEqual(result.channels, 8) - self.assertEqual(len(result.parameters), len(specs)) - zeros = [] - sizes = [] - for entry, (p, _, _), (dense, cls, kwargs) in zip( - result.parameters, specs, expected - ): - full = p.full_tensor() - self.assertEqual(full.eq(0), dense.eq(0)) - self.assertEqual(full, dense, rtol=0, atol=0) - self.assertIs(entry.parameter, p) - zeros.append(int(dense.eq(0).sum())) - sizes.append(dense.numel()) - self.assertEqual(entry.zero_elts, zeros[-1]) - self.assertEqual(entry.numel, sizes[-1]) - self.assertEqual( - int(_channel_view(full, cls, kwargs).eq(0).all(1).sum()), chosen - ) - self.assertEqual(result.zero_elts, sum(zeros)) - self.assertEqual(result.numel, sum(sizes)) - - def _parity(self, originals, mesh, placements, route="fast", budgets=(0.26,)): - for score_type in ("rms", "l2", "param_cost"): - for budget in budgets: - with self.subTest( - score_type=score_type, budget=budget, placements=placements - ): - expected, chosen = _dense_reference(originals, budget) - actual = self._distribute(originals, mesh, placements) - forced = self._distribute(originals, mesh, placements) - a = self._apply(actual, budget, score_type, route) - b = self._apply(forced, budget, score_type, "forced") - self._check_result(a, actual, expected, chosen) - self._check_result(b, forced, expected, chosen) - self.assertEqual( - (a.zero_elts, a.numel, a.zero_channels, a.channels), - (b.zero_elts, b.numel, b.zero_channels, b.channels), - ) - - -class TestPATCoupledShardedParity(_CoupledParityMixin, _GlooMultiProcessTestCase): - @property - def world_size(self): - return 4 - - def test_mixed_readers_writers_and_replicated_axes(self): - self._init_process_group() - try: - originals = _fixture( - (Dim1Grouper, Dim0Grouper, Dim1Grouper, Dim0Grouper, Dim0Grouper) - ) - mesh = DeviceMesh("cpu", [0, 1, 2, 3]) - self._parity(originals, mesh, (Shard(0),), budgets=(0.0, 0.26, 1.0)) - double_specs = [(p.double(), cls, kwargs) for p, cls, kwargs in originals] - self._parity(double_specs, mesh, (Shard(0),)) - mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) - for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))): - self._parity(originals, mesh, placements) - finally: - dist.destroy_process_group() - - def test_readers_only_uneven_rows_and_writers_only(self): - self._init_process_group() - try: - mesh = DeviceMesh("cpu", [0, 1, 2, 3]) - for groupers in ( - (Dim1Grouper,) * 3, - (Dim0Grouper,) * 3, - (Dim1Grouper, Dim0Grouper), - ): - self._parity(_fixture(groupers), mesh, (Shard(0),)) - finally: - dist.destroy_process_group() - - def test_unsupported_layouts_and_groupers_materialize(self): - self._init_process_group() - try: - originals = _fixture((Dim1Grouper, Dim0Grouper)) - mesh = DeviceMesh("cpu", [0, 1, 2, 3]) - self._parity(originals, mesh, (Shard(1),), "fallback") - kwargs_specs = [ - (p, cls, {"start_dim": 0, "end_dim": 1}) for p, cls, _ in originals - ] - self._parity(kwargs_specs, mesh, (Shard(0),), "fallback") - head_specs = [ - ( - p.repeat_interleave(2, dim=1 if cls is Dim1Grouper else 0), - AttentionHeadGrouperDim1 - if cls is Dim1Grouper - else AttentionHeadGrouperDim0, - {"num_heads": 8}, - ) - for p, cls, _ in originals - ] - self._parity(head_specs, mesh, (Shard(0),), "fallback") - mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) - self._parity(originals, mesh, (Shard(0), Shard(0)), "fallback") - finally: - dist.destroy_process_group() - - def test_tied_and_near_tied_selection_agrees_on_all_ranks(self): - self._init_process_group() - try: - mesh = DeviceMesh("cpu", [[0, 1], [2, 3]]) - for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))): - for near_tied in (False, True): - for route in ("fast", "forced"): - with self.subTest( - placements=placements, near_tied=near_tied, route=route - ): - scales = torch.ones(8) - if near_tied: - scales += ( - torch.arange(8) * torch.finfo(torch.float32).eps - ) - originals = [ - (scales.repeat(8, 1), Dim1Grouper, {}), - (scales[:, None].repeat(1, 4), Dim0Grouper, {}), - ] - specs = self._distribute(originals, mesh, placements) - result = self._apply(specs, 0.5, "rms", route) - mask = specs[0][0].to_local().eq(0).all(0) - self.assertEqual(int(mask.sum()), 4) - shard_dim = next( - i - for i, p in enumerate(placements) - if isinstance(p, Shard) - ) - shard_rank = mesh.get_coordinate()[shard_dim] - self.assertEqual( - specs[1][0].to_local().eq(0).all(1), - mask.chunk(2)[shard_rank], - ) - record = ( - mask.tolist(), - result.zero_channels, - result.zero_elts, - result.numel, - ) - records = [None] * self.world_size - dist.all_gather_object(records, record) - self.assertEqual(records, [record] * self.world_size) - # Dense tie-breaking need not match; shared masks must. - expected = [ - (p.clone(), cls, kwargs) for p, cls, kwargs in originals - ] - expected[0][0][:, mask] = 0.0 - expected[1][0][mask] = 0.0 - self._check_result(result, specs, expected, 4) - finally: - dist.destroy_process_group() - - -class TestPATCoupledUnevenShards(_CoupledParityMixin, _GlooMultiProcessTestCase): - @property - def world_size(self): - return 3 - - def test_uneven_writers_fall_back_but_uneven_readers_are_fast(self): - self._init_process_group() - try: - mesh = DeviceMesh("cpu", [0, 1, 2]) - self._parity( - _fixture((Dim1Grouper, Dim0Grouper, Dim0Grouper)), - mesh, - (Shard(0),), - "fallback", - ) - self._parity(_fixture((Dim1Grouper,) * 3), mesh, (Shard(0),)) - finally: - dist.destroy_process_group() - - -if __name__ == "__main__": - run_tests() diff --git a/test/prototype/pat/test_pat_subset_mesh.py b/test/prototype/pat/test_pat_subset_mesh.py index bef79fbb012..0f9ecf3b91b 100644 --- a/test/prototype/pat/test_pat_subset_mesh.py +++ b/test/prototype/pat/test_pat_subset_mesh.py @@ -54,7 +54,7 @@ def _step(optimizer, params): param.grad = torch.zeros_like(param) optimizer.step() - def test_coupled_subset_mesh_fast_and_materialized_state(self): + def test_coupled_subset_mesh_materialized_state(self): self._init_process_group() try: mesh = DeviceMesh("cpu", [1, 2]) @@ -72,130 +72,95 @@ def test_coupled_subset_mesh_fast_and_materialized_state(self): counts = [int(p.eq(0).sum()) for p in expected] total = sum(p.numel() for p in expected) groupers = (Dim1Grouper, Dim0Grouper) - for forced in (False, True): - with self.subTest(forced=forced): - 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 - route_patch = ( - mock.patch.object( - executor, "_coupled_fast_path", return_value=None - ) - if forced - else mock.patch.object( - executor, - "_coupled_prox_sharded", - wraps=executor._coupled_prox_sharded, - ) - ) - with ( - route_patch as route, - 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] - ) - if forced: - 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"]) - else: - self.assertEqual(route.call_count, 3 if participant else 0) - self.assertEqual(redistributed.call_count, 0) + 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(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, + optimizer.relative_sparsity, sum(counts) / total ) - 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]) + 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( - records[1], (True, 4, sum(counts), sum(counts) / total) + 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() diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 589331c34d0..755e4e4beb4 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -137,7 +137,7 @@ An existing non-SVD prox can add a nested `coupled` dictionary containing `coupl 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 2-D DTensors with exact default `Dim0Grouper`/`Dim1Grouper` layouts, one `Shard(0)` mesh axis, replication elsewhere, and evenly split writers, coupling uses local masking and at most three mesh-scoped collectives per cluster, independent of tensor count. Other layouts fall back to full materialization; peak dense memory is then proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover four-rank 1-D and 2-D meshes, both replicated/sharded axis orders, three-rank uneven fallback, and a subset excluding rank zero. Reduction order can perturb tied or near-tied selection across layouts; exact cross-layout mask parity is only asserted for well-separated scores. These correctness tests do not establish NCCL/FSDP2 throughput or model-training performance. +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. diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py index 76e9ae7a841..fb9f63dc0d3 100644 --- a/torchao/prototype/pat/optim/prox_executor.py +++ b/torchao/prototype/pat/optim/prox_executor.py @@ -11,14 +11,12 @@ from typing import Any, Sequence import torch -import torch.distributed as dist from torch import Tensor from torch.distributed.tensor import distribute_tensor from torch.distributed.tensor.experimental import local_map from torch.distributed.tensor.placement_types import Partial, Shard from ..distributed_utils import _is_dtensor -from ..group.dim import Dim0Grouper, Dim1Grouper from ..group.grouper import ElemGrouper, LayerGrouper from ..utils import get_index_linspace from .group_lasso import ProxGroupLasso, ProxGroupLassoVectorized @@ -359,147 +357,6 @@ def _validate_coupled_specs(specs) -> bool: return True -def _shard_mesh_dim(p: Tensor) -> int | None: - """Find a single plain Shard(0) mesh axis, with replication elsewhere.""" - found = None - for mesh_dim, placement in enumerate(p.placements): - if type(placement) is Shard: - if placement.dim != 0 or found is not None: - return None - found = mesh_dim - elif not placement.is_replicate(): - return None - return found - - -def _coupled_fast_path(specs) -> tuple[Any, int, int] | None: - """Find eligible mesh-local reader reductions and writer gathers. - - Readers (Dim1) are row-sharded and hold partial norms for every channel. - Writers (Dim0) own complete channels, requiring evenly sized row blocks. - Reshaped groupers and other layouts use the materializing reference path. - """ - if not specs: - return None - mesh = None - mesh_dim = None - n_channels = None - defaults = {"start_dim": 1, "end_dim": -1} - for p, grouper_cls, grouper_kwargs in specs: - if grouper_cls not in (Dim0Grouper, Dim1Grouper): - return None - if any( - key not in defaults or value != defaults[key] - for key, value in grouper_kwargs.items() - ): - return None - if not _is_dtensor(p) or p.dim() != 2: - return None - p_mesh_dim = _shard_mesh_dim(p) - if p_mesh_dim is None: - return None - if mesh is None: - mesh, mesh_dim = p.device_mesh, p_mesh_dim - elif p.device_mesh != mesh or p_mesh_dim != mesh_dim: - return None - channels = p.size(1) if grouper_cls is Dim1Grouper else p.size(0) - if n_channels is None: - n_channels = channels - elif channels != n_channels: - return None - n_shard = mesh.size(mesh_dim) - coordinate = mesh.get_coordinate() - if coordinate is None: - return None - if any(cls is Dim0Grouper and p.size(0) % n_shard for p, cls, _ in specs): - return None - return mesh.get_group(mesh_dim), coordinate[mesh_dim], n_shard - - -def _coupled_prox_sharded( - specs, prox_map, min_sparsity, shard_group, shard_rank, n_shard -) -> CoupledProxResult: - """Apply a coupled mask with at most three mesh-scoped collectives. - - Aggregate all readers' partials together, gather all writers' channel - blocks together, then reduce one vector of literal-zero counts. Replica - mesh dimensions never participate in these collectives. - """ - reader_sq = None - writer_blocks = [] - group_size_sum = 0 - for p, cls, _ in specs: - local = p.to_local() - if cls is Dim1Grouper: - part = local.float().square().sum(dim=0) - reader_sq = part if reader_sq is None else reader_sq + part - group_size_sum += p.size(0) - else: - writer_blocks.append(local.float().square().sum(dim=1)) - group_size_sum += p.size(1) - - if reader_sq is not None: - dist.all_reduce(reader_sq, group=shard_group) - writer_sq = None - if writer_blocks: - stacked = torch.stack(writer_blocks).contiguous() - n_writers, local_channels = stacked.shape - gathered = torch.empty( - n_shard * n_writers, - local_channels, - dtype=stacked.dtype, - device=stacked.device, - ) - dist.all_gather_into_tensor(gathered, stacked, group=shard_group) - writer_sq = ( - gathered.view(n_shard, n_writers, local_channels) - .permute(1, 0, 2) - .reshape(n_writers, n_shard * local_channels) - .sum(dim=0) - ) - sq_norm_sum = ( - reader_sq - if writer_sq is None - else (writer_sq if reader_sq is None else reader_sq + writer_sq) - ) - scores = prox_map.combine_scores(sq_norm_sum, group_size_sum) - n_channels = scores.numel() - n_zero = math.ceil(min_sparsity * n_channels) - if n_zero <= 0: - zero_idx = torch.empty(0, dtype=torch.long, device=scores.device) - elif n_zero >= n_channels: - zero_idx = torch.arange(n_channels, device=scores.device) - else: - _, zero_idx = torch.topk(scores, k=n_zero, largest=False, sorted=False) - - zero_counts = [] - for p, cls, _ in specs: - local = p.to_local() - if cls is Dim1Grouper: - local[:, zero_idx] = 0.0 - else: - offset = shard_rank * local.size(0) - local_idx = ( - zero_idx[(zero_idx >= offset) & (zero_idx < offset + local.size(0))] - - offset - ) - local[local_idx] = 0.0 - zero_counts.append((local == 0).sum()) - counts = torch.stack(zero_counts).to(torch.long) - dist.all_reduce(counts, group=shard_group) - parameters = tuple( - ParameterSparsity(parameter=p, zero_elts=zeros, numel=p.numel()) - for (p, _, _), zeros in zip(specs, counts.tolist()) - ) - return CoupledProxResult( - parameters=parameters, - zero_elts=sum(p.zero_elts for p in parameters), - numel=sum(p.numel for p in parameters), - zero_channels=zero_idx.numel(), - channels=n_channels, - ) - - def apply_coupled_prox( specs: Sequence[tuple[Tensor, Any, dict[str, Any]]], prox_map, @@ -516,10 +373,9 @@ def apply_coupled_prox( Reported element counts include all final literal zeros, including zeros from a prior prox; callers must account each participating parameter once. - Eligible Shard(0) layouts avoid full tensors. The fallback caches every - grouped full tensor, so peak dense memory is the sum of clustered tensors. - Floating-point reduction order can affect selection at tied or near-tied - score boundaries across layouts; selection is consistent within a mesh. + DTensor execution caches every grouped full tensor, so peak dense memory + is the sum of clustered tensors. Every mesh participant selects from the + same materialized scores; nonparticipants skip processing. """ if ( materialization_policy @@ -537,9 +393,6 @@ def apply_coupled_prox( return empty if not _validate_coupled_specs(specs): return empty - fast_path = _coupled_fast_path(specs) if dist.is_initialized() else None - if fast_path is not None: - return _coupled_prox_sharded(specs, prox_map, min_sparsity, *fast_path) with contextlib.ExitStack() as stack: entries = [] From 0766f44197bab55b9d54feb446e147c7166fc090 Mon Sep 17 00:00:00 2001 From: lisjin Date: Thu, 15 Oct 2026 09:53:45 -0700 Subject: [PATCH 4/4] Update (base update) [ghstack-poisoned] --- .../pat/test_pat_direct_projection.py | 69 +++++++++++++++++++ .../pat/test_pat_global_min_sparsity.py | 41 +++++++++++ test/prototype/pat/test_pat_grouper.py | 42 +++++++++-- torchao/prototype/pat/README.md | 4 +- torchao/prototype/pat/group/attention.py | 7 +- torchao/prototype/pat/optim/pruneopt.py | 42 +++++++++-- 6 files changed, 190 insertions(+), 15 deletions(-) diff --git a/test/prototype/pat/test_pat_direct_projection.py b/test/prototype/pat/test_pat_direct_projection.py index f03c6ae1446..f3e8fc3465b 100644 --- a/test/prototype/pat/test_pat_direct_projection.py +++ b/test/prototype/pat/test_pat_direct_projection.py @@ -113,6 +113,75 @@ def _moment_keys(name): class TestDirectProjection(common_utils.TestCase): + def test_direct_mode_rejects_soft_prox_before_warmup(self): + for prox_type in ("ProxLasso", "ProxGroupLasso", "ProxNuclearNorm"): + for builder in (False, True): + with self.subTest(prox_type=prox_type, builder=builder): + param = torch.nn.Parameter(torch.ones(4, 4)) + group = { + "params": [param], + "prox_type": prox_type, + "group_type": ( + "SVDGrouper" + if prox_type == "ProxNuclearNorm" + else "Dim0Grouper" + ), + } + base = torch.optim.SGD([group], lr=0.1) + with self.assertRaisesRegex( + ValueError, + f"latent_weights=False.*hard constraints.*{prox_type}", + ): + if builder: + build_prune_optimizer( + base, + prune_reg_lambda=1.0, + prune_warmup_steps=10, + latent_weights=False, + ) + else: + PruneOptimizer( + base, + reg_lambda=1.0, + warmup_steps=10, + latent_weights=False, + ) + + def test_direct_mode_supports_hard_constraints(self): + for prox_type in ( + "MinSparsityConstraint", + "GlobalMinSparsityConstraint", + "NMSparseConstraint", + "MinRankConstraint", + ): + with self.subTest(prox_type=prox_type): + is_rank = prox_type == "MinRankConstraint" + param = torch.nn.Parameter( + torch.diag(torch.arange(1.0, 5.0)) + if is_rank + else torch.arange(1.0, 17.0).reshape(4, 4) + ) + group = { + "params": [param], + "prox_type": prox_type, + "group_type": "SVDGrouper" if is_rank else "Dim0Grouper", + "min_sparsity": 0.5, + "n_nonzero": 2, + } + optimizer = PruneOptimizer( + torch.optim.SGD([group], lr=0.1), + reg_lambda=1.0, + latent_weights=False, + ) + for _ in range(2): + param.grad = torch.zeros_like(param) + optimizer.step() + self._assert_no_latent(optimizer) + if is_rank: + self.assertEqual(int(torch.linalg.matrix_rank(param)), 2) + else: + self.assertEqual(int(param.eq(0).sum()), 8) + def _assert_no_latent(self, optimizer): self.assertNotIn("_state", optimizer.__dict__) for state in optimizer.state.values(): diff --git a/test/prototype/pat/test_pat_global_min_sparsity.py b/test/prototype/pat/test_pat_global_min_sparsity.py index 881b252152a..b08e3c969e6 100644 --- a/test/prototype/pat/test_pat_global_min_sparsity.py +++ b/test/prototype/pat/test_pat_global_min_sparsity.py @@ -168,6 +168,47 @@ def test_optimizer_forwards_reference(self): 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): diff --git a/test/prototype/pat/test_pat_grouper.py b/test/prototype/pat/test_pat_grouper.py index 3ce6843fc98..e40712f80e5 100644 --- a/test/prototype/pat/test_pat_grouper.py +++ b/test/prototype/pat/test_pat_grouper.py @@ -144,11 +144,45 @@ def test_optimizer_config(self): 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) - for kwargs in ({}, {"num_heads": 2, "head_dim": 3}): - invalid = {"group_type": group["group_type"], **kwargs} - with self.assertRaisesRegex(ValueError, "exactly one"): - opt._get_grouper_kwargs(invalid) + + 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): diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index 240c9815e5d..bcbbfae3c85 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -51,6 +51,8 @@ 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. + 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. ## Groupers and proximal maps @@ -93,7 +95,7 @@ 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` 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 | diff --git a/torchao/prototype/pat/group/attention.py b/torchao/prototype/pat/group/attention.py index cea1b667f76..caf3766323a 100644 --- a/torchao/prototype/pat/group/attention.py +++ b/torchao/prototype/pat/group/attention.py @@ -14,9 +14,10 @@ def _resolve_heads( ) -> tuple[int, int]: if (num_heads is None) == (head_dim is None): raise ValueError("specify exactly one of num_heads or head_dim") - name, value = ( - ("head_dim", head_dim) if head_dim is not None else ("num_heads", num_heads) - ) + 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: diff --git a/torchao/prototype/pat/optim/pruneopt.py b/torchao/prototype/pat/optim/pruneopt.py index c50fcc6b303..a9b6277f30c 100644 --- a/torchao/prototype/pat/optim/pruneopt.py +++ b/torchao/prototype/pat/optim/pruneopt.py @@ -23,6 +23,15 @@ from ..utils import instantiate_module from .prox_executor import apply_global_prox, apply_prox_to_param +_HARD_CONSTRAINT_PROX_TYPES = frozenset( + { + "GlobalMinSparsityConstraint", + "MinRankConstraint", + "MinSparsityConstraint", + "NMSparseConstraint", + } +) + class PruneOptimizer(Optimizer): """Wraps a base optimizer to apply proximal updates that induce sparsity @@ -39,6 +48,7 @@ class PruneOptimizer(Optimizer): overridden per parameter group. latent_weights: Keep dense latent weights for AProx updates by default. If False, step directly from projected weights without latent storage. + Direct mode supports only hard-constraint proximal maps. Reconstruct the same mode when loading an optimizer checkpoint. """ @@ -57,15 +67,27 @@ 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 self.latent_weights = latent_weights for group in self.regularized_param_groups(): + if ( + not self.latent_weights + and group["prox_type"] not in _HARD_CONSTRAINT_PROX_TYPES + ): + raise ValueError( + "latent_weights=False supports only hard constraints; " + f"got prox_type={group['prox_type']!r}. " + "Use latent_weights=True for soft proximal maps." + ) 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; " @@ -262,15 +284,9 @@ def _build_prox_artifacts(self, group: dict[str, Any]): def should_prune(self, group: dict[str, Any], step: int) -> bool: """Run the group's prox map every ``prox_freq`` steps after warmup.""" - hard_constraints = { - "GlobalMinSparsityConstraint", - "MinRankConstraint", - "MinSparsityConstraint", - "NMSparseConstraint", - } if ( step == self.healing_start_step - 1 - and group.get("prox_type") in hard_constraints + and group.get("prox_type") in _HARD_CONSTRAINT_PROX_TYPES ): # Materialize the final hard mask immediately before healing, # regardless of prox_freq alignment. @@ -281,6 +297,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")