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 0000000000..d912d9ce3e --- /dev/null +++ b/test/prototype/pat/test_pat_coupled_min_sparsity.py @@ -0,0 +1,714 @@ +# 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, + ) + if prox_type == "ProxGroupLasso" and not latent: + with self.assertRaisesRegex( + ValueError, + "latent_weights=False supports only hard constraints", + ): + _optimizer([group, untouched], lr=1.0, latent_weights=latent) + continue + 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_subset_mesh.py b/test/prototype/pat/test_pat_subset_mesh.py index f3a8cd45ab..0f9ecf3b91 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,117 @@ def _step(optimizer, params): param.grad = torch.zeros_like(param) optimizer.step() + def test_coupled_subset_mesh_materialized_state(self): + self._init_process_group() + try: + mesh = DeviceMesh("cpu", [1, 2]) + participant = mesh.get_coordinate() is not None + scores = torch.tensor( + [100.0, 1.0, 10000.0, 10.0, 1000.0, 2.0, 100000.0, 20.0] + ) + originals = [scores.repeat(5, 1), scores[:, None].repeat(1, 3)] + originals[0][-1, 0] = 0.0 + originals[1][0, -1] = 0.0 + selected = torch.tensor([1, 3, 5, 7]) + expected = [p.clone() for p in originals] + expected[0][:, selected] = 0.0 + expected[1][selected] = 0.0 + counts = [int(p.eq(0).sum()) for p in expected] + total = sum(p.numel() for p in expected) + groupers = (Dim1Grouper, Dim0Grouper) + direct = [ + distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None) + for p in originals + ] + params = [ + torch.nn.Parameter( + distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None) + ) + for p in originals + ] + groups = [ + { + "params": [p], + "group_type": cls.__name__, + "prox_type": "CoupledMinSparsityConstraint", + "couple_key": "residual", + "min_sparsity": 0.5, + } + for p, cls in zip(params, groupers) + ] + optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.0)) + optimizer.relative_sparsity = 0.375 + optimizer.relative_factored_frac = 0.25 + with mock.patch.object( + executor, "distribute_tensor", wraps=executor.distribute_tensor + ) as redistributed: + result = executor.apply_coupled_prox( + [(p, cls, {}) for p, cls in zip(direct, groupers)], + CoupledMinSparsityConstraint(0.0, 0.5), + 0.5, + ) + for _ in range(2): + self._step(optimizer, params) + if participant: + self.assertEqual( + optimizer.relative_sparsity, sum(counts) / total + ) + else: + self.assertEqual(optimizer.relative_sparsity, 0.375) + self.assertEqual(optimizer.relative_factored_frac, 0.25) + for p in params: + self.assertNotIn("sparsity_frac", optimizer.state[p]) + self.assertEqual(redistributed.call_count, 6 if participant else 0) + for call in redistributed.call_args_list: + self.assertIn("src_data_rank", call.kwargs) + self.assertIsNone(call.kwargs["src_data_rank"]) + if participant: + self.assertEqual(result.zero_channels, 4) + self.assertEqual(result.channels, 8) + self.assertEqual(result.zero_elts, sum(counts)) + self.assertEqual(result.numel, total) + self.assertEqual(len(result.parameters), 2) + for i, (p, q, dense) in enumerate(zip(direct, params, expected)): + self.assertIs(result.parameters[i].parameter, p) + self.assertEqual(result.parameters[i].zero_elts, counts[i]) + self.assertEqual(result.parameters[i].numel, dense.numel()) + self.assertEqual(p.full_tensor(), dense, rtol=0, atol=0) + self.assertEqual(q.full_tensor(), dense, rtol=0, atol=0) + self.assertEqual( + optimizer.state[q]["sparsity_frac"], counts[i] / dense.numel() + ) + self.assertEqual(optimizer.relative_factored_frac, 0.0) + record = ( + True, + result.zero_channels, + result.zero_elts, + optimizer.relative_sparsity, + ) + else: + self.assertEqual(result.parameters, ()) + self.assertEqual( + ( + result.zero_elts, + result.numel, + result.zero_channels, + result.channels, + ), + (0, 0, 0, 0), + ) + record = ( + False, + optimizer.relative_sparsity, + optimizer.relative_factored_frac, + ) + records = [None] * self.world_size + dist.all_gather_object(records, record) + self.assertEqual(records[0], (False, 0.375, 0.25)) + self.assertEqual(records[1], records[2]) + self.assertEqual(records[1], (True, 4, sum(counts), sum(counts) / total)) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + def test_subset_mesh_state_values_and_metrics(self): self._init_process_group() try: diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md index bcbbfae3c8..f466a2c28b 100644 --- a/torchao/prototype/pat/README.md +++ b/torchao/prototype/pat/README.md @@ -51,7 +51,7 @@ Pruning configs are dictionaries that define which parameter groups to prune and `PruneOptimizer(..., latent_weights=False)` and `build_prune_optimizer(..., latent_weights=False)` step directly from projected parameters instead of restoring PAT's dense latent weights. The default remains `True`. Direct projection removes PAT's latent copy but changes the optimization trajectory; it does not remove the base optimizer's momentum or AdamW moments. Gradients or momentum may revive zeros during pruning; the existing healing mask freezes them once healing starts. -Direct mode supports only `MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint`, which ignore cumulative `gamma`. Other proximal maps are rejected at optimizer construction. Soft maps such as `ProxLasso`, `ProxGroupLasso`, and `ProxNuclearNorm` require `latent_weights=True` because applying cumulative shrinkage repeatedly to projected weights would over-regularize them. +Direct mode supports only `MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `CoupledMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint`, which ignore cumulative `gamma`. Other proximal maps are rejected at optimizer construction. Soft maps such as `ProxLasso`, `ProxGroupLasso`, and `ProxNuclearNorm` require `latent_weights=True` because applying cumulative shrinkage repeatedly to projected weights would over-regularize them. The mode is a constructor setting, not part of the delegated optimizer state dict. When resuming, reconstruct the same mode before loading both model and optimizer checkpoints. Cross-mode resume is unsupported: direct checkpoints do not contain the latent weights needed by default-mode restoration. @@ -84,6 +84,7 @@ Attention-head groupers accept exactly one positive integer `num_heads` or `head | `ProxNuclearNorm` | Soft-threshold singular values to shrink rank smoothly | | `MinSparsityConstraint(min_sparsity)` | Hard-zero the smallest-L2 `ceil(min_sparsity * n_groups)` groups in each tensor | | `GlobalMinSparsityConstraint(min_sparsity)` | Hard-zero one jointly scored group-count budget across all tensors in an optimizer parameter group | +| `CoupledMinSparsityConstraint(min_sparsity)` | Hard-zero the same shared-channel indices across all tensors with one `couple_key` | | `MinRankConstraint(min_sparsity)` | Hard-zero the smallest `ceil(min_sparsity * k)` singular values in each matrix | | `NMSparseConstraint(n_nonzero)` | Keep the largest-magnitude `n_nonzero` elements per group | @@ -95,18 +96,50 @@ Attention-head groupers accept exactly one positive integer `num_heads` or `head | **Row sparsity** | `Dim0Grouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Use the hard constraint for an exact target or group Lasso for a smooth regularization knob | | **Column sparsity** | `Dim1Grouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Drops input channels of a Linear layer | | **Conv filter sparsity** | `ConvFilterGrouper` | `MinSparsityConstraint` or `ProxGroupLasso` | Drops complete `(c_out, c_in)` filter slices | -| **Attention head sparsity** | `AttentionHeadGrouperDim0` and/or `AttentionHeadGrouperDim1` | `MinSparsityConstraint` | Configure matching `num_heads` or `head_dim` for the selected projections | +| **Attention head sparsity** | `AttentionHeadGrouperDim0` and/or `AttentionHeadGrouperDim1` | `MinSparsityConstraint` | Configure `num_heads` or fixed `head_dim` for the selected projections | | **Global structured sparsity** | Any supported non-SVD grouper | `GlobalMinSparsityConstraint` | Shares one group-count budget across every tensor in the optimizer parameter group; `score_type: rms` is the recommended default | +| **Shared residual-channel sparsity** | `Dim1Grouper` on readers and `Dim0Grouper` on writers | `CoupledMinSparsityConstraint` | One caller-defined cluster uses the same selected channel indices everywhere | | **Low-rank approximation, smooth** | `SVDGrouper` or `PackedSVDGrouper` | `ProxNuclearNorm` | Uses regularization-controlled rank decay | | **Low-rank approximation, exact target** | `SVDGrouper` or `PackedSVDGrouper` | `MinRankConstraint(min_sparsity)` | Zeros `ceil(min_sparsity * k)` and retains `k - ceil(min_sparsity * k)` singular values per matrix | -`MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint` are hard-zero maps: they ignore `reg_lambda` and `gamma` and are driven by their target argument. Set `min_sparsity_schedule: true` to ramp a minimum-sparsity or minimum-rank target cubically from the end of warmup to `healing_start_step`. Regardless of `prox_freq` alignment, PAT applies a hard constraint once at `healing_start_step - 1` so healing freezes the final target mask rather than an earlier mask. +`MinSparsityConstraint`, `GlobalMinSparsityConstraint`, `CoupledMinSparsityConstraint`, `MinRankConstraint`, and `NMSparseConstraint` are hard-zero maps: they ignore `reg_lambda` and `gamma` and are driven by their target argument. Set `min_sparsity_schedule: true` to ramp a minimum-sparsity or minimum-rank target cubically from the end of warmup to `healing_start_step`. Regardless of `prox_freq` alignment, PAT applies a hard constraint once at `healing_start_step - 1` so healing freezes the final target mask rather than an earlier mask. `GlobalMinSparsityConstraint` computes one budget as `ceil(min_sparsity * total_groups)` for each optimizer parameter group, not one budget per tensor and not a parameter-count budget. It jointly ranks the groups exposed by the configured grouper, and all parameters in that optimizer group must produce scores on the same device. Use `score_type: rms` by default when tensors have different group sizes because it normalizes L2 magnitude by `sqrt(group_size)`. Raw `l2` tends to favor retaining larger groups because their norms grow with group size, while `param_cost` divides by the full group size and more strongly favors removing groups that save more parameters. Padded `KElementGrouper` views are rejected because padding would distort both scoring and accounting; choose a `k` that divides every grouped dimension. For recipes whose initialization scales magnitudes by the inverse square root of the number of groups, set the optional `score_group_count_ref` positive integer in a global pruning group. This multiplies each tensor's scores by `sqrt(n_groups / score_group_count_ref)` before joint selection. It is separate from RMS normalization by group size, is not automatically enabled by `head_dim`, and is disabled by default. Different positive reference magnitudes produce the same mathematical ranking because the reference contributes one common scale factor across candidates. -For DTensor parameters, global selection requires full materialization of every grouped tensor on every rank before the shared top-k decision, followed by scattering the selected masks back to the original placements. The current internal `CACHE_FULL_TENSORS` policy gathers each DTensor once and retains all dense copies until selection and write-back finish, so peak dense memory is the sum of the tensors in the optimizer parameter group. A future lower-memory policy could release each copy after scoring, but would need to gather each DTensor again to apply the selected mask. Every rank gathers the same tensors and performs the same deterministic selection, so masks are expected to agree, but current PAT CI covers the DTensor API only at world size one rather than exercising true multi-rank execution. +For DTensor parameters, global selection requires full materialization of every grouped tensor on every rank before the shared top-k decision, followed by scattering the selected masks back to the original placements. The current internal `CACHE_FULL_TENSORS` policy gathers each DTensor once and retains all dense copies until selection and write-back finish, so peak dense memory is the sum of the tensors in the optimizer parameter group. A future lower-memory policy could release each copy after scoring, but would need to gather each DTensor again to apply the selected mask. Every rank gathers the same tensors and performs the same deterministic selection, so masks are expected to agree, and CPU Gloo regression tests exercise actual multi-rank execution, including subset meshes. + +### Coupled shared-axis pruning + +Global pruning shares a budget but can remove different indices in each tensor. Coupled pruning sums squared norms at each shared channel index, scores the combined norm (`rms`, `l2`, or `param_cost`), and removes one index set of size `ceil(min_sparsity * channels)` from every tensor with the same `couple_key`, even across optimizer parameter groups. + +```python +import torch +from torchao.prototype.pat.optim import PruneOptimizer + +reader = torch.nn.Parameter(torch.randn(4, 8)) # input-channel axis +writer = torch.nn.Parameter(torch.randn(8, 4)) # output-channel axis +shared = { + "prox_type": "CoupledMinSparsityConstraint", + "couple_key": "residual", + "min_sparsity": 0.25, +} +groups = [ + {**shared, "params": [reader], "group_type": "Dim1Grouper"}, + {**shared, "params": [writer], "group_type": "Dim0Grouper"}, +] +optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.01), healing_start_step=8) +for p in (reader, writer): + p.grad = torch.zeros_like(p) +optimizer.step() +``` + +An existing non-SVD prox can add a nested `coupled` dictionary containing `couple_key`, `group_type`, and `min_sparsity`. Its own per-parameter/global projection runs first, followed by coupling. The nested stage inherits only `min_sparsity_schedule`, `reg_lambda`, and `score_type`; specify its own target and optional `prox_freq` (default 1). All members must agree on effective target, schedule, score type, and coupled frequency. Scheduled coupling requires a finite `healing_start_step`; the final hard mask is applied at `healing_start_step - 1` even off cadence. Metrics count each participating parameter once using final literal zeros, and retain cached totals on skipped updates. SVD/coupled composition and padded layouts are unsupported. + +The caller must enumerate every relevant tensor and its channel axis. Coupling does not discover residual dependencies, handle normalization/bias dependencies automatically, resize modules, or produce a compact inference model. + +For DTensors, coupled selection materializes every clustered tensor on each mesh participant, selects one shared mask, and redistributes locally with `src_data_rank=None`. Peak dense memory is proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover the materializing reference on a subset excluding rank zero, including uneven row sharding and cached nonparticipant metrics. These correctness tests do not establish NCCL/FSDP2 throughput or model-training performance. SVD decompositions can dominate optimizer cost on wide tensors. Set `prox_freq: N` on a parameter group to run its grouper and proximal map every `N` optimizer steps. SVD groups using the hard `MinRankConstraint` reapply the proximal map during healing by default so the base optimizer cannot refill removed singular values. Soft SVD maps such as `ProxNuclearNorm` retain their existing behavior unless `prox_through_heal: true` is set explicitly, while `MinRankConstraint` can opt out with `prox_through_heal: false`. Setting `prox_through_heal: true` on a non-SVD grouper is invalid and rejected during optimizer construction. diff --git a/torchao/prototype/pat/__init__.py b/torchao/prototype/pat/__init__.py index 206be4fbb5..ebc6fe6752 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 c96688cd64..7801ee0020 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 01deeb8062..b4a0ef6965 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 5e4614f098..9c873e86b1 100644 --- a/torchao/prototype/pat/optim/prox_executor.py +++ b/torchao/prototype/pat/optim/prox_executor.py @@ -53,6 +53,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 +287,186 @@ 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 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. + + 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 + 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 + + 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 a9b6277f30..19c528471b 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,10 +22,19 @@ _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") _HARD_CONSTRAINT_PROX_TYPES = frozenset( { + "CoupledMinSparsityConstraint", "GlobalMinSparsityConstraint", "MinRankConstraint", "MinSparsityConstraint", @@ -62,6 +72,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, ( @@ -87,13 +102,15 @@ def __init__( 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; " "the ramp ends when the mask freezes." ) + self._coupled_clusters() + for group in self.regularized_param_groups(): + self._get_grouper_kwargs(group) self.relative_sparsity = 0 self.relative_factored_frac = 0 @@ -249,6 +266,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, ( @@ -346,6 +455,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: @@ -425,6 +538,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) @@ -442,14 +562,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 = ( @@ -473,6 +595,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 @@ -512,6 +637,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():