[Good First Issue][NNCF]: Add ONNX support of data-free Weight Compression Algorithm

[kshpv]

Context

NNCF supports OpenVINO, Torch and TorchFX backends for weight compression algorithm - nncf.compress_weights(). The goal of this issue is to expand support to the ONNX backend. The structure of the NNCF code is designed in a way that this could be done quite straightforward, but it requires attention to detail.

Very important: the runtime target is OpenVINOExecutionProvider provider for ONNXRuntime

What needs to be done?

The task is to implement data-free int8 and uint8 Weight Compression algorithm support which includes:

1. Implement WeightCompressionAlgoBackend for ONNX:
   We already have this implemented for OpenVINO, Torch, and TorchFX, so you can use those as references.
   Some methods like insert_adapters, _get_statistics_for_weights_compression, and dump_parameters can be skipped.
   The goal is to make sure we can run nncf.compress_weights(onnx_model) and get a ONNX model with compressed weights in int8, uin8 formats.

2. Test the Compression:
   Ensure that running nncf.compress_weights(onnx_model) actually produces a compressed ONNX model.

3. Add Initial Tests:
   This is super important to prove that the algorithm works correctly.
   There are two types of tests we need:
   Conformance Tests: Add a tinyllama_data_free case for ONNX, similar to what we have for OpenVINO. Note: the test, a good starting point is to read readme
   Unit Tests: We'll need to add some unit tests. Note: this is a unitetsts for OpenVINO, Torch and TorchFX

This can be split into some subtasks to develop/review faster. This is up to you and can be discussed.
If you have any questions or need guidance, feel free to ask in the comments or reach out to the maintainers.

Example Pull Requests

Adding support of data free for Torch - #2333
Adding support of data free for TorchFX - #2891

Resources

• Contribution guide - start here!
• Intel DevHub Discord channel - engage in discussions, ask questions and talk to OpenVINO developers
• How to link your Pull Request to an issue

Contact points

@kshpv

The description is not full and will be updated

[XueSongTap]

Hi @kshpv I'd like to work on this issue. I have experience with ONNX and can help implement the weight compression algorithm support for the ONNX backend.

[alexsu52]

@XueSongTap, thank you for your interest! Assigned to you.

[kshpv]

Very important requirements to keep in mind:

1. The weights compression scheme should be as follows: DequantizeLinear(weight, ...) -> MatMul/Conv/....

2. The block_size attribute is only supported for ONNX opset >= 21. This leads to the next requirement, Weight Compression will only support ONNX model starting from opset 21. For the rest of the models algorithm should not execute with appropriate message.

[kshpv]

I would like to assist with Step 2 - Test the Compression:

I have prepared a set of scripts that can be utilized with the following workflow:

1. Export the ONNX model, such as BERT, using the optimum and transformers libraries.
2. Measure the execution time of the original model over N iterations.
3. Compress the ONNX model using NNCF (Neural Network Compression Framework).
4. Measure the execution time of the compressed model over N iterations.

Let me know when we are good with the step 1

[alexsu52]

@XueSongTap, how long do you estimate it will take you to complete this task?

[XueSongTap]

Hi @alexsu52 based on the requirements and my current understanding, I estimate it will take around 2-3 weeks to complete this task:

Week 1:

• Study existing implementations (OpenVINO, Torch backends)
• Start implementing basic WeightCompressionAlgoBackend for ONNX

Week 2:

• Complete core implementation
• Add initial unit tests
• Handle ONNX opset version checks and requirements

Week 3 (if needed):

• Implement conformance tests
• Performance testing and optimization
• Documentation and code review fixes

I plan to provide regular updates and can adjust the timeline based on feedback. Please let me know if this timeline works or if you have different expectations.

[alexsu52]

@XueSongTap, sounds good! We look forward to updates from you.

If you have any questions, don't hesitate to ask them.

[alexsu52]

Hello @XueSongTap, could you provide your status? When are you going to open PR with core implementation?

[XueSongTap]

hi@alexsu52 I've mostly completed the core feature implementation for NNCF weight-compression with ONNX support. Currently, I'm in the self-testing phase but encountering some environment-related issues that I'm working to resolve.

The specific error I'm facing is:

Traceback (most recent call last):
  File "/home/yxc/code/2502/nncf/tests/test_weigh_compression.py", line 3, in <module>
    from nncf.quantization import compress_weights
  File "/home/yxc/code/2502/nncf/nncf/__init__.py", line 43, in <module>
    from nncf.quantization import QuantizationPreset as QuantizationPreset
  File "/home/yxc/code/2502/nncf/nncf/quantization/__init__.py", line 13, in <module>
    from nncf.quantization.quantize_model import compress_weights as compress_weights
  File "/home/yxc/code/2502/nncf/nncf/quantization/quantize_model.py", line 36, in <module>
    from nncf.quantization.algorithms.hyperparameter_tuner.algorithm import HyperparameterTuner
  File "/home/yxc/code/2502/nncf/nncf/quantization/algorithms/hyperparameter_tuner/algorithm.py", line 19, in <module>
    from nncf.common.factory import NNCFGraphFactory
  File "/home/yxc/code/2502/nncf/nncf/common/factory.py", line 19, in <module>
    from nncf.common.tensor_statistics import aggregator
  File "/home/yxc/code/2502/nncf/nncf/common/tensor_statistics/aggregator.py", line 25, in <module>
    from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
  File "/home/yxc/code/2502/nncf/nncf/common/tensor_statistics/statistic_point.py", line 16, in <module>
    from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
  File "/home/yxc/code/2502/nncf/nncf/experimental/common/tensor_statistics/collectors.py", line 20, in <module>
    import nncf.tensor.functions as fns
  File "/home/yxc/code/2502/nncf/nncf/tensor/functions/__init__.py", line 86, in <module>
    _initialize_backends()
  File "/home/yxc/code/2502/nncf/nncf/tensor/functions/__init__.py", line 78, in _initialize_backends
    import nncf.tensor.functions.torch_io
  File "/home/yxc/code/2502/nncf/nncf/tensor/functions/torch_io.py", line 15, in <module>
    from safetensors.torch import load_file as pt_load_file
  File "/home/yxc/micromamba/envs/nncf/lib/python3.11/site-packages/safetensors/torch.py", line 11, in <module>
    def storage_ptr(tensor: torch.Tensor) -> int:
                            ^^^^^^^^^^^^
AttributeError: module 'torch' has no attribute 'Tensor'

This appears to be a PyTorch installation issue. I'm currently troubleshooting this problem. Once I resolve these environment issues and complete self-testing, I'll submit the PR with the core implementation. I estimate this will another week to complete.

[alexsu52]

Hi @XueSongTap, thanks for the very valuable update. Could you open the draft PR with current implementation? I understand that your implementation is not ready yet, but in any case, it will be easier for us to support you in solving problems.

Could you please provide the command you use to run the tests and how you set up PYTHONPATH?