diff --git a/src/coreai_opt/quantization/spec/errors.py b/src/coreai_opt/_utils/errors.py similarity index 79% rename from src/coreai_opt/quantization/spec/errors.py rename to src/coreai_opt/_utils/errors.py index 2d537e3..23af82c 100644 --- a/src/coreai_opt/quantization/spec/errors.py +++ b/src/coreai_opt/_utils/errors.py @@ -3,6 +3,8 @@ # Use of this source code is governed by a BSD-3-Clause license that can # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause +"""Errors shared across compression domains (quantization, pruning, ...).""" + class _BlockSizeMismatchError(ValueError): """Raised when a tensor dimension is not divisible by the block size.""" diff --git a/src/coreai_opt/pruning/spec/__init__.py b/src/coreai_opt/pruning/spec/__init__.py index 354345f..9e9df32 100644 --- a/src/coreai_opt/pruning/spec/__init__.py +++ b/src/coreai_opt/pruning/spec/__init__.py @@ -6,11 +6,13 @@ """Pruning spec components: specs, schemes, and parametrizations.""" from .prune import PruneImplBase, _MagnitudePruneImpl -from .scheme import ChannelStructured, PruningScheme, Unstructured +from .scheme import BlockStructured, ChannelStructured, NMStructured, PruningScheme, Unstructured from .spec import PruningSpec, default_weight_pruning_spec __all__ = [ + "BlockStructured", "ChannelStructured", + "NMStructured", "PruneImplBase", "PruningScheme", "PruningSpec", diff --git a/src/coreai_opt/pruning/spec/prune.py b/src/coreai_opt/pruning/spec/prune.py index 7333fda..1918fce 100644 --- a/src/coreai_opt/pruning/spec/prune.py +++ b/src/coreai_opt/pruning/spec/prune.py @@ -7,7 +7,6 @@ from __future__ import annotations -import math from abc import abstractmethod from typing import TYPE_CHECKING, Any @@ -17,10 +16,9 @@ PartialConstructor as _PartialConstructor, with_args as _with_args, ) -from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis from coreai_opt.config.spec import CompressionSimulatorBase -from .scheme import ChannelStructured, PruningScheme +from .scheme import PruningScheme if TYPE_CHECKING: # Imported only for type checking — runtime would be a circular import via @@ -117,10 +115,11 @@ def with_args(cls, **kwargs: Any) -> _PartialConstructor[PruneImplBase]: @PruneImplBase.register("default") class _MagnitudePruneImpl(PruneImplBase): - """Magnitude-based pruning supporting unstructured and channel-structured schemes. + """Magnitude-based pruning that delegates to the configured pruning scheme. Prunes a given tensor to target sparsity by zero-ing out the smallest-magnitude - elements until desired sparsity is achieved. + elements, per whatever structural pattern ``pruning_scheme`` defines (unstructured, + channel-structured, block-structured, N:M-structured, or future schemes). """ @staticmethod @@ -129,7 +128,7 @@ def compute_mask( sparsity: float, pruning_scheme: PruningScheme, ) -> torch.Tensor: - """Compute a magnitude-based mask respecting the pruning scheme. + """Compute a magnitude-based mask by delegating to the pruning scheme. Args: weight (torch.Tensor): The weight tensor. @@ -139,60 +138,4 @@ def compute_mask( Returns: torch.Tensor: Binary mask (1 = keep, 0 = prune). """ - if sparsity == 0.0: - return torch.ones_like(weight) - if sparsity >= 1.0: - return torch.zeros_like(weight) - - # TODO: Replace this with generic abstractions - if isinstance(pruning_scheme, ChannelStructured): - return _MagnitudePruneImpl._compute_channel_mask(weight, sparsity, pruning_scheme.axis) - return _MagnitudePruneImpl._compute_unstructured_mask(weight, sparsity) - - @staticmethod - def _compute_unstructured_mask(weight: torch.Tensor, sparsity: float) -> torch.Tensor: - """Element-wise magnitude pruning.""" - num_elements = weight.numel() - num_keep = num_elements - math.floor(num_elements * sparsity) - abs_weight = weight.abs() - _, topk_indices = torch.topk(abs_weight.flatten(), num_keep) - mask = torch.zeros(num_elements, dtype=weight.dtype, device=weight.device) - mask[topk_indices] = 1.0 - return mask.reshape(weight.shape) - - @staticmethod - def _compute_channel_mask( - weight: torch.Tensor, - sparsity: float, - axis: int, - ) -> torch.Tensor: - """Channel-structured magnitude pruning along *axis*. - - Channel importance is measured by L1 norm. The least-important - channels are pruned entirely. - """ - if not (-weight.ndim <= axis < weight.ndim): - raise ValueError( - f"Invalid axis. Should be in range [{-weight.ndim}, {weight.ndim}), but got {axis}" - ) - axis = _normalize_axis(axis, weight.ndim) - - num_channels = weight.shape[axis] - num_prune = math.floor(num_channels * sparsity) - - if num_prune == 0: - return torch.ones_like(weight) - if num_prune >= num_channels: - return torch.zeros_like(weight) - - reduce_dims = [d for d in range(weight.ndim) if d != axis] - channel_norms = weight.abs().sum(dim=reduce_dims) - - num_keep = num_channels - num_prune - _, keep_indices = torch.topk(channel_norms, num_keep, largest=True) - channel_mask = torch.zeros(num_channels, dtype=weight.dtype, device=weight.device) - channel_mask[keep_indices] = 1.0 - - shape = [1] * weight.ndim - shape[axis] = num_channels - return channel_mask.view(shape).expand_as(weight) + return pruning_scheme.compute_mask(weight, sparsity) diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py index 73e1d6b..7df0f38 100644 --- a/src/coreai_opt/pruning/spec/scheme.py +++ b/src/coreai_opt/pruning/spec/scheme.py @@ -7,18 +7,32 @@ from __future__ import annotations +import math +from abc import abstractmethod from typing import Any, Literal -from pydantic import BaseModel, ConfigDict, Field, model_serializer +import torch +from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator +from coreai_opt._utils.errors import _BlockSizeMismatchError from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin +from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis class PruningScheme(BaseModel, _ConfigRegistryMixin): """Base class for pruning scheme specifications. A pruning scheme defines the structural pattern of sparsity applied - to a tensor. Subclasses represent different ways of structuring the pruning. + to a tensor, and knows how to turn ``(weight, sparsity)`` into a binary + mask. Call the public :meth:`compute_mask` to get a mask; subclasses + implement the abstract :meth:`_compute_mask`, which handles sparsity + strictly between 0 and 1 (the 0.0 / 1.0 edge cases are handled once, in + the base class). + + The sole exception is :class:`NMStructured`, whose achieved sparsity is + fixed by construction (``n / m``) rather than a free parameter — it + overrides :meth:`compute_mask` directly and ignores the ``sparsity`` + argument entirely. Attributes: axis (int | None): The axis along which structured pruning is applied. @@ -51,6 +65,32 @@ def _serialize_model(self) -> dict[str, Any]: return data + def compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + """Compute a binary pruning mask for the given weight tensor. + + Args: + weight (torch.Tensor): The weight tensor to compute a mask for. + sparsity (float): Fraction of elements to prune, in [0, 1]. + + Returns: + torch.Tensor: Binary mask with the same shape as *weight* (1 = keep, + 0 = prune). + """ + if sparsity == 0.0: + return torch.ones_like(weight) + if sparsity >= 1.0: + return torch.zeros_like(weight) + return self._compute_mask(weight, sparsity) + + @abstractmethod + def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + """Compute a mask for *sparsity* strictly between 0 and 1. + + Subclasses implement the scheme-specific masking logic here; the + 0.0 / 1.0 edge cases are already handled by :meth:`compute_mask`. + """ + ... + @PruningScheme.register("unstructured") class Unstructured(PruningScheme): @@ -62,17 +102,170 @@ class Unstructured(PruningScheme): axis: Literal[None] = None + def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + num_elements = weight.numel() + num_keep = num_elements - math.floor(num_elements * sparsity) + abs_weight = weight.abs() + _, topk_indices = torch.topk(abs_weight.flatten(), num_keep) + mask = torch.zeros(num_elements, dtype=weight.dtype, device=weight.device) + mask[topk_indices] = 1.0 + return mask.reshape(weight.shape) + @PruningScheme.register("channel_structured") class ChannelStructured(PruningScheme): """Channel-structured pruning scheme. Entire channels (slices along ``axis``) are pruned or kept together. - Channel importance is determined by the pruning algorithm (e.g. L1 norm - of each channel for magnitude-based pruning). + Channel importance is determined by L1 norm of each channel. Note: ``axis`` can be negatively indexed as per standard Python style indexing. """ axis: int = Field(default=0, description="Axis along which channels are pruned.") + + def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + if not (-weight.ndim <= self.axis < weight.ndim): + raise ValueError( + f"Invalid axis. Should be in range [{-weight.ndim}, {weight.ndim}), " + f"but got {self.axis}" + ) + axis = _normalize_axis(self.axis, weight.ndim) + + num_channels = weight.shape[axis] + num_prune = math.floor(num_channels * sparsity) + + if num_prune == 0: + return torch.ones_like(weight) + if num_prune >= num_channels: + return torch.zeros_like(weight) + + reduce_dims = [d for d in range(weight.ndim) if d != axis] + channel_norms = weight.abs().sum(dim=reduce_dims) + + num_keep = num_channels - num_prune + _, keep_indices = torch.topk(channel_norms, num_keep, largest=True) + channel_mask = torch.zeros(num_channels, dtype=weight.dtype, device=weight.device) + channel_mask[keep_indices] = 1.0 + + shape = [1] * weight.ndim + shape[axis] = num_channels + return channel_mask.view(shape).expand_as(weight) + + +@PruningScheme.register("block_structured") +class BlockStructured(PruningScheme): + """Block-structured pruning scheme. + + Generalizes :class:`ChannelStructured` to prune contiguous blocks of + ``block_size`` slices along ``axis`` together, ranked by L2 norm + (``ChannelStructured`` is equivalent to ``block_size=1``, ranked by L1 + norm — the two remain separate registered schemes). + + The size of ``weight`` along ``axis`` must be evenly divisible by + ``block_size``. + """ + + axis: int = Field(default=0, description="Axis along which blocks are formed.") + block_size: int = Field(gt=0, description="Number of contiguous slices per block along axis.") + + def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + num_along_axis = weight.shape[self.axis] + if num_along_axis % self.block_size != 0: + raise _BlockSizeMismatchError( + f"Tensor size {num_along_axis} along axis {self.axis} is not " + f"divisible by block_size {self.block_size}. Full tensor shape: " + f"{tuple(weight.shape)}" + ) + + num_blocks = num_along_axis // self.block_size + num_prune = math.floor(num_blocks * sparsity) + + if num_prune == 0: + return torch.ones_like(weight) + if num_prune >= num_blocks: + return torch.zeros_like(weight) + + # axis is now at position 0; all other dims keep their relative order. + moved = torch.movedim(weight, self.axis, 0) + other_dims = moved.shape[1:] + grouped = moved.view(num_blocks, self.block_size, *other_dims) + block_norms = grouped.pow(2).sum(dim=tuple(range(1, grouped.ndim))).sqrt() + + num_keep = num_blocks - num_prune + _, keep_indices = torch.topk(block_norms, num_keep, largest=True) + block_mask = torch.zeros(num_blocks, dtype=weight.dtype, device=weight.device) + block_mask[keep_indices] = 1.0 + + mask = block_mask.repeat_interleave(self.block_size) + mask = mask.view(num_along_axis, *([1] * len(other_dims))).expand( + num_along_axis, *other_dims + ) + # Restore axis to its original position. + return torch.movedim(mask, 0, self.axis) + + +@PruningScheme.register("n_m_structured") +class NMStructured(PruningScheme): + """N:M structured pruning scheme. + + Zeroes exactly ``n`` smallest-magnitude elements out of every contiguous + group of ``m`` elements along ``axis`` — a hardware-friendly sparsity + pattern with a fixed sparsity ratio of ``n / m``. + + Unlike other schemes, the achieved sparsity is fixed by construction and + does not depend on ``PruningSpec.target_sparsity``: :meth:`compute_mask` + overrides the base class directly and **ignores** its ``sparsity`` + argument. + + The size of ``weight`` along ``axis`` must be evenly divisible by ``m``. + """ + + axis: int = Field(default=0, description="Axis along which N:M groups are formed.") + n: int = Field(ge=0, description="Number of smallest-magnitude elements zeroed per group of m.") + m: int = Field(gt=0, description="Group size along axis.") + + @model_validator(mode="after") + def _validate_n_lt_m(self) -> NMStructured: + if self.n >= self.m: + raise ValueError(f"n ({self.n}) must be less than m ({self.m})") + return self + + def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + """Unreachable: :meth:`compute_mask` is overridden directly below and never + delegates here. Defined only to satisfy the base class's abstract method. + """ + raise NotImplementedError( + "NMStructured overrides compute_mask directly; _compute_mask is unused." + ) + + def compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: + """Compute the N:M mask. + + ``sparsity`` is accepted for interface compatibility with + :class:`PruneImplBase` but is ignored — achieved sparsity is fixed + at ``n / m`` by construction. + """ + num_along_axis = weight.shape[self.axis] + if num_along_axis % self.m != 0: + raise _BlockSizeMismatchError( + f"Tensor size {num_along_axis} along axis {self.axis} is not " + f"divisible by m {self.m}. Full tensor shape: {tuple(weight.shape)}" + ) + + if self.n == 0: + return torch.ones_like(weight) + + # axis is now at position -1; all other dims keep their relative order. + moved = torch.movedim(weight, self.axis, -1) + original_shape = moved.shape + grouped = moved.reshape(-1, self.m) + + prune_idx = torch.argsort(grouped.abs(), dim=1, stable=True)[:, : self.n] + mask = torch.ones_like(grouped) + mask.scatter_(1, prune_idx, 0.0) + + mask = mask.reshape(original_shape) + # Restore axis to its original position. + return torch.movedim(mask, -1, self.axis) diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index 1c0bfaa..fbf93f7 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/src/coreai_opt/quantization/spec/fake_quantize.py @@ -15,6 +15,7 @@ from torch.autograd import Function from torchao.quantization.pt2e import FakeQuantizeBase +from coreai_opt._utils.errors import _BlockSizeMismatchError from coreai_opt._utils.spec_utils import ( PartialConstructor as _PartialConstructor, with_args as _with_args, @@ -27,7 +28,6 @@ ) from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes -from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError from coreai_opt.quantization.spec.qscheme import QuantizationScheme from .granularity import QuantizationGranularity diff --git a/src/coreai_opt/quantization/spec/granularity.py b/src/coreai_opt/quantization/spec/granularity.py index 859bac3..d926af8 100644 --- a/src/coreai_opt/quantization/spec/granularity.py +++ b/src/coreai_opt/quantization/spec/granularity.py @@ -11,9 +11,9 @@ import torch from pydantic import BaseModel, ConfigDict, Field, model_serializer +from coreai_opt._utils.errors import _BlockSizeMismatchError from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis -from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError class QuantizationGranularity(BaseModel, _ConfigRegistryMixin): diff --git a/tests/fixtures/pruning.py b/tests/fixtures/pruning.py index f4392dc..454a856 100644 --- a/tests/fixtures/pruning.py +++ b/tests/fixtures/pruning.py @@ -11,7 +11,13 @@ from coreai_opt import ExportBackend from coreai_opt.pruning import MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec -from coreai_opt.pruning.spec import ChannelStructured, PruningScheme, Unstructured +from coreai_opt.pruning.spec import ( + BlockStructured, + ChannelStructured, + NMStructured, + PruningScheme, + Unstructured, +) @dataclass @@ -21,7 +27,8 @@ class ParametrizedPruneConfigs: Attributes: config: MagnitudePrunerConfig instance. target_sparsity: Target sparsity fraction. - pruning_scheme: PruningScheme instance (Unstructured or ChannelStructured). + pruning_scheme: PruningScheme instance (Unstructured, ChannelStructured, + BlockStructured, or NMStructured). backend: Export backend (CoreML or CoreAI). """ @@ -53,7 +60,12 @@ def from_prune_params( params=[ (target_sparsity, pruning_scheme, backend) for target_sparsity in [0.25, 0.5, 0.75] - for pruning_scheme in [Unstructured(), ChannelStructured(axis=0)] + for pruning_scheme in [ + Unstructured(), + ChannelStructured(axis=0), + BlockStructured(axis=0, block_size=2), + NMStructured(axis=0, n=1, m=2), + ] for backend in [ExportBackend.CoreML, ExportBackend.CoreAI] ], ids=lambda p: f"sparsity:{p[0]}-scheme:{p[1].__class__.__name__}-backend:{p[2].value}", diff --git a/tests/pruning/test_magnitude_pruner.py b/tests/pruning/test_magnitude_pruner.py index b3b60e3..d20e3be 100644 --- a/tests/pruning/test_magnitude_pruner.py +++ b/tests/pruning/test_magnitude_pruner.py @@ -23,7 +23,13 @@ OpMagnitudePrunerConfig, PolynomialDecaySchedule, ) -from coreai_opt.pruning.spec import ChannelStructured, PruneImplBase, Unstructured +from coreai_opt.pruning.spec import ( + BlockStructured, + ChannelStructured, + NMStructured, + PruneImplBase, + Unstructured, +) @pytest.fixture @@ -553,6 +559,181 @@ def test_channel_structured_duplicate_norms(self) -> None: f"Expected exactly 3/5 channels pruned despite tied norms, got {num_pruned}" ) + def test_block_structured_pruning_hand(self) -> None: + """Hand-written 8x4 tensor with block pruning (axis=0, block_size=2) at 50%. + + Block magnitudes [5, 1, 4, 2] are interleaved so the two pruned blocks + (norms 1 and 2) are not contiguous — this keeps the case distinct from + plain channel-structured pruning of a contiguous leading/trailing slice. + """ + model = nn.Linear(4, 8, bias=False) + with torch.no_grad(): + model.weight.copy_( + torch.tensor( + [ + [5.0, 5.0, 5.0, 5.0], + [5.0, 5.0, 5.0, 5.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [4.0, 4.0, 4.0, 4.0], + [4.0, 4.0, 4.0, 4.0], + [2.0, 2.0, 2.0, 2.0], + [2.0, 2.0, 2.0, 2.0], + ] + ) + ) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.5, + pruning_scheme=BlockStructured(axis=0, block_size=2), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 4),)) + + expected = torch.tensor( + [ + [5.0, 5.0, 5.0, 5.0], + [5.0, 5.0, 5.0, 5.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [4.0, 4.0, 4.0, 4.0], + [4.0, 4.0, 4.0, 4.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ] + ) + assert torch.equal(model.weight.detach(), expected) + + def test_block_structured_non_leading_axis(self) -> None: + """Block pruning along axis=1 (input-feature dim) groups column pairs.""" + model = nn.Linear(6, 4, bias=False) + with torch.no_grad(): + model.weight.copy_(torch.tensor([[1.0, 1.0, 5.0, 5.0, 3.0, 3.0]]).expand(4, 6)) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.4, + pruning_scheme=BlockStructured(axis=1, block_size=2), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 6),)) + + expected = torch.tensor([[0.0, 0.0, 5.0, 5.0, 3.0, 3.0]]).expand(4, 6) + assert torch.equal(model.weight.detach(), expected) + + def test_block_structured_duplicate_norms(self) -> None: + """Blocks with tied L2 norms are still pruned to the exact target count. + + Block magnitudes [1, 2, 2, 2, 5] (5 blocks of 2 rows each) at 60% + sparsity should prune exactly 3 of 5 blocks, despite the norm ties. + """ + magnitudes = [1.0, 2.0, 2.0, 2.0, 5.0] + rows = [torch.full((4,), m) for m in magnitudes for _ in range(2)] + model = nn.Linear(4, 10, bias=False) + with torch.no_grad(): + model.weight.copy_(torch.stack(rows)) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.6, + pruning_scheme=BlockStructured(axis=0, block_size=2), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 4),)) + + weight = model.weight.detach() + num_pruned_blocks = sum( + 1 + for block_start in range(0, 10, 2) + if weight[block_start : block_start + 2].eq(0).all() + ) + assert num_pruned_blocks == 3 + + def test_nm_structured_pruning_hand(self) -> None: + """NMStructured(n=2, m=4) enforces exact 2:4 sparsity per group, ignoring target_sparsity. + + Each row is one group of 4 (axis=1): with strictly increasing values, + the 2 smallest elements per row are the exact ones zeroed. + """ + model = nn.Linear(4, 8, bias=False) + original = torch.arange(1, 33, dtype=torch.float32).reshape(8, 4) + with torch.no_grad(): + model.weight.copy_(original) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.9, + pruning_scheme=NMStructured(axis=1, n=2, m=4), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 4),)) + + expected = original.clone() + expected[:, :2] = 0.0 + assert torch.equal(model.weight.detach(), expected) + + def test_nm_structured_multi_group_per_axis_line(self) -> None: + """An 8-element axis with m=4, n=2 forms two independent groups per row.""" + model = nn.Linear(8, 1, bias=False) + with torch.no_grad(): + model.weight.copy_(torch.tensor([[4.0, 1.0, 3.0, 2.0, 8.0, 5.0, 7.0, 6.0]])) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={"weight": PruningSpec(pruning_scheme=NMStructured(axis=1, n=2, m=4))} + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 8),)) + + # Group 1 = [4, 1, 3, 2]: two smallest (1, 2) zeroed -> keep 4, 3. + # Group 2 = [8, 5, 7, 6]: two smallest (5, 6) zeroed -> keep 8, 7. + expected = torch.tensor([[4.0, 0.0, 3.0, 0.0, 8.0, 0.0, 7.0, 0.0]]) + assert torch.equal(model.weight.detach(), expected) + + def test_nm_structured_n_zero_keeps_everything(self) -> None: + """NMStructured(n=0) is a degenerate no-op: the weight is left untouched.""" + model = nn.Linear(8, 2, bias=False) + original = torch.arange(1, 17, dtype=torch.float32).reshape(2, 8) + with torch.no_grad(): + model.weight.copy_(original) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.9, + pruning_scheme=NMStructured(axis=1, n=0, m=4), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 8),)) + + assert torch.equal(model.weight.detach(), original) + def test_backprop_through_pruned_weights(self, linear_100x100_unique: nn.Linear) -> None: """Gradients flow through unpruned entries and are zero where the mask is zero.""" pruner = MagnitudePruner( diff --git a/tests/pruning/test_pruning_config_and_spec.py b/tests/pruning/test_pruning_config_and_spec.py index 0eee33a..285e8b7 100644 --- a/tests/pruning/test_pruning_config_and_spec.py +++ b/tests/pruning/test_pruning_config_and_spec.py @@ -9,13 +9,16 @@ import torch import torch.nn as nn +from coreai_opt._utils.errors import _BlockSizeMismatchError from coreai_opt.pruning.config import ( MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, OpMagnitudePrunerConfig, ) from coreai_opt.pruning.spec import ( + BlockStructured, ChannelStructured, + NMStructured, PruneImplBase, PruningScheme, PruningSpec, @@ -42,8 +45,17 @@ def test_default_spec(self) -> None: (0.75, Unstructured()), (0.5, ChannelStructured(axis=0)), (0.9, ChannelStructured(axis=1)), + (0.5, BlockStructured(axis=0, block_size=2)), + (0.9, NMStructured(axis=0, n=1, m=4)), + ], + ids=[ + "25%-unstructured", + "75%-unstructured", + "50%-channel-ax0", + "90%-channel-ax1", + "50%-block-ax0", + "90%-nm-ax0", ], - ids=["25%-unstructured", "75%-unstructured", "50%-channel-ax0", "90%-channel-ax1"], ) def test_custom_spec(self, target_sparsity: float, pruning_scheme: PruningScheme) -> None: """Custom spec values are accepted and stored correctly.""" @@ -93,6 +105,74 @@ def test_pruning_scheme_round_trip( assert isinstance(spec.pruning_scheme, expected_type) assert spec.pruning_scheme.axis == expected_axis + @pytest.mark.parametrize( + "scheme_dict,expected_axis,expected_block_size", + [ + ({"type": "block_structured", "axis": 0, "block_size": 2}, 0, 2), + ({"type": "block_structured", "axis": 1, "block_size": 4}, 1, 4), + ], + ids=["block-ax0-size2", "block-ax1-size4"], + ) + def test_block_structured_round_trip( + self, scheme_dict: dict, expected_axis: int, expected_block_size: int + ) -> None: + """BlockStructured constructed from dict resolves correctly.""" + spec = PruningSpec(pruning_scheme=scheme_dict) + assert isinstance(spec.pruning_scheme, BlockStructured) + assert spec.pruning_scheme.axis == expected_axis + assert spec.pruning_scheme.block_size == expected_block_size + + @pytest.mark.parametrize( + "scheme_dict,expected_axis,expected_n,expected_m", + [ + ({"type": "n_m_structured", "axis": 0, "n": 1, "m": 4}, 0, 1, 4), + ({"type": "n_m_structured", "axis": 1, "n": 2, "m": 8}, 1, 2, 8), + ], + ids=["nm-ax0-1of4", "nm-ax1-2of8"], + ) + def test_nm_structured_round_trip( + self, scheme_dict: dict, expected_axis: int, expected_n: int, expected_m: int + ) -> None: + """NMStructured constructed from dict resolves correctly.""" + spec = PruningSpec(pruning_scheme=scheme_dict) + assert isinstance(spec.pruning_scheme, NMStructured) + assert spec.pruning_scheme.axis == expected_axis + assert spec.pruning_scheme.n == expected_n + assert spec.pruning_scheme.m == expected_m + + def test_block_structured_invalid_block_size_raises(self) -> None: + """block_size <= 0 raises ValueError at construction time.""" + with pytest.raises(ValueError): + BlockStructured(axis=0, block_size=0) + + def test_block_structured_not_divisible_raises_via_spec(self) -> None: + """A tensor size not divisible by block_size raises _BlockSizeMismatchError.""" + scheme = BlockStructured(axis=0, block_size=2) + with pytest.raises(_BlockSizeMismatchError): + scheme.compute_mask(torch.randn(5, 4), sparsity=0.5) + + def test_nm_structured_not_divisible_raises_via_spec(self) -> None: + """A tensor size not divisible by m raises _BlockSizeMismatchError.""" + scheme = NMStructured(axis=0, n=1, m=4) + with pytest.raises(_BlockSizeMismatchError): + scheme.compute_mask(torch.randn(5, 4), sparsity=0.0) + + @pytest.mark.parametrize( + "n,m", + [(4, 4), (5, 4)], + ids=["n-equal-m", "n-greater-than-m"], + ) + def test_nm_structured_invalid_n_m_raises_via_spec(self, n: int, m: int) -> None: + """n >= m raises ValueError, including when constructed via dict through the spec.""" + with pytest.raises(ValueError, match="must be less than"): + PruningSpec(pruning_scheme={"type": "n_m_structured", "n": n, "m": m}) + + def test_nm_structured_n_zero_valid_construction(self) -> None: + """n=0 is a valid degenerate configuration regardless of m.""" + scheme = NMStructured(n=0, m=4) + assert scheme.n == 0 + assert scheme.m == 4 + def test_pruning_algo_round_trip(self) -> None: """Default pruning_algo string resolves to the class object.""" spec = PruningSpec(pruning_algo="default") diff --git a/tests/quantization/test_quantization_spec.py b/tests/quantization/test_quantization_spec.py index 8dc9efe..347b1d5 100644 --- a/tests/quantization/test_quantization_spec.py +++ b/tests/quantization/test_quantization_spec.py @@ -9,6 +9,7 @@ from pydantic import ValidationError from torchao.quantization import MappingType as TorchAOMappingType +from coreai_opt._utils.errors import _BlockSizeMismatchError from coreai_opt.quantization import QuantizationSpec from coreai_opt.quantization.spec import ( PerBlockGranularity, @@ -19,7 +20,6 @@ default_activation_quantization_spec, default_weight_quantization_spec, ) -from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError from coreai_opt.quantization.spec.fake_quantize import ( _DefaultFakeQuantizeImpl, )