openvinotoolkit / openvinotoolkit/nncf

[Feature request] Structured M:N (2:4) magnitude-based pruning

Open
#4,222 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.2k
Forks
305
Avg merge
1d 9h
Merged PRs (30d)
27

Description

Is your feature request related to a problem? Please describe.

NNCF currently supports pruning via nncf.prune with the modes UNSTRUCTURED_MAGNITUDE_LOCAL, UNSTRUCTURED_MAGNITUDE_GLOBAL and UNSTRUCTURED_REGULARIZATION_BASED (see Pruning usage docs). All of these produce unstructured sparsity: zeros are scattered arbitrarily through the weight tensors.

There is currently no support for structured M:N sparsity (for example 2:4, two zeros in every consecutive group of four weights). As discussed in #3810, where this direction was raised as an opportunity for contribution, M:N patterns are widely supported by hardware and runtimes as a sparsity format, and they can be produced with a magnitude criterion using the same kind of workflow NNCF already has.

Describe the solution you'd like

Add structured M:N magnitude-based pruning to the existing nncf.prune flow, starting with 2:4:

  • For every applicable group of N consecutive weights along a chosen structural dimension, zero out the M weights with the smallest magnitudes and keep the remaining N-M weights. The invariant is enforced independently for every group, rather than only globally.
  • For 2:4 this means each group of four weights contains exactly two zeros and two retained values, selected by magnitude.

Design points I would like feedback on before/while implementing:

  1. API representation. Following the existing PruneMode enum convention, something like a new mode STRUCTURED_MAGNITUDE_2_4 (similar to how CompressWeightsMode enumerates format variants such as MXFP4, NVFP4, ...), with an optional extensible sparsity_pattern=(m, n) parameter kept internal for now. Alternatively a fully generic (m, n) public parameter from the start.
  2. Supported layer types. Initially nn.Linear (incl. MatMul/Addmm weights) and nn.Conv2d/Conv1d/Conv3d weights. Depthwise convolutions and embeddings would be excluded, and unsupported metatypes would be skipped with a warning so the rest of the model is still pruned.
  3. Grouping dimension. For linear/matmul weights, group along the output-channel (row) dimension, with in-features contiguous. For convolutions, reuse the existing channel-axes helpers so groups follow the output-channel × spatial position layout. The grouping should be built from the same channel/reduction axis information the weight-compression code already uses (get_weight_compression_reduction_axes).
  4. Non-divisible dimensions. The cleanest policy seems to be raising nncf.ValidationError with a clear message when the structural dimension is not divisible by N, consistent with the fail-fast validation style used elsewhere (e.g. weight compression). Happy to follow a different precedent if one exists.
  5. Interaction with existing features. ignored_scope filtering should keep working unchanged; the new mode should produce statistics via nncf.pruning_statistic; nncf.strip(..., strip_format=StripFormat.IN_PLACE) should bake the masks into the weights; checkpoint save/load via get_config/load_from_config should work through the existing StatefulModuleInterface mechanism.
  6. Testing. Group-level invariant tests (not just overall zero ratio), magnitude selection with positive/negative and tied magnitudes, non-divisible dimensions, ignored scopes, save/load, strip, statistics, and integration tests on Linear/Conv models.
Describe alternatives you've considered
  • A standalone external script/tool outside of NNCF: does not help NNCF users and cannot leverage the existing graph analysis, ignored scopes and statistics.
  • Supporting only a global sparsity ratio and calling it 2:4: this would not guarantee the pattern per group and would not be consumable by hardware expecting the format.
Additional context

This feature follows up on the structured sparsity direction raised in #3810. I searched the current issues and PRs and did not find an equivalent feature request or work in progress for M:N / 2:4 sparsity; if one exists, please point me to it and I will close this.

The goal is to reuse the existing pruning infrastructure (function-hook masks, get_prunable_parameters, pruning statistics, strip) rather than introducing a parallel pruning framework. No hardware acceleration or new kernels are proposed — only producing the weight pattern itself, which is a prerequisite for acceleration on runtimes/hardware that support the format. I'm happy to open a PR implementing this once the general direction is confirmed.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the existing nncf.prune flow, PruneMode conventions, get_prunable_parameters, and get_weight_compression_reduction_axes; the pruning usage docs are at docs/usage/training_time_compression/pruning/Usage.md. Review how pruning statistics, ignored_scope, strip with StripFormat.IN_PLACE, and get_config/load_from_config work. Done means 2:4 group invariants and magnitude selection work for the specified Linear and convolution cases, including the listed validation, persistence, stripping, statistics, and integration coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.