Add support for structured sparsity (block structure & n:m structure) - #61
Add support for structured sparsity (block structure & n:m structure)#61u-simha wants to merge 4 commits into
Conversation
Introduces BlockStructured and NMStructured PruningScheme implementations alongside the existing Unstructured and ChannelStructured, and refactors _MagnitudePruneImpl to delegate mask computation to PruningScheme.compute_mask() instead of switching on scheme type. _BlockSizeMismatchError moves to a shared coreai_opt/_utils/errors.py so both quantization and pruning specs can raise it.
| @@ -0,0 +1,10 @@ | |||
| # Copyright 2026 Apple Inc. | |||
There was a problem hiding this comment.
Will you update the pruning tutorial to show structured sparsity in this PR?
There was a problem hiding this comment.
I don't plan to update the tutorial, but I did plan to update the doc page, specifically the config page to highlight the options. Let me add that onto this PR itself.
Removes the pruning and quantization spec errors.py shims that just re-exported _utils.errors._BlockSizeMismatchError, per review feedback on apple#61 — all call sites now import the shared error directly instead of going through a per-domain copy.
State the axis-divisibility constraint directly instead of contrasting with another implementation's padding behavior.
…dd-structured-sparsity # Conflicts: # src/coreai_opt/pruning/spec/prune.py # src/coreai_opt/pruning/spec/scheme.py # src/coreai_opt/quantization/spec/fake_quantize.py # src/coreai_opt/quantization/spec/granularity.py
| return self._compute_mask(weight, sparsity) | ||
|
|
||
| @abstractmethod | ||
| def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor: |
There was a problem hiding this comment.
Given that PruningScheme is a public class meant for users to inherit from, I think we should have any abstractmethods users must override be public as well. Otherwise the private naming suggests it's a method we can choose to freely modify, but would cause any user defined code to break if we do
| 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 |
There was a problem hiding this comment.
Instead of just ignoring target_sparsity completely which could be misleading (it is a public object which user logic may depend on), could we instead have it be derived from doing n/m?
And then two things could come out as a result:
- No need to override the existing
compute_maskmethod (it can follow the early exit logic of sparsity = 0 or 1) - When defining the abstract
_compute_maskmethod, instead of always using strict n/m and ignoring sparsity, it can still use sparsity to determine the number of elements within each block to zero out. This would allow N:M structured to also be trainable via a schedule. We could have a statement likenum_prune = self.n if sparsity >= self.n / self.m else math.floor(sparsity * self.m)to protect against numerical inaccuracies in doingsparsity * self.m, to guarantee that we always end up with at least n elements zeroed out at the end of training.
We could make target_sparsity a float | None in PruningSpec where it is required if the pruning scheme is not NMStructured, and raise an error if the user creates a spec with both target_sparsity provided as well as pruning_scheme = NMStructured.
| 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() |
There was a problem hiding this comment.
Nit, this could be rewritten as torch.linalg.vector_norm(grouped, dim=...)
3 participants
What
This PR adds support for structured sparsity, specifically block structured & n:m structured. Certain hardware can accelerate sparsity further by skipping blocks of 0 when loading the weights into memory.
How
prune.pyand abstracted cleanly intoscheme.pywhere each pruning scheme computes the mask based on the sparsityBlockSizeMismatcherror to be shared across pruning & quantizationTesting