Added support of 2-3 bit symmetric mode to LoRA-QAT - #4158
Open
andreyanufr wants to merge 18 commits into
Open
andreyanufr wants to merge 18 commits into
andreyanufr wants to merge 18 commits into
Conversation
2) Removed 6bit from repacking.
2) Added optional possibility to save torch checkpoint with dequantized weights.
2) Added README and requirements.txt
2) Removed debug prints.
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved critical and moderate issues remain in low-bit stripping, OpenVINO repacking, and the example workflow.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (6)
Identify zero-point constants by graph role, not shape · New Add 2/3-bit DQ packing support or reject validation · New Include gate_proj in producer search · New Handle meta-device offloaded weights during equalization · New Use logger instead of unconditional print calls · New Strengthen repacking test with graph/output validation · New
What changed in this PR
Adds INT2/INT3 symmetric LoRA-QAT support and OpenVINO weight repacking.
Changes:
- Adds 2/3-bit quantization mappings and strip validation.
- Implements and exports
repack_weights. - Adds tests, models, and a distillation QAT example with documentation.
| File | Summary and review notes |
|---|---|
tests/openvino/native/test_model_repacking.py |
Adds repacking coverage; assertions should validate graph correctness and numerical equivalence (nit, 2 votes). |
tests/openvino/native/models.py |
Adds a repacking test model. |
src/nncf/torch/function_hook/strip.py |
2/3-bit DQ still uses 4-bit decompression and packing (critical, 3 votes). |
src/nncf/quantization/quantize_model.py |
Exposes the repacking API. |
src/nncf/quantization/algorithms/weight_compression/torch_backend.py |
Adds INT2/INT3 LoRA quantization mappings. |
src/nncf/quantization/algorithms/weight_compression/openvino_backend.py |
Shape heuristic skips valid signed compressed weights, including 1-D tensors (critical, 2 votes). |
src/nncf/quantization/algorithms/weight_compression/algorithm.py |
Coordinates repacking; direct print calls should use the project logger (nit, 4 votes). |
src/nncf/quantization/__init__.py |
Exports repack_weights. |
src/nncf/openvino/quantization/quantize_model.py |
Provides the OpenVINO repacking implementation. |
src/nncf/__init__.py |
Exports the public API. |
examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt |
Defines example dependencies. |
examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md |
Documents the example workflow. |
examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py |
Omits gate_proj equalization (moderate, 2 votes), skips meta-device weights (moderate, 3 votes), drops partial accumulation gradients (moderate, 1 vote), and does not return its documented metric (nit, 1 vote). |
.ci/cspell_dict.txt |
Adds example vocabulary. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+136
to
+138
| if len(shape) == 1 or shape[-1] == 1: | ||
| # Probably zero point | ||
| return False, 0, 0 |
Comment on lines
+495
to
+497
| if w.device.type != "meta": | ||
| return w.data | ||
| return None |
Comment on lines
+34
to
+38
| repacked_model = nncf.repack_weights(model) | ||
|
|
||
| element_types = get_constant_element_types(repacked_model) | ||
| assert ov.Type.u3 in element_types, "Expected u3 constant after repacking" | ||
| assert ov.Type.u2 in element_types, "Expected u2 constant after repacking" |
2) Added arg for deterministic cuda execution. This arg is optional because it can slow down the training.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.



Changes
LoRA-QAT now support two new modes: CompressWeightsMode.INT3_SYM and CompressWeightsMode.INT2_SYM.
New function repack_weights was added to repack 4-8 bit weights to lower bit representation if possible.
Reason for changes
Support of 2-3 bit in compression in QAT and pytorch backend.
Related tickets
CVS-180191
Tests
In progress.
Weight compression - success