Merge pull request pytorch#22 from pytorch-labs/jcaip/sparsity

[sparse] add sparsity, add wanda sparsifier to ao

Author: jcaip

test/sparsity/test_wanda.py

@@ -0,0 +1,114 @@
+import logging
+import unittest
+
+import torch
+from torch import nn
+from torchao.sparsity import WandaSparsifier
+from torch.ao.pruning import FakeSparsity
+from torch.nn.utils.parametrize import is_parametrized
+from torch.testing._internal.common_pruning import SimpleLinear
+from torch.testing._internal.common_utils import TestCase
+
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+
+
+class TestWandaSparsifier(TestCase):
+    """
+    Test Wanda Sparsifier
+    """
+
+    def test_prepare(self):
+        model = SimpleLinear()
+        sparsifier = WandaSparsifier()
+        sparsifier.prepare(model, config=None)
+        for g in sparsifier.groups:
+            module = g["module"]
+            # Check mask exists
+            assert hasattr(module.parametrizations["weight"][0], "mask")
+            # Check parametrization exists and is correct
+            assert is_parametrized(module, "weight")
+            assert type(module.parametrizations.weight[0]) == FakeSparsity
+            # check activation observer is present
+            assert hasattr(module, "activation_post_process")
+
+    def test_squash_mask(self):
+        # check observers and parameterizations removed
+        model = SimpleLinear()
+        sparsifier = WandaSparsifier()
+        sparsifier.prepare(model, config=None)
+        sparsifier.squash_mask()
+        for g in sparsifier.groups:
+            module = g["module"]
+            assert not is_parametrized(module, "weight")
+            assert not hasattr(module, "mask")
+            assert not hasattr(module, "activation_post_process")
+
+    def test_one_layer_mlp_2x4(self):
+        model = nn.Sequential(nn.Linear(8, 1))
+        weights = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]])
+        model[0].weight.data.copy_(weights.data)
+        X = torch.ones(1, 8)
+
+        sparsifier = WandaSparsifier(semi_structured_block_size=4)
+        sparsifier.prepare(model, config=None)
+
+        model(X)
+
+        sparsifier.step()
+        sparsifier.squash_mask()
+
+        sparsity = (model[0].weight == 0).float().mean()
+        assert sparsity == 0.5
+
+        expected_fc = torch.tensor([[0, 0, 3, 4, 0, 0, 7, 8]], dtype=torch.float32)
+        assert torch.allclose(model[0].weight.data, expected_fc, rtol=1e-05, atol=1e-07)
+
+    def test_one_layer_mlp_unstructured(self):
+        model = nn.Sequential(nn.Linear(4, 1))
+        weights = torch.tensor([[1, 2, 3, 4]], dtype=torch.float32)
+        model[0].weight.data.copy_(weights.data)
+        X = torch.tensor([[100, 10, 1, 0.1]], dtype=torch.float32)
+
+        sparsifier = WandaSparsifier(sparsity_level=0.5)
+        sparsifier.prepare(model, config=None)
+
+        model(X)
+
+        sparsifier.step()
+        sparsifier.squash_mask()
+
+        sparsity = (model[0].weight == 0).float().mean()
+        assert sparsity == 0.5
+
+        expected_fc = torch.tensor([[1, 2, 0, 0]], dtype=torch.float32)
+        assert torch.allclose(model[0].weight.data, expected_fc, rtol=1e-05, atol=1e-07)
+
+    def test_two_layer_mlp_unstructured(self):
+        model = nn.Sequential(
+            nn.Linear(128, 200), nn.ReLU(), nn.Linear(200, 10)
+        )  # C_in by C_out
+        X1 = torch.randn(100, 128)  # B1 by C_in
+        X2 = torch.randn(50, 128)  # B2 by C_in
+
+        sparsifier = WandaSparsifier(sparsity_level=0.5)
+        sparsifier.prepare(model, config=None)
+
+        model(X1)
+        model(X2)
+        sparsifier.step()
+
+        cnt = 0
+        for m in model.modules():
+            if isinstance(m, nn.Linear):
+                cnt += 1
+                sparsity_level = (m.weight == 0).float().mean()
+                assert (
+                    sparsity_level == 0.5
+                ), f"sparsity for linear layer {cnt} should be 0.5"
+
+        sparsifier.squash_mask()
+
+if __name__ == "__main__":
+    unittest.main()

torchao/sparsity/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .wanda import WandaSparsifier  # noqa: F403
+from .utils import PerChannelNormObserver  # noqa: F403
+
+__all__ = [
+    "WandaSparsifier",
+    "PerChannelNormObserver"
+]

torchao/sparsity/utils.py

@@ -0,0 +1,48 @@
+import torch
+from torch.ao.quantization.observer import UniformQuantizationObserverBase
+
+__all__ = ["PerChannelNormObserver"]
+
+# Observers
+class PerChannelNormObserver(UniformQuantizationObserverBase):
+    """
+    A custom observer that computes the L2 norm of each channel and stores it in a buffer.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        # init with fixed qparams for quantization flow
+        super().__init__(
+            dtype=torch.quint8,
+            qscheme=torch.per_channel_affine,
+            reduce_range=False,
+            quant_min=None,
+            quant_max=None,
+            eps=torch.finfo(torch.float32).eps,
+            **kwargs
+        )
+        # set averaging constant so quantization flow knows observer is memoryless.
+        self.averaging_constant = 1.0
+        self.register_buffer("norm", torch.tensor([]))
+
+    def forward(self, x_orig):
+        if x_orig.numel() == 0:
+            return x_orig
+        x = x_orig.detach()  # avoid keeping autograd tape
+
+        # channel_ax is always the last dimension
+        new_axis_list = [i for i in range(x.dim())]  # noqa: C416
+        new_axis_list[0], new_axis_list[-1] = new_axis_list[-1], new_axis_list[0]
+        y = x.permute(new_axis_list)
+        y = torch.flatten(y, start_dim=1)
+        norm = torch.norm(y, dim=1) ** 2
+
+        if self.norm.numel() == 0:
+            self.norm.resize_(norm.shape)
+            self.norm.copy_(norm)
+        else:
+            self.norm += norm
+
+        return x_orig
+
+    def calculate_qparams(self):
+        raise NotImplementedError("PerChannelNormObserver is designed to store activations only. ")

torchao/sparsity/wanda.py

@@ -0,0 +1,110 @@
+
+import warnings
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import nn
+from torch.ao.pruning import BaseSparsifier
+from torch.ao.quantization import default_placeholder_observer, QConfig
+from torch.ao.quantization.quantize import _remove_qconfig
+from .utils import PerChannelNormObserver
+
+__all__ = ["WandaSparsifier"]
+
+
+class WandaSparsifier(BaseSparsifier):
+    r"""Wanda sparsifier
+
+    Wanda (Pruning by Weights and activations), proposed in https://arxiv.org/abs/2306.11695
+    is an activation aware pruning method. The sparsifier removes weights based on the product
+    of the input activation norm and the weight magnitude.
+
+    This sparsifier is controlled by three variables:
+    1. `sparsity_level` defines the number of *sparse blocks* that are zeroed-out;
+
+    Args:
+        sparsity_level: The target level of sparsity;
+        model: The model to be sparsified;
+    """
+
+    def __init__(
+        self,
+        sparsity_level: float = 0.5,
+        semi_structured_block_size: Optional[int] = None,
+    ):
+        defaults = {
+            "sparsity_level": sparsity_level,
+            "semi_structured_block_size": semi_structured_block_size,
+        }
+        if semi_structured_block_size is not None:
+            m = semi_structured_block_size
+            warnings.warn(
+                f"WandaSparsifier got semi_structured_bock_size={m}, sparsity_level fixed to 50% ({m // 2}:{m}) sparsity"
+            )
+        super().__init__(defaults=defaults)
+
+    def prepare(self, model: nn.Module, config: List[Dict]) -> None:
+        # activation: use PerChannelNormObserver
+        # use no-op placeholder weight observer
+        model.qconfig = QConfig(
+            activation=PerChannelNormObserver, weight=default_placeholder_observer
+        )  # type: ignore[assignment]
+        torch.ao.quantization.prepare(model, inplace=True)
+
+        # call superclass prepare
+        super().prepare(model, config)
+
+    def update_mask(  # type: ignore[override]
+        self, module: nn.Module, tensor_name: str, sparsity_level: float, **kwargs
+    ) -> None:
+        r"""Pruning function for WandaSparsifier
+
+        The activation statistics is retrieved first in the `act_per_input` variable.
+        Then the Wanda pruning metric is computed. The weight matrix is then pruned
+        by comparing this metric across the whole current layer.
+        """
+
+        # Step 1: get the tensor and the mask from the parametrizations
+        mask = getattr(module.parametrizations, tensor_name)[0].mask
+        tensor = getattr(module.parametrizations, tensor_name).original
+        activation_norm_per_channel = module.activation_post_process.norm
+
+        # Step 2: Calculate Wx
+        pruning_metric = torch.abs(tensor) * activation_norm_per_channel
+
+        # defaults for unstructured sparsity
+        block_size = pruning_metric.numel()
+        num_specified = int(block_size * sparsity_level)
+        # if set to use semi-structured, ignore sparsity_level
+        if kwargs.get("semi_structured_block_size", None) is not None:
+            block_size = kwargs["semi_structured_block_size"]
+            num_specified = block_size // 2
+
+        # get indicies to prune
+        pruning_inds = pruning_metric.view(-1, block_size).argsort(dim=1)[
+            :, :num_specified
+        ]
+        # update mask
+        mask.data.view(-1, block_size).scatter_(
+            1, pruning_inds, torch.zeros_like(pruning_inds, dtype=mask.dtype)
+        )
+
+    def squash_mask(
+        self,
+        params_to_keep: Optional[Tuple[str, ...]] = None,
+        params_to_keep_per_layer: Optional[Dict[str, Tuple[str, ...]]] = None,
+        *args,
+        **kwargs,
+    ):
+        # remove quantization config
+        for config in self.groups:
+            module = config["module"]
+            tensor_name = config["tensor_name"]
+            _remove_qconfig(module)
+
+        # remove parameterizations
+        super().squash_mask(
+            params_to_keep=params_to_keep,
+            params_to_keep_per_layer=params_to_keep_per_layer,
+        )