From 32e2e5dff00f1e4207eead4d8ec6d9ca5c524bdf Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Mon, 27 Jul 2026 17:39:52 +0200 Subject: [PATCH 01/27] Added rapcking of weights to support torch compression for 2-3 bit in nncf side. --- .../torch/distillation_qat_with_lora/main.py | 6 +- .../repack_weights.py | 8 +++ .../openvino/quantization/quantize_model.py | 10 +++ .../weight_compression/algorithm.py | 28 ++++++++ .../weight_compression/constants.py | 2 +- .../weight_compression/openvino_backend.py | 64 +++++++++++++++++++ .../weight_compression/torch_backend.py | 2 + src/nncf/quantization/quantize_model.py | 23 +++++++ src/nncf/torch/function_hook/strip.py | 2 +- 9 files changed, 140 insertions(+), 5 deletions(-) create mode 100644 examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/main.py b/examples/llm_compression/torch/distillation_qat_with_lora/main.py index 1f6dcf8e964..14d84bc5474 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora/main.py @@ -162,7 +162,7 @@ def set_trainable(model: nn.Module, lora_lr: float, fq_lr: float) -> list[dict[s adapters_to_train = [] hook_storage = get_hook_storage(model) for _, module in hook_storage.named_hooks(): - if isinstance(module, (AsymmetricLoraQuantizer, SymmetricLoraQuantizer)) and (module.num_bits == 4): + if isinstance(module, (AsymmetricLoraQuantizer, SymmetricLoraQuantizer)) and (module.num_bits < 4): module.enable_gradients() params = module.get_trainable_params() adapters = module.get_adapters() @@ -298,7 +298,7 @@ def get_argument_parser() -> argparse.ArgumentParser: help="Learning rate for fine-tuning. " "For larger models (over 3 billion parameters), a learning rate of 5e-5 is recommended.", ) - parser.add_argument("--epochs", type=int, default=10, help="Number of epochs.") + parser.add_argument("--epochs", type=int, default=1, help="Number of epochs.") parser.add_argument("--batch_size", type=int, default=32, help="Size of training batch.") parser.add_argument( "--microbatch_size", @@ -321,7 +321,7 @@ def main(argv) -> float: device = "cuda" torch_dtype = torch.bfloat16 compression_config = dict( - mode=CompressWeightsMode.INT4_ASYM, + mode=CompressWeightsMode.INT3_SYM, group_size=64, awq=not args.basic_init, scale_estimation=not args.basic_init, diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py b/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py new file mode 100644 index 00000000000..24aa024ddb6 --- /dev/null +++ b/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py @@ -0,0 +1,8 @@ +from optimum.intel.openvino import OVModelForCausalLM + +from nncf.quantization.quantize_model import repack_weights + +model = OVModelForCausalLM.from_pretrained("output/last") +model.model = repack_weights(model.model) + +model.save_pretrained("output/last_repacked") diff --git a/src/nncf/openvino/quantization/quantize_model.py b/src/nncf/openvino/quantization/quantize_model.py index 19031fb1674..39312e0636c 100644 --- a/src/nncf/openvino/quantization/quantize_model.py +++ b/src/nncf/openvino/quantization/quantize_model.py @@ -422,3 +422,13 @@ def compress_weights_impl( statistics_points = statistics_aggregator.statistic_points return compression_algorithm.apply(model, graph, statistics_points, dataset) + + +def repack_weights_impl(model: ov.Model) -> ov.Model: + """ + Implementation of the `repack_weights()` method for the OpenVINO backend. + """ + model = remove_friendly_name_duplicates(model) + graph = build_graph(model) + + return WeightCompression.repack_weights(model, graph) diff --git a/src/nncf/quantization/algorithms/weight_compression/algorithm.py b/src/nncf/quantization/algorithms/weight_compression/algorithm.py index 20c3de05102..c80cd4bf299 100644 --- a/src/nncf/quantization/algorithms/weight_compression/algorithm.py +++ b/src/nncf/quantization/algorithms/weight_compression/algorithm.py @@ -1423,3 +1423,31 @@ def _get_statistics_for_weights_compression( for node in matmul_nodes: statistics[node.node_name] = copy.deepcopy(stats) return statistics + + @staticmethod + def repack_weights(model: TModel, graph: NNCFGraph) -> TModel: + """ + Repacks compressed weight constants to a lower-bit representation when possible. + + Iterates over all nodes in the graph and attempts to replace compressed constants + (i8, u8, i4, u4) with a more compact symmetric representation (i2, i3, i6) if the + actual value range allows it. + + :param model: Backend-specific model with compressed weights. + :param graph: NNCFGraph instance corresponding to the model. + :return: The model with repacked weight constants. + """ + backend = get_backend(model) + if backend != BackendType.OPENVINO: + msg = f"Unsupported type of backend: {backend}" + raise nncf.UnsupportedBackendError(msg) + + from nncf.quantization.algorithms.weight_compression.openvino_backend import OVWeightCompressionAlgoBackend + + backend_entity = OVWeightCompressionAlgoBackend(model) + + for node in graph.topological_sort(): + is_repacked = backend_entity.try_repack(node, graph) + if is_repacked: + print(f"Repacked node: {node.node_name}") + return model diff --git a/src/nncf/quantization/algorithms/weight_compression/constants.py b/src/nncf/quantization/algorithms/weight_compression/constants.py index 449c1b67ff9..6ee76f6155c 100644 --- a/src/nncf/quantization/algorithms/weight_compression/constants.py +++ b/src/nncf/quantization/algorithms/weight_compression/constants.py @@ -113,7 +113,7 @@ } -MIN_INPUT_SIZE_FOR_OPTIMIZED_COMPRESSION = 10000 +MIN_INPUT_SIZE_FOR_OPTIMIZED_COMPRESSION = 2 # 10000 OPTIMIZED_COMPRESSION_COMPATIBLE_INT_MODES = ( CompressWeightsMode.INT8_ASYM, diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index 0af6caec739..6d0c42a7790 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -101,6 +101,70 @@ def noop_metatypes(self) -> list[type[OperatorMetatype]]: def is_node_with_weights(node: NNCFNode, graph: NNCFGraph) -> bool: return node.layer_attributes and node.layer_attributes.constant_attributes + def is_compressed_constant(self, node: NNCFNode) -> bool: + res = node.metatype == om.OVConstantMetatype and node.node_name in self.name_to_node_mapping + + if res: + ov_node = self.name_to_node_mapping[node.node_name] + res = ov_node.tensor_view.element_type in [ov.Type.i8, ov.Type.u8, ov.Type.i4, ov.Type.u4] + return res + + def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: + if not self.is_compressed_constant(node): + return False + + next_nodes = graph.get_next_nodes(node) + if len(next_nodes) != 1 or next_nodes[0].metatype != om.OVConvertMetatype: + return False + + ov_node = self.name_to_node_mapping[node.node_name] + bits = 8 if ov_node.tensor_view.element_type in [ov.Type.i8, ov.Type.u8] else 4 + asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4] + + # only repack for sym types are supported for now: i6, i3, i2 + if asym: + return False + + weight_tensor = Tensor(get_const_value_as_ov_tensor(ov_node)) + weight_tensor = weight_tensor.as_numpy_tensor().astype(TensorDataType.int32) + shape = list(ov_node.tensor_view.shape) + + if len(shape) == 1 or shape[-1] == 1: + # Probably zero point + return False + + max_val = max(int(weight_tensor.max().item()), abs(int(weight_tensor.min().item()))) + optimal_bits = int(max_val).bit_length() + + supported_optimal_bits = [2, 3, 6] + + # If the optimal bits is equal to the current bits, no need to repack + if optimal_bits == bits or optimal_bits not in supported_optimal_bits: + return False + + const_node_name = node.node_name + compression_dtype = ov.Type.u2 if optimal_bits == 2 else ov.Type.u3 if optimal_bits == 3 else ov.Type.u6 + + # Repack the weight tensor to the optimal bits and change constant node with new + # representation [u2/u3/u6] - zero_point per tensor + offset = 2 ** (optimal_bits - 1) + compressed_tensor = weight_tensor + offset + compressed_const = create_ov_const_from_tensor(compressed_tensor, compression_dtype, name=const_node_name) + converted_const = opset.convert(compressed_const, ov.Type.f16) + + zero_point_const = opset.constant(offset, dtype=ov.Type.i8, name=f"{const_node_name}/zero_point") + zero_point_const = opset.convert(zero_point_const, ov.Type.f16) + + converted_const = opset.subtract( + converted_const, zero_point_const, name=f"{const_node_name}/zero_point/subtract" + ) + + # replace the convert after old const value with the new converted_const + ov_node = self.name_to_node_mapping[next_nodes[0].node_name] + self._replace_node(ov_node, converted_const) + + return True + @staticmethod def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: NNCFGraph) -> tuple[int] | None: channel_axes = get_weight_channel_axes(node_with_weight) diff --git a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py index 0a53e872709..031ae01853a 100644 --- a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py @@ -263,6 +263,8 @@ def get_fq_insertion_command( compression_config = wc_params.compression_config # default mapping for 4bit weight compression and FQ_LORA format, no need to add lora adapters for 8bit weight mode_vs_schema_map = { + CompressWeightsMode.INT2_SYM: QuantizationScheme.SYMMETRIC_LORA, + CompressWeightsMode.INT3_SYM: QuantizationScheme.SYMMETRIC_LORA, CompressWeightsMode.INT4_ASYM: QuantizationScheme.ASYMMETRIC_LORA, CompressWeightsMode.INT4_SYM: QuantizationScheme.SYMMETRIC_LORA, CompressWeightsMode.INT8_ASYM: QuantizationScheme.ASYMMETRIC, diff --git a/src/nncf/quantization/quantize_model.py b/src/nncf/quantization/quantize_model.py index cdb2e0bdba3..36fa38c1559 100644 --- a/src/nncf/quantization/quantize_model.py +++ b/src/nncf/quantization/quantize_model.py @@ -782,3 +782,26 @@ def quantize_with_tune_hyperparams( quantized_model = hyperparameter_tuner.apply(model, validation_dataset) return quantized_model + + +@api(canonical_alias="nncf.repack_weights") +def repack_weights( + model: TModel, +) -> TModel: + """ + Looking for 4 and 8 bit weights in OV model and repack them if maximal absolute value corresponds + to the supported type with lower bits. + + :param model: A model to be repacked. + :type model: TModel + :return: The non-trainable model with repacked weights or the same model. + """ + backend = get_backend(model) + + if backend != BackendType.OPENVINO: + msg = f"Unsupported type of backend: {backend}" + raise nncf.UnsupportedBackendError(msg) + + from nncf.openvino.quantization.quantize_model import repack_weights_impl + + return repack_weights_impl(model) diff --git a/src/nncf/torch/function_hook/strip.py b/src/nncf/torch/function_hook/strip.py index 2a7d64b4221..31c87421561 100644 --- a/src/nncf/torch/function_hook/strip.py +++ b/src/nncf/torch/function_hook/strip.py @@ -126,7 +126,7 @@ def replace_quantizer_to_compressed_weight_with_decompressor(model: TModel) -> T msg = "" if hook_module._qspec.half_range or hook_module._qspec.narrow_range: msg += "Unexpected parameters of quantizers on strip: half_range and narrow_range should be False.\n" - if hook_module.num_bits not in [4, 8]: + if hook_module.num_bits not in [2, 3, 4, 8]: msg += f"Unsupported number of bits {hook_module.num_bits} for the quantizer {hook_module}.\n" if msg: raise nncf.ValidationError(msg) From 029d4cfb8a080f39afe6b59c720d941d6484602d Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 30 Jul 2026 15:55:42 +0200 Subject: [PATCH 02/27] 1) Removed debug code. 2) Removed 6bit from repacking. --- .../distillation_qat_with_lora/repack_weights.py | 11 +++++++++++ .../algorithms/weight_compression/constants.py | 2 +- .../algorithms/weight_compression/openvino_backend.py | 8 ++++---- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py b/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py index 24aa024ddb6..f56ca27d633 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py @@ -1,3 +1,14 @@ +# Copyright (c) 2026 Intel Corporation +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from optimum.intel.openvino import OVModelForCausalLM from nncf.quantization.quantize_model import repack_weights diff --git a/src/nncf/quantization/algorithms/weight_compression/constants.py b/src/nncf/quantization/algorithms/weight_compression/constants.py index 6ee76f6155c..449c1b67ff9 100644 --- a/src/nncf/quantization/algorithms/weight_compression/constants.py +++ b/src/nncf/quantization/algorithms/weight_compression/constants.py @@ -113,7 +113,7 @@ } -MIN_INPUT_SIZE_FOR_OPTIMIZED_COMPRESSION = 2 # 10000 +MIN_INPUT_SIZE_FOR_OPTIMIZED_COMPRESSION = 10000 OPTIMIZED_COMPRESSION_COMPATIBLE_INT_MODES = ( CompressWeightsMode.INT8_ASYM, diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index 6d0c42a7790..922ab268d7b 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -121,7 +121,7 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: bits = 8 if ov_node.tensor_view.element_type in [ov.Type.i8, ov.Type.u8] else 4 asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4] - # only repack for sym types are supported for now: i6, i3, i2 + # only repack for sym types are supported for now: i3, i2 if asym: return False @@ -136,17 +136,17 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: max_val = max(int(weight_tensor.max().item()), abs(int(weight_tensor.min().item()))) optimal_bits = int(max_val).bit_length() - supported_optimal_bits = [2, 3, 6] + supported_optimal_bits = [2, 3] # If the optimal bits is equal to the current bits, no need to repack if optimal_bits == bits or optimal_bits not in supported_optimal_bits: return False const_node_name = node.node_name - compression_dtype = ov.Type.u2 if optimal_bits == 2 else ov.Type.u3 if optimal_bits == 3 else ov.Type.u6 + compression_dtype = ov.Type.u2 if optimal_bits == 2 else ov.Type.u3 # Repack the weight tensor to the optimal bits and change constant node with new - # representation [u2/u3/u6] - zero_point per tensor + # representation [u2/u3] - zero_point per tensor offset = 2 ** (optimal_bits - 1) compressed_tensor = weight_tensor + offset compressed_const = create_ov_const_from_tensor(compressed_tensor, compression_dtype, name=const_node_name) From 03ceee72754db9d062db9f5bab97f0c21a5916ef Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 30 Jul 2026 16:39:01 +0200 Subject: [PATCH 03/27] Fixed style bugs. --- src/nncf/quantization/__init__.py | 1 + src/nncf/quantization/quantize_model.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/nncf/quantization/__init__.py b/src/nncf/quantization/__init__.py index cf5a6c6133e..0a1e91cabcb 100644 --- a/src/nncf/quantization/__init__.py +++ b/src/nncf/quantization/__init__.py @@ -14,3 +14,4 @@ from nncf.quantization.quantize_model import compress_weights as compress_weights from nncf.quantization.quantize_model import quantize as quantize from nncf.quantization.quantize_model import quantize_with_accuracy_control as quantize_with_accuracy_control +from nncf.quantization.quantize_model import repack_weights as repack_weights diff --git a/src/nncf/quantization/quantize_model.py b/src/nncf/quantization/quantize_model.py index 36fa38c1559..5d2fab5ece4 100644 --- a/src/nncf/quantization/quantize_model.py +++ b/src/nncf/quantization/quantize_model.py @@ -804,4 +804,4 @@ def repack_weights( from nncf.openvino.quantization.quantize_model import repack_weights_impl - return repack_weights_impl(model) + return repack_weights_impl(model) # type: ignore From db5a21ab61e3736a5448b74456f3298045a58ecb Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 30 Jul 2026 20:08:46 +0200 Subject: [PATCH 04/27] Added test for weight repacking. --- src/nncf/__init__.py | 1 + tests/openvino/native/models.py | 48 +++++++++++++++++++ tests/openvino/native/test_model_repacking.py | 33 +++++++++++++ 3 files changed, 82 insertions(+) create mode 100644 tests/openvino/native/test_model_repacking.py diff --git a/src/nncf/__init__.py b/src/nncf/__init__.py index 475d7e836fc..2a1ee31e549 100644 --- a/src/nncf/__init__.py +++ b/src/nncf/__init__.py @@ -54,6 +54,7 @@ from nncf.quantization import compress_weights as compress_weights from nncf.quantization import quantize as quantize from nncf.quantization import quantize_with_accuracy_control as quantize_with_accuracy_control +from nncf.quantization import repack_weights as repack_weights from nncf.quantization.advanced_parameters import ( AdvancedAccuracyRestorerParameters as AdvancedAccuracyRestorerParameters, ) diff --git a/tests/openvino/native/models.py b/tests/openvino/native/models.py index 5f841e690fc..ba0eb6bfa6b 100644 --- a/tests/openvino/native/models.py +++ b/tests/openvino/native/models.py @@ -1463,3 +1463,51 @@ def _create_ov_model(self): model = ov.Model([result_node, result_node1], [input_node], name="ParallelEdgesOutputPortIdModel") return model + + +class ModelForRepack(OVReferenceModel): + HIDDEN_DIM = 16 + INPUT_SHAPE = [2, 24, HIDDEN_DIM] # [B, SeqLen, HiddenDim] + + @staticmethod + def get_weights(weights_data, w_bits, name): + if w_bits not in [4, 8, 32]: + msg = f"Unsupported weight bits: {w_bits}. Supported values are 4, 8, and 32." + raise ValueError(msg) + if w_bits == 32: + return opset.constant(weights_data, dtype=np.float32, name=name) + qw = opset.constant(weights_data, dtype=ov.Type.i8 if w_bits == 8 else ov.Type.i4, name="qw_" + name) + qw = opset.convert(qw, destination_type=np.float32) + + scale = opset.constant( + np.ones((weights_data.shape[-2], 1), dtype=np.float32), dtype=np.float32, name="scale_" + name + ) + + return qw * scale + + def _create_ov_model( + self, + input_shape=None, + ): + self._input_shape = self.INPUT_SHAPE if input_shape is None else input_shape + + input_1 = opset.parameter(self._input_shape, name="Input") + + model_config = [ + {"w_bits": 32, "data_range": [-128, 128]}, + {"w_bits": 8, "data_range": [-8, 7]}, + {"w_bits": 4, "data_range": [-4, 3]}, + {"w_bits": 4, "data_range": [-2, 1]}, + ] + + x = opset.relu(input_1) + for i, config in enumerate(model_config): + data = np.random.randint( + config["data_range"][0], config["data_range"][1] + 1, size=(self.HIDDEN_DIM, self.HIDDEN_DIM) + ).astype(np.int32) + weights = ModelForRepack.get_weights(data, w_bits=config["w_bits"], name=f"weights_{i + 1}") + x = opset.matmul(x, weights, transpose_a=False, transpose_b=True, name=f"MatMul_{i + 1}") + + result = x + model = ov.Model([result], [input_1]) + return model diff --git a/tests/openvino/native/test_model_repacking.py b/tests/openvino/native/test_model_repacking.py new file mode 100644 index 00000000000..0786f896b0d --- /dev/null +++ b/tests/openvino/native/test_model_repacking.py @@ -0,0 +1,33 @@ +# Copyright (c) 2026 Intel Corporation +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import openvino as ov + +import nncf +from tests.openvino.native.models import ModelForRepack + + +def get_constant_element_types(model: ov.Model) -> list[ov.Type]: + """Collects element types of all Constant nodes in the model.""" + types = [] + for op in model.get_ops(): + if op.get_type_name() == "Constant": + types.append(op.get_element_type()) + return types + + +def test_repack_weights_produces_u3_and_u2(): + model = ModelForRepack().ov_model + 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" From 707c35fe9b9cbe821ee7f914f2fc06a416193339 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 30 Jul 2026 20:10:55 +0200 Subject: [PATCH 05/27] Added test for weight repacking. --- tests/openvino/native/test_model_repacking.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/openvino/native/test_model_repacking.py b/tests/openvino/native/test_model_repacking.py index 0786f896b0d..cdde111eb66 100644 --- a/tests/openvino/native/test_model_repacking.py +++ b/tests/openvino/native/test_model_repacking.py @@ -26,6 +26,11 @@ def get_constant_element_types(model: ov.Model) -> list[ov.Type]: def test_repack_weights_produces_u3_and_u2(): model = ModelForRepack().ov_model + + element_types = get_constant_element_types(model) + assert ov.Type.u3 not in element_types, "Did not expect u3 constant before repacking" + assert ov.Type.u2 not in element_types, "Did not expect u2 constant before repacking" + repacked_model = nncf.repack_weights(model) element_types = get_constant_element_types(repacked_model) From 2c45ac39567278dd2d97e938633ca6406d8e674b Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Tue, 15 Sep 2026 10:27:51 +0200 Subject: [PATCH 06/27] Train 2-3 bit quantizers. --- .../main.py | 487 ++++++++++++++++++ 1 file changed, 487 insertions(+) create mode 100644 examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py new file mode 100644 index 00000000000..04ee019fa89 --- /dev/null +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -0,0 +1,487 @@ +# Copyright (c) 2026 Intel Corporation +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import shutil +import sys +import warnings +from datetime import datetime +from pathlib import Path +from pprint import pprint +from typing import Any + +import faulthandler, signal +faulthandler.register(signal.SIGUSR1, all_threads=True) + +import torch +import torch.nn.functional as F +import transformers +from datasets import load_dataset +from lm_eval import simple_evaluate +from lm_eval.models.optimum_lm import OptimumLM +from optimum.exporters.openvino.convert import export_from_model +from optimum.intel.openvino import OVModelForCausalLM +from optimum.modeling_base import OptimizedModel +from torch import Tensor +from torch import nn +from torch.jit import TracerWarning +from torch.utils.tensorboard import SummaryWriter +from transformers import AutoModelForCausalLM +from transformers import AutoTokenizer + +import nncf +from nncf.common.logging.track_progress import track +from nncf.data.dataset import Dataset +from nncf.parameters import CompressionFormat +from nncf.parameters import CompressWeightsMode +from nncf.parameters import StripFormat +from nncf.quantization.advanced_parameters import AdvancedAWQParameters +from nncf.quantization.advanced_parameters import AdvancedCompressionParameters +from nncf.quantization.quantize_model import compress_weights +from nncf.torch import load_from_config +from nncf.torch.function_hook.wrapper import get_hook_storage +from nncf.torch.quantization.layers import AsymmetricLoraQuantizer +from nncf.torch.quantization.layers import SymmetricLoraQuantizer +from nncf.quantization.quantize_model import repack_weights + +warnings.filterwarnings("ignore", category=TracerWarning) + + +def get_wikitext2(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device) -> list[Tensor]: + """ + Loads and processes the Wikitext-2 dataset for training. + + :param num_samples: Number of samples to generate. + :param seqlen: Sequence length for each sample. + :param tokenizer: Tokenizer to encode the text. + :param device: Device to move the tensors to (e.g., 'cpu' or 'cuda'). + :return: A list of tensors containing the tokenized text samples. + """ + traindata = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train") + limit = num_samples * seqlen // 4 # ~1k for 128 samples with seqlen=32 to be aligned with optimum + text = "".join([" \n" if s == "" else s for s in traindata["text"][:limit]]) + trainenc = tokenizer(text, return_tensors="pt") + trainloader = [] + for _ in range(num_samples): + # Crop a sequence of tokens of length seqlen starting at a random position + i = torch.randint(0, trainenc.input_ids.shape[1] - seqlen - 1, (1,)).item() + j = i + seqlen + inp = trainenc.input_ids[:, i:j].to(device) + trainloader.append(inp) + return trainloader + + +def measure_perplexity( + optimum_model: OptimizedModel, + max_length: int | None = None, + limit: int | float | None = None, +) -> float: + """ + Measure perplexity on the Wikitext dataset, via rolling loglikelihoods for a given model. + + :param optimum_model: A model to be evaluated. + :param max_length: The maximum sequence length for evaluation. + :param limit: Limit the number of examples per task (only use this for testing). + If <1, limit is a percentage of the total number of examples. + :return: The similarity score as a float. + """ + task = "wikitext" + print("#" * 50 + " Evaluate via lm-eval-harness " + "#" * 50) + lm_obj = OptimumLM(pretrained=optimum_model, max_length=max_length) + results = simple_evaluate(lm_obj, tasks=[task], limit=limit, log_samples=False) + return results["results"][task]["word_perplexity,none"] + + +@torch.no_grad() +def calc_hiddens(model: nn.Module, dataloader: list[Tensor]) -> list[Tensor]: + """ + Calculate the hidden states for each input in the dataloader using the given model. + + :param model: The model used to calculate the hidden states. + :param dataloader: The dataloader providing the inputs to the model. + :return: A list of hidden states for each input in the dataloader. + """ + orig_hiddens = [] + for data in track(dataloader, description="Calculating original hiddens"): + model_input = get_model_input(data) + orig_hiddens.append(model.model(**model_input).last_hidden_state) + torch.cuda.empty_cache() + return orig_hiddens + + +def get_model_input(input_ids: Tensor) -> dict[str, Tensor]: + """ + Prepares the model input dictionary with input IDs, attention mask, and position IDs. + + :param input_ids: Tensor containing the input IDs. + :return: A dictionary with keys "input_ids", "attention_mask", and "position_ids", + each mapping to their respective tensors. + """ + attention_mask = torch.ones_like(input_ids) + position_ids = torch.cumsum(attention_mask, axis=1) - 1 + return {"input_ids": input_ids, "attention_mask": attention_mask, "position_ids": position_ids} + + +def kl_div(student_hiddens: torch.Tensor, teacher_hiddens: torch.Tensor) -> torch.Tensor: + """ + Computes the Kullback-Leibler divergence loss between the student and teacher hidden states. + The input tensors are expected to have the same shape, and the last dimension represents the number of classes. + + :param student_hiddens: The hidden states from the student model. + :param teacher_hiddens: The hidden states from the teacher model. + :returns: The computed KL divergence loss. + """ + num_classes = student_hiddens.shape[-1] + return F.kl_div( + input=F.log_softmax(student_hiddens.view(-1, num_classes), dim=-1), + target=F.log_softmax(teacher_hiddens.view(-1, num_classes), dim=-1), + log_target=True, + reduction="batchmean", + ) + + +def set_trainable( + model: nn.Module, lora_lr: float, fq_lr: float, scales_only: bool = False +) -> list[dict[str, Any]]: + """ + Sets the trainable parameters of the model for quantization-aware training with LoRA (Low-Rank Adaptation). + + This function disables gradients for all parameters in the model, then selectively enables gradients for + specific quantizers (AsymmetricLoraQuantizer, SymmetricLoraQuantizer) that have 4-bit quantization. + It collects the trainable parameters and adapters from these quantizers and returns them in a format + suitable for an optimizer. + + :param model: The model to be trained. + :param lora_lr: Learning rate for the LoRA adapters. + :param fq_lr: Learning rate for the quantizer scales. + :param scales_only: If True, only quantizer scales are trainable; LoRA adapters are frozen. + :return: A list of dictionaries containing the parameters to be optimized and their corresponding learning rates. + """ + model.requires_grad_(False) + scales_to_train = [] + adapters_to_train = [] + hook_storage = get_hook_storage(model) + for _, module in hook_storage.named_hooks(): + if isinstance(module, (AsymmetricLoraQuantizer, SymmetricLoraQuantizer)) and (module.num_bits < 4): + module.enable_gradients() + params = module.get_trainable_params() + adapters = module.get_adapters() + if scales_only: + for adapter in adapters.values(): + adapter.requires_grad_(False) + else: + adapters_to_train.extend(adapters.values()) + scales_to_train.extend(param for name, param in params.items() if name not in adapters) + + params = list(model.parameters()) + trainable_params = sum(p.numel() for p in params if p.requires_grad) + all_param = sum(p.numel() for p in params) + print( + f"trainable params: {trainable_params:,d} || " + f"all params: {all_param:,d} || " + f"trainable%: {100 * trainable_params / all_param:.4f}" + ) + model.train() + if scales_only: + return [{"params": scales_to_train, "lr": fq_lr}] + return [{"params": adapters_to_train, "lr": lora_lr}, {"params": scales_to_train, "lr": fq_lr}] + + +def save_checkpoint(model: nn.Module, ckpt_file: Path, model_state: bool = True) -> None: + """ + Stores the current state of a quantized model to a checkpoint file. + + :param model: The model whose state will be saved to checkpoint. + :param ckpt_file: Path to store the checkpoint file. + :param model_state: Whether to save the complete model weights in addition to NNCF state. Required when using + AWQ method which fuses scaling factors into weights. When False, only NNCF configuration and state are saved, + as they're maintained separately from the model's weights. + """ + hook_storage = get_hook_storage(model) + ckpt = {"nncf_state_dict": hook_storage.state_dict(), "nncf_config": nncf.torch.get_config(model)} + if model_state: + ckpt["model_state"] = model.state_dict() + torch.save(ckpt, ckpt_file) + + +def load_checkpoint(model: nn.Module, ckpt_file: Path) -> nn.Module: + """ + Loads the state of a tuned model from a checkpoint. This function restores the placement of Fake Quantizers (FQs) + with absorbable LoRA adapters and loads their parameters. + + :param model: The model to load the checkpoint into. + :param ckpt_file: Path to the checkpoint file. + :returns: The model with the loaded NNCF state from checkpoint. + """ + ckpt = torch.load(ckpt_file, weights_only=False, map_location="cpu") + model = load_from_config(model, ckpt["nncf_config"]) + if "model_state" in ckpt: + model.load_state_dict(ckpt["model_state"]) + hook_storage = get_hook_storage(model) + hook_storage.load_state_dict(ckpt["nncf_state_dict"]) + return model + + +@torch.no_grad() +def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path) -> OVModelForCausalLM: + """ + Create a wrapper of OpenVINO model from the checkpoint for evaluation on CPU via WWB. + + :param pretrained: The name or path of the pretrained model. + :param ckpt_file: The path to the checkpoint file to load the model weights and NNCF configurations. + :param ir_dir: The directory where the OpenVINO model will be saved. + :return: A wrapper of OpenVINO model ready for evaluation. + """ + model_to_eval = AutoModelForCausalLM.from_pretrained(pretrained, torch_dtype=torch.float32, device_map="cpu") + model_to_eval = load_checkpoint(model_to_eval, ckpt_file) + model_to_eval = nncf.strip(model_to_eval, do_copy=False, strip_format=StripFormat.DQ) + export_from_model(model_to_eval, ir_dir, device="cpu") + + model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) + model_to_eval.model = repack_weights(model_to_eval.model) + model_to_eval.save_pretrained(ir_dir/"repacked") + + + return OVModelForCausalLM.from_pretrained( + model_id=ir_dir/"repacked", + trust_remote_code=True, + load_in_8bit=False, + compile=True, + ) + + +def limit_type(astr: str): + value = float(astr) + if value < 0 or value > 1: + msg = "value not in range [0,1]" + raise argparse.ArgumentTypeError(msg) + return value + + +def get_argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(add_help=True) + + # Model params + parser.add_argument( + "--pretrained", + type=str, + #default="Qwen/Qwen3-8B", + default="meta-llama/Llama-3.2-1B-Instruct", + help="The model id or path of a pretrained HF model configuration.", + ) + parser.add_argument( + "--bits", + type=int, + default=3, + help="Number of bits for weight compression (e.g., 2, 3, 4, 8).", + ) + parser.add_argument( + "--output_dir", + type=Path, + default="output", + help="Path to the directory for storing logs, tuning checkpoint, compressed model, validation references.", + ) + parser.add_argument( + "--resume", + action="store_true", + help="Whether to start from previously saved checkpoint. If not specified or checkpoint does not exist, " + "start from scratch by post-training weight compression initialization.", + ) + parser.add_argument("--lora_rank", type=int, default=256, help="Rank of lora adapters") + parser.add_argument( + "--basic_init", + action="store_true", + help="Whether to initialize quantization with basic min-max round-to-nearest schema. By default, advanced " + "data-aware post-training methods are used: AWQ + Scale Estimation. These methods typically provide better " + "accuracy, but require a calibration dataset and additional initialization time " + "(~20 sec for 1B and ~80 sec for 8B models).", + ) + + # Data params + parser.add_argument("--num_train_samples", type=int, default=2048, help="Number of training samples") + parser.add_argument("--train_seqlen", type=int, default=1024, help="Train data context length.") + parser.add_argument("--eval_seqlen", type=int, default=2048, help="Evaluation data context length.") + parser.add_argument( + "--limit", + type=limit_type, + default=None, + help="A percentage of the total number of examples for evaluation. " + "Should be on the range [0,1]. If None, all samples will be used.", + ) + + # Training params + parser.add_argument( + "--lr", + type=float, + default=1e-4, + help="Learning rate for fine-tuning. " + "For larger models (over 3 billion parameters), a learning rate of 5e-5 is recommended.", + ) + parser.add_argument("--epochs", type=int, default=1, help="Number of epochs.") + parser.add_argument( + "--scale_finetune_epochs", + type=int, + default=0, + help="Number of additional epochs after the main training loop during which only the quantizer scales are " + "finetuned (LoRA adapters are frozen). Set to 0 to skip this phase.", + ) + parser.add_argument("--batch_size", type=int, default=32, help="Size of training batch.") + parser.add_argument( + "--microbatch_size", + type=int, + default=8, + help="Size of each training microbatch. Gradients will be accumulated until the batch size is reached.", + ) + return parser + + +def main(argv) -> float: + """ + Fine-tunes the specified model and returns the difference between initial and best validation perplexity in Torch, + and the test perplexity for best model exported to OpenVINO. + """ + parser = get_argument_parser() + args = parser.parse_args(argv) + assert torch.cuda.is_available() + transformers.set_seed(42) + device = "cuda" + torch_dtype = torch.bfloat16 + compression_config = dict( + mode=CompressWeightsMode.INT3_SYM if args.bits == 3 else CompressWeightsMode.INT2_SYM, + group_size=64, + awq=not args.basic_init, + scale_estimation=not args.basic_init, + compression_format=CompressionFormat.FQ_LORA, + ) + pprint({"CLI arguments": vars(args), "Major compression parameters": compression_config}) + compression_config["advanced_parameters"] = AdvancedCompressionParameters( + awq_params=AdvancedAWQParameters(prefer_data_aware_scaling=not args.basic_init), + lora_adapter_rank=args.lora_rank, + ) + # Configure output and log files. + output_dir = Path(args.output_dir) + tensorboard_dir = output_dir / "tb" / datetime.now().strftime("%Y-%m-%d__%H-%M-%S") + last_dir = output_dir / f"last_bits_{args.bits}" + if not args.resume: + shutil.rmtree(last_dir, ignore_errors=True) + for path in [output_dir, tensorboard_dir, last_dir]: + path.mkdir(exist_ok=True, parents=True) + ckpt_file = last_dir / "nncf_checkpoint.pth" + print(f"To visualize the loss and validation metrics, open Tensorboard using the logs from: {tensorboard_dir}") + tb = SummaryWriter(tensorboard_dir, "QAT with absorbable LoRA") + + # Load original model and tokenizer. + model = AutoModelForCausalLM.from_pretrained(args.pretrained, torch_dtype=torch_dtype, device_map="auto") + tokenizer = AutoTokenizer.from_pretrained(args.pretrained) + + # Prepare training and calibration data + train_loader = get_wikitext2( + num_samples=args.num_train_samples, seqlen=args.train_seqlen, tokenizer=tokenizer, device=device + ) + if args.basic_init: + example_input = {k: v.to(device) for k, v in model.dummy_inputs.items()} + dataset = Dataset([example_input]) + else: + calib_loader = get_wikitext2(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) + dataset = Dataset(map(get_model_input, calib_loader)) + + # Pre-compute hiddens of teacher model for distillation loss. + orig_hiddens = calc_hiddens(model, train_loader) + + # Create or load model to tune with Fake Quantizers and absorbable LoRA adapters. + if args.resume and ckpt_file.exists(): + model = load_checkpoint(model, ckpt_file) + else: + model = compress_weights(model, dataset=dataset, **compression_config) + save_checkpoint(model, ckpt_file, model_state=not args.basic_init) + fq_lr = args.lr / 10 + weight_decay = args.lr + param_to_train = set_trainable(model, lora_lr=args.lr, fq_lr=fq_lr) + opt = torch.optim.AdamW(param_to_train, weight_decay=weight_decay) + + # Run tuning with distillation loss and validation after each epoch. + grad_accumulation_steps = args.batch_size // args.microbatch_size + num_samples = len(train_loader) + epoch_samples = num_samples - num_samples % args.microbatch_size + microbatches_per_epoch = epoch_samples // args.microbatch_size + + def run_training( + num_epochs: int, optimizer: torch.optim.Optimizer, phase_desc: str, start_total_steps: int + ) -> int: + loss_numerator = grad_steps = 0 + total_steps = start_total_steps + for epoch in range(num_epochs): + batch_indices_epoch = torch.randperm(num_samples)[:epoch_samples].chunk(microbatches_per_epoch) + epoch_tracker = track(batch_indices_epoch, description=f"{phase_desc} {epoch}") + for indices in epoch_tracker: + indices = indices.tolist() + + def form_batch(inputs: list[Tensor], model_input: bool): + batch = torch.cat([inputs[i] for i in indices], dim=0) + return get_model_input(batch) if model_input else batch.to(device=device, dtype=torch_dtype) + + # Compute distillation loss between logits of the original model and the model with FQ + LoRA. + inputs = form_batch(train_loader, model_input=True) + with torch.no_grad(): + targets = model.lm_head(form_batch(orig_hiddens, model_input=False)) + if hasattr(model.config, "final_logit_softcapping"): # Gemma has post-processing after lm_head + fls = model.config.final_logit_softcapping + if fls is not None: + targets = targets / fls + targets = torch.tanh(targets) + targets = targets * fls + outputs = model(**inputs).logits + loss = kl_div(outputs, targets.to(dtype=torch_dtype, device=device)) + + # Perform an optimization step after accumulating gradients over multiple minibatches. + loss_numerator += loss.item() + grad_steps += 1 + if not torch.isfinite(loss).item(): + err = f"Fine-tuning loss is {loss}" + raise ValueError(err) + (loss / grad_accumulation_steps).backward() + if grad_steps == grad_accumulation_steps: + optimizer.step() + optimizer.zero_grad() + aggregated_loss = loss_numerator / grad_steps + loss_numerator = grad_steps = 0 + total_steps += 1 + tb.add_scalar("loss", aggregated_loss, total_steps) + epoch_tracker.update( + advance=0, description=f"{phase_desc} {epoch} | loss={aggregated_loss:.4f}" + ) + + save_checkpoint(model, ckpt_file, model_state=not args.basic_init) + return total_steps + + total_steps = run_training(args.epochs, opt, phase_desc="Train epoch", start_total_steps=0) + + # Optional scales-only finetuning phase: LoRA adapters are frozen, only quantizer scales are trained. + if args.scale_finetune_epochs > 0: + scale_params = set_trainable(model, lora_lr=args.lr, fq_lr=fq_lr, scales_only=True) + opt = torch.optim.AdamW(scale_params, weight_decay=weight_decay) + run_training( + args.scale_finetune_epochs, opt, phase_desc="Scales-only epoch", start_total_steps=total_steps + ) + + del model + # Export the best tuned model to OpenVINO and evaluate it using LM-Evaluation-Harness. + model_for_eval = export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) + ov_perplexity = measure_perplexity(model_for_eval, args.eval_seqlen, args.limit) + tb.add_scalar("ov_perplexity", ov_perplexity, 0) + print( + f"The finetuned model has been exported to OpenVINO and saved to: {last_dir}\n" + f"The word perplexity on wikitext (test) = {ov_perplexity:.4f}" + ) + return ov_perplexity + + +if __name__ == "__main__": + main(sys.argv[1:]) From b0f02404a67f23a9f36e4644989403b1ab262a8f Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 16 Sep 2026 08:41:40 +0200 Subject: [PATCH 07/27] Revert experimental changes. --- .../torch/distillation_qat_with_lora/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/main.py b/examples/llm_compression/torch/distillation_qat_with_lora/main.py index 14d84bc5474..1f6dcf8e964 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora/main.py @@ -162,7 +162,7 @@ def set_trainable(model: nn.Module, lora_lr: float, fq_lr: float) -> list[dict[s adapters_to_train = [] hook_storage = get_hook_storage(model) for _, module in hook_storage.named_hooks(): - if isinstance(module, (AsymmetricLoraQuantizer, SymmetricLoraQuantizer)) and (module.num_bits < 4): + if isinstance(module, (AsymmetricLoraQuantizer, SymmetricLoraQuantizer)) and (module.num_bits == 4): module.enable_gradients() params = module.get_trainable_params() adapters = module.get_adapters() @@ -298,7 +298,7 @@ def get_argument_parser() -> argparse.ArgumentParser: help="Learning rate for fine-tuning. " "For larger models (over 3 billion parameters), a learning rate of 5e-5 is recommended.", ) - parser.add_argument("--epochs", type=int, default=1, help="Number of epochs.") + parser.add_argument("--epochs", type=int, default=10, help="Number of epochs.") parser.add_argument("--batch_size", type=int, default=32, help="Size of training batch.") parser.add_argument( "--microbatch_size", @@ -321,7 +321,7 @@ def main(argv) -> float: device = "cuda" torch_dtype = torch.bfloat16 compression_config = dict( - mode=CompressWeightsMode.INT3_SYM, + mode=CompressWeightsMode.INT4_ASYM, group_size=64, awq=not args.basic_init, scale_estimation=not args.basic_init, From ab544bb892b62505b5ec04f33f1abe279da642ba Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 16 Sep 2026 16:00:16 +0200 Subject: [PATCH 08/27] Fixed style and removed unused code. --- .../repack_weights.py | 19 -- .../main.py | 193 ++++++++++++------ .../weight_compression/torch_backend.py | 9 +- 3 files changed, 136 insertions(+), 85 deletions(-) delete mode 100644 examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py b/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py deleted file mode 100644 index f56ca27d633..00000000000 --- a/examples/llm_compression/torch/distillation_qat_with_lora/repack_weights.py +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright (c) 2026 Intel Corporation -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from optimum.intel.openvino import OVModelForCausalLM - -from nncf.quantization.quantize_model import repack_weights - -model = OVModelForCausalLM.from_pretrained("output/last") -model.model = repack_weights(model.model) - -model.save_pretrained("output/last_repacked") diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 04ee019fa89..e879f3fdb08 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -17,9 +17,6 @@ from pprint import pprint from typing import Any -import faulthandler, signal -faulthandler.register(signal.SIGUSR1, all_threads=True) - import torch import torch.nn.functional as F import transformers @@ -45,11 +42,11 @@ from nncf.quantization.advanced_parameters import AdvancedAWQParameters from nncf.quantization.advanced_parameters import AdvancedCompressionParameters from nncf.quantization.quantize_model import compress_weights +from nncf.quantization.quantize_model import repack_weights from nncf.torch import load_from_config from nncf.torch.function_hook.wrapper import get_hook_storage from nncf.torch.quantization.layers import AsymmetricLoraQuantizer from nncf.torch.quantization.layers import SymmetricLoraQuantizer -from nncf.quantization.quantize_model import repack_weights warnings.filterwarnings("ignore", category=TracerWarning) @@ -111,7 +108,7 @@ def calc_hiddens(model: nn.Module, dataloader: list[Tensor]) -> list[Tensor]: orig_hiddens = [] for data in track(dataloader, description="Calculating original hiddens"): model_input = get_model_input(data) - orig_hiddens.append(model.model(**model_input).last_hidden_state) + orig_hiddens.append(model.model(**model_input).last_hidden_state.cpu()) torch.cuda.empty_cache() return orig_hiddens @@ -147,9 +144,7 @@ def kl_div(student_hiddens: torch.Tensor, teacher_hiddens: torch.Tensor) -> torc ) -def set_trainable( - model: nn.Module, lora_lr: float, fq_lr: float, scales_only: bool = False -) -> list[dict[str, Any]]: +def set_trainable(model: nn.Module, lora_lr: float, fq_lr: float, scales_only: bool = False) -> list[dict[str, Any]]: """ Sets the trainable parameters of the model for quantization-aware training with LoRA (Low-Rank Adaptation). @@ -246,11 +241,10 @@ def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path) -> OVMode model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) model_to_eval.model = repack_weights(model_to_eval.model) - model_to_eval.save_pretrained(ir_dir/"repacked") - + model_to_eval.save_pretrained(ir_dir / "repacked") return OVModelForCausalLM.from_pretrained( - model_id=ir_dir/"repacked", + model_id=ir_dir / "repacked", trust_remote_code=True, load_in_8bit=False, compile=True, @@ -272,7 +266,7 @@ def get_argument_parser() -> argparse.ArgumentParser: parser.add_argument( "--pretrained", type=str, - #default="Qwen/Qwen3-8B", + # default="Qwen/Qwen3-8B", default="meta-llama/Llama-3.2-1B-Instruct", help="The model id or path of a pretrained HF model configuration.", ) @@ -280,8 +274,9 @@ def get_argument_parser() -> argparse.ArgumentParser: "--bits", type=int, default=3, - help="Number of bits for weight compression (e.g., 2, 3, 4, 8).", - ) + choices=[2, 3], + help="Number of bits for weight compression (2 or 3).", + ) parser.add_argument( "--output_dir", type=Path, @@ -342,6 +337,91 @@ def get_argument_parser() -> argparse.ArgumentParser: return parser +def run_training( + model: nn.Module, + train_loader: list[Tensor], + orig_hiddens: list[Tensor], + optimizer: torch.optim.Optimizer, + ckpt_file: Path, + tb: SummaryWriter, + num_epochs: int, + phase_desc: str, + start_total_steps: int, + *, + device: torch.device | str, + torch_dtype: torch.dtype, + grad_accumulation_steps: int, + num_samples: int, + epoch_samples: int, + microbatches_per_epoch: int, + model_state: bool, +) -> int: + """ + Run the distillation-based training loop for the compressed model. + + :param model: The model being tuned. + :param train_loader: Training samples used for distillation. + :param orig_hiddens: Teacher hidden states computed from the original model. + :param optimizer: Optimizer used for updates. + :param ckpt_file: Path to save the checkpoint after each epoch. + :param tb: TensorBoard writer for loss metrics. + :param num_epochs: Number of epochs to run. + :param phase_desc: Human-readable description of the current epoch phase. + :param start_total_steps: The initial global optimizer step count. + :param device: Device to place batch tensors on. + :param torch_dtype: The dtype used for the computation. + :param grad_accumulation_steps: Number of microbatches before an optimizer step. + :param num_samples: Number of training samples. + :param epoch_samples: The available number of samples in the epoch after rounding. + :param microbatches_per_epoch: Number of microbatch chunks in the epoch. + :param model_state: Whether to save full model weights in the checkpoint. + :return: The total number of optimizer steps performed. + """ + loss_numerator = grad_steps = 0 + total_steps = start_total_steps + for epoch in range(num_epochs): + batch_indices_epoch = torch.randperm(num_samples)[:epoch_samples].chunk(microbatches_per_epoch) + epoch_tracker = track(batch_indices_epoch, description=f"{phase_desc} {epoch}") + for indices in epoch_tracker: + indices = indices.tolist() + + def form_batch(inputs: list[Tensor], model_input: bool): + batch = torch.cat([inputs[i] for i in indices], dim=0) + return get_model_input(batch) if model_input else batch.to(device=device, dtype=torch_dtype) + + # Compute distillation loss between logits of the original model and the model with FQ + LoRA. + inputs = form_batch(train_loader, model_input=True) + with torch.no_grad(): + targets = model.lm_head(form_batch(orig_hiddens, model_input=False)) + if hasattr(model.config, "final_logit_softcapping"): # Gemma has post-processing after lm_head + fls = model.config.final_logit_softcapping + if fls is not None: + targets = targets / fls + targets = torch.tanh(targets) + targets = targets * fls + outputs = model(**inputs).logits + loss = kl_div(outputs, targets.to(dtype=torch_dtype, device=device)) + + # Perform an optimization step after accumulating gradients over multiple minibatches. + loss_numerator += loss.item() + grad_steps += 1 + if not torch.isfinite(loss).item(): + err = f"Fine-tuning loss is {loss}" + raise ValueError(err) + (loss / grad_accumulation_steps).backward() + if grad_steps == grad_accumulation_steps: + optimizer.step() + optimizer.zero_grad() + aggregated_loss = loss_numerator / grad_steps + loss_numerator = grad_steps = 0 + total_steps += 1 + tb.add_scalar("loss", aggregated_loss, total_steps) + epoch_tracker.update(advance=0, description=f"{phase_desc} {epoch} | loss={aggregated_loss:.4f}") + + save_checkpoint(model, ckpt_file, model_state=model_state) + return total_steps + + def main(argv) -> float: """ Fine-tunes the specified model and returns the difference between initial and best validation perplexity in Torch, @@ -412,63 +492,46 @@ def main(argv) -> float: epoch_samples = num_samples - num_samples % args.microbatch_size microbatches_per_epoch = epoch_samples // args.microbatch_size - def run_training( - num_epochs: int, optimizer: torch.optim.Optimizer, phase_desc: str, start_total_steps: int - ) -> int: - loss_numerator = grad_steps = 0 - total_steps = start_total_steps - for epoch in range(num_epochs): - batch_indices_epoch = torch.randperm(num_samples)[:epoch_samples].chunk(microbatches_per_epoch) - epoch_tracker = track(batch_indices_epoch, description=f"{phase_desc} {epoch}") - for indices in epoch_tracker: - indices = indices.tolist() - - def form_batch(inputs: list[Tensor], model_input: bool): - batch = torch.cat([inputs[i] for i in indices], dim=0) - return get_model_input(batch) if model_input else batch.to(device=device, dtype=torch_dtype) - - # Compute distillation loss between logits of the original model and the model with FQ + LoRA. - inputs = form_batch(train_loader, model_input=True) - with torch.no_grad(): - targets = model.lm_head(form_batch(orig_hiddens, model_input=False)) - if hasattr(model.config, "final_logit_softcapping"): # Gemma has post-processing after lm_head - fls = model.config.final_logit_softcapping - if fls is not None: - targets = targets / fls - targets = torch.tanh(targets) - targets = targets * fls - outputs = model(**inputs).logits - loss = kl_div(outputs, targets.to(dtype=torch_dtype, device=device)) - - # Perform an optimization step after accumulating gradients over multiple minibatches. - loss_numerator += loss.item() - grad_steps += 1 - if not torch.isfinite(loss).item(): - err = f"Fine-tuning loss is {loss}" - raise ValueError(err) - (loss / grad_accumulation_steps).backward() - if grad_steps == grad_accumulation_steps: - optimizer.step() - optimizer.zero_grad() - aggregated_loss = loss_numerator / grad_steps - loss_numerator = grad_steps = 0 - total_steps += 1 - tb.add_scalar("loss", aggregated_loss, total_steps) - epoch_tracker.update( - advance=0, description=f"{phase_desc} {epoch} | loss={aggregated_loss:.4f}" - ) - - save_checkpoint(model, ckpt_file, model_state=not args.basic_init) - return total_steps - - total_steps = run_training(args.epochs, opt, phase_desc="Train epoch", start_total_steps=0) + total_steps = run_training( + model=model, + train_loader=train_loader, + orig_hiddens=orig_hiddens, + optimizer=opt, + ckpt_file=ckpt_file, + tb=tb, + num_epochs=args.epochs, + phase_desc="Train epoch", + start_total_steps=0, + device=device, + torch_dtype=torch_dtype, + grad_accumulation_steps=grad_accumulation_steps, + num_samples=num_samples, + epoch_samples=epoch_samples, + microbatches_per_epoch=microbatches_per_epoch, + model_state=not args.basic_init, + ) # Optional scales-only finetuning phase: LoRA adapters are frozen, only quantizer scales are trained. if args.scale_finetune_epochs > 0: scale_params = set_trainable(model, lora_lr=args.lr, fq_lr=fq_lr, scales_only=True) opt = torch.optim.AdamW(scale_params, weight_decay=weight_decay) run_training( - args.scale_finetune_epochs, opt, phase_desc="Scales-only epoch", start_total_steps=total_steps + model=model, + train_loader=train_loader, + orig_hiddens=orig_hiddens, + optimizer=opt, + ckpt_file=ckpt_file, + tb=tb, + num_epochs=args.scale_finetune_epochs, + phase_desc="Scales-only epoch", + start_total_steps=total_steps, + device=device, + torch_dtype=torch_dtype, + grad_accumulation_steps=grad_accumulation_steps, + num_samples=num_samples, + epoch_samples=epoch_samples, + microbatches_per_epoch=microbatches_per_epoch, + model_state=not args.basic_init, ) del model diff --git a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py index 031ae01853a..510ebe7e01b 100644 --- a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py @@ -440,7 +440,14 @@ def transform_model( model = model.model transformation_layout = TransformationLayout() - is_all_8bit = all(wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters) + + if hasattr(weight_compression_parameters, "sequence"): + # avoid track + is_all_8bit = all( + wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters.sequence + ) + else: + is_all_8bit = all(wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters) for wc_params in weight_compression_parameters: compression_config = wc_params.compression_config if compression_config.mode in [ From fc719949a9dd5e87bb1f5c1a4e370b5ea798dbb4 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 17 Sep 2026 10:52:46 +0200 Subject: [PATCH 09/27] 1)Added optional equalization for MLP layer. 2) Added optional possibility to save torch checkpoint with dequantized weights. --- .../main.py | 277 +++++++++++++++++- 1 file changed, 274 insertions(+), 3 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index e879f3fdb08..8f8dfd71dd4 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -251,6 +251,24 @@ def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path) -> OVMode ) +@torch.no_grad() +def export_to_dequantized_torch(pretrained: str, ckpt_file: Path, pt_dir: Path): + """ + Replace the quantized weights of the model with dequantized weights and save it to the specified directory. + + :param pretrained: The name or path of the pretrained model. + :param ckpt_file: The path to the checkpoint file to load the model weights and NNCF configurations. + :param pt_dir: The directory where the dequantized PyTorch model will be saved. + :return: None. The dequantized model is saved to the specified directory. + """ + tokenizer = AutoTokenizer.from_pretrained(pretrained) + model_to_eval = AutoModelForCausalLM.from_pretrained(pretrained, torch_dtype=torch.float32, device_map="cpu") + model_to_eval = load_checkpoint(model_to_eval, ckpt_file) + model_to_eval = nncf.strip(model_to_eval, do_copy=False, strip_format=StripFormat.IN_PLACE) + model_to_eval.save_pretrained(pt_dir) + tokenizer.save_pretrained(pt_dir) + + def limit_type(astr: str): value = float(astr) if value < 0 or value > 1: @@ -298,6 +316,17 @@ def get_argument_parser() -> argparse.ArgumentParser: "accuracy, but require a calibration dataset and additional initialization time " "(~20 sec for 1B and ~80 sec for 8B models).", ) + parser.add_argument( + "--equalize_mlp", + action="store_true", + help="Whether to equalize the scales of MLP layers (down_proj and up_proj/gate_proj with preceding LayerNorm) " + "before quantization. This can improve the accuracy of low-bit quantization.", + ) + parser.add_argument( + "--save_pt", + action="store_true", + help="Whether to save the model with dequantization after training. It is useful for fast evaluation.", + ) # Data params parser.add_argument("--num_train_samples", type=int, default=2048, help="Number of training samples") @@ -337,6 +366,238 @@ def get_argument_parser() -> argparse.ArgumentParser: return parser +# ---------------------------------------------------------------------- # +# MLP equalization (average up_proj/gate_proj input scale absorbed into layer norm weights) +# ---------------------------------------------------------------------- # +def _find_mlp_groups(model: nn.Module) -> list[tuple[nn.Module, nn.Linear, list[nn.Linear]]]: + """ + Collect ``(parent, down_proj, [producers])`` triples where ``producers`` + are the sibling linears whose outputs are consumed by ``down_proj`` along + its input-channel dimension. + + Recognized layouts: + * Llama-style: ``down_proj`` consumes ``up_proj`` * SiLU(``gate_proj``); + producers = ``[up_proj, gate_proj]``. + * Generic: only ``up_proj`` present -> producers = ``[up_proj]``. + """ + groups: list[tuple[nn.Module, nn.Linear, list[nn.Linear]]] = [] + for parent in model.modules(): + down = getattr(parent, "down_proj", None) + if not isinstance(down, nn.Linear): + continue + producers: list[nn.Linear] = [] + for attr in ("up_proj",): + sib = getattr(parent, attr, None) + if isinstance(sib, nn.Linear) and sib.out_features == down.in_features: + producers.append(sib) + if producers: + groups.append((parent, down, producers)) + return groups + + +def _find_up_gate_groups(model: nn.Module) -> list[tuple[nn.Module, nn.Linear, list[nn.Linear]]]: + """ + Collect ``(parent, down_proj, [producers])`` triples where ``producers`` + are the sibling linears whose outputs are consumed by ``down_proj`` along + its input-channel dimension. + + Recognized layouts: + * Llama-style: ``down_proj`` consumes ``up_proj`` * SiLU(``gate_proj``); + producers = ``[up_proj, gate_proj]``. + * Generic: only ``up_proj`` present -> producers = ``[up_proj]``. + """ + groups: list[tuple[nn.Module, nn.Linear, list[nn.Linear]]] = [] + for parent in model.modules(): + if not hasattr(parent, "mlp"): + continue + mlp = getattr(parent, "mlp") + gate = getattr(mlp, "gate_proj", None) + if not isinstance(gate, nn.Linear): + continue + + up = getattr(mlp, "up_proj", None) + if not isinstance(up, nn.Linear): + continue + + # Identify the LayerNorm whose output is the direct input to gate_proj / up_proj. + # Check Gemma3-style first (pre_feedforward_layernorm), then fall back to + # Llama/Qwen-style (post_attention_layernorm). + producer = None + for attr in ("pre_feedforward_layernorm", "post_attention_layernorm"): + sib = getattr(parent, attr, None) + if sib is not None: + producer = sib + break + + if producer: + groups.append((up, gate, producer)) + return groups + + +def _get_weight_data(module: nn.Module) -> torch.Tensor | None: + """Return the actual weight tensor for *module*, even when the parameter is on + the meta device (i.e. CPU-offloaded by accelerate's device_map='auto'). + + For meta parameters, the real data lives in the AlignDevicesHook's + ``weights_map`` dict-like object; we return a copy on CPU so arithmetic works. + Returns ``None`` if the actual data cannot be found. + """ + w = module.weight + if w.device.type != "meta": + return w.data + hook = getattr(module, "_hf_hook", None) + if hook is not None: + wm = getattr(hook, "weights_map", None) + if wm is not None: + try: + return wm["weight"].cpu() + except Exception: + pass + return None + + +def _set_weight_data(module: nn.Module, data: torch.Tensor) -> None: + """Write *data* back into *module*'s weight, even when it is on meta device. + + For non-meta parameters: in-place copy to the parameter's device. + For meta (offloaded) parameters: replace the entry in the hook's weights_map. + """ + w = module.weight + if w.device.type != "meta": + w.data.copy_(data.to(device=w.device, dtype=w.dtype)) + return + hook = getattr(module, "_hf_hook", None) + if hook is not None: + wm = getattr(hook, "weights_map", None) + if wm is not None: + try: + original = wm["weight"] + wm["weight"] = data.to(device=original.device, dtype=original.dtype) + except Exception: + pass + + +# rescale scale to [min, max] to avoid extreme values that cause instability during training or quantization +def align_scale(s: Tensor, min=0.1, max=1.0) -> Tensor: + min_s = s.min() + max_s = s.max() + if max_s - min_s < 1e-5: + return torch.clamp(s, min=min, max=max) + s = (s - min_s) / (max_s - min_s) * (max - min) + min + return s + + +@torch.no_grad() +def equalize_up_gate_with_layernorm(model: nn.Module, eps: float = 1e-5, use_align_scale: bool = True) -> int: + groups = _find_up_gate_groups(model) + if not groups: + return 0 + + n_done = 0 + for up, gate, producer in groups: + w_gate = _get_weight_data(gate) + w_up = _get_weight_data(up) + w_prod = _get_weight_data(producer) + if w_gate is None or w_up is None or w_prod is None: + continue + + s_gate = w_gate.abs().mean(dim=0).clamp_min(eps).to(dtype=w_gate.dtype) + s_up = w_up.abs().mean(dim=0).clamp_min(eps).to(dtype=w_up.dtype) + + s_gate = s_gate / s_gate.norm(p=2, dim=0, keepdim=True) + s_up = s_up / s_up.norm(p=2, dim=0, keepdim=True) + + # up_proj theoretically more sensitive to quantization + s = 0.1 * s_gate + 0.9 * s_up + if use_align_scale: + s = align_scale(s, min=0.1, max=1.0) + # Divide gate/up input columns by s. + print("Max val before equalization gate:", w_gate.abs().max().item()) + print("Max val before equalization up:", w_up.abs().max().item()) + _set_weight_data(gate, w_gate * (1.0 / s.unsqueeze(0))) + _set_weight_data(up, w_up * (1.0 / s.unsqueeze(0))) + print("Max val after equalization gate:", _get_weight_data(gate).abs().max().item()) + print("Max val after equalization up:", _get_weight_data(up).abs().max().item()) + + # Scale producer (LayerNorm/RMSNorm) so its output is multiplied by s. + s_dev = s.to(dtype=w_prod.dtype) + # Gemma3 uses (1 + weight) RMSNorm (weight initialized to zeros): + # effective multiplier = (1 + weight) + # to scale output by s: (1 + w_new) = s*(1 + w_old) → w_new = s*(1+w_old) - 1 + # Standard RMSNorm (e.g. Llama) uses weight*x (weight initialized to ones): + # effective multiplier = weight + # to scale output by s: w_new = s * w_old + if "gemma" in type(producer).__name__.lower(): + _set_weight_data(producer, s_dev * (1.0 + w_prod) - 1.0) + else: + _set_weight_data(producer, w_prod * s_dev) + if hasattr(producer, "bias") and producer.bias is not None and producer.bias.device.type != "meta": + producer.bias.data.mul_(s_dev.to(producer.bias.device, producer.bias.dtype)) + n_done += 1 + return n_done + + +@torch.no_grad() +def equalize_down_proj( + model: nn.Module, + eps: float = 1e-5, + use_align_scale: bool = True, +) -> int: + """ + Equalize each ``down_proj`` layer by absorbing the per-input-channel + activation magnitude into its producers (``up_proj`` and, when present, + ``gate_proj``). + + For every MLP block let ``s = mean(|x|, dim=batch_seq)`` measured at the + input of ``down_proj`` over the calibration set. Then: + + * ``down_proj.weight /= s[None, :]`` (divide along input channels) + * For each producer ``L`` (e.g. ``up_proj``, ``gate_proj``): + ``L.weight *= s[:, None]`` (scale output channels) + ``L.bias *= s`` (if a bias exists) + + Mathematically, ``down(up(x) * silu(gate(x))) = down((up(x)*s) * (silu(gate(x)*s)/s))`` + is *not* exact for the SiLU branch in general, but in practice this + pre-quantization equalization (cf. SmoothQuant / AWQ) significantly + flattens the weight magnitudes seen by the per-group quantizer. The + transformation is exact when no SiLU is present (``producers == [up_proj]``). + + :param model: Model whose MLP blocks expose ``down_proj`` (and optional + ``up_proj``/``gate_proj`` siblings) as direct attributes. Must be + called on plain ``nn.Linear`` layers (i.e. **before** wrapping them + with :class:`QuantizedLoraLinear`). + :param eps: Lower bound for ``s`` to avoid division by zero. + :return: Number of equalized MLP groups. + """ + groups = _find_mlp_groups(model) + if not groups: + return 0 + + n_done = 0 + for _, down, producers in groups: + w_down = _get_weight_data(down) + if w_down is None: + continue + s = w_down.abs().mean(dim=0).clamp_min(eps).to(dtype=w_down.dtype) + if use_align_scale: + s = align_scale(s, min=0.1, max=1.0) + print("Max val before equalization:", w_down.abs().max().item()) + _set_weight_data(down, w_down * (1.0 / s.unsqueeze(0))) + print("Max val after equalization:", _get_weight_data(down).abs().max().item()) + + # Scale producer output rows by s. + for prod in producers: + w_prod = _get_weight_data(prod) + if w_prod is None: + continue + s_dev = s.to(dtype=w_prod.dtype) + _set_weight_data(prod, w_prod * s_dev.unsqueeze(1)) + if prod.bias is not None and prod.bias.device.type != "meta": + prod.bias.data.mul_(s_dev.to(prod.bias.device, prod.bias.dtype)) + n_done += 1 + return n_done + + def run_training( model: nn.Module, train_loader: list[Tensor], @@ -479,6 +740,11 @@ def main(argv) -> float: if args.resume and ckpt_file.exists(): model = load_checkpoint(model, ckpt_file) else: + if args.equalize_mlp: + n_eq = equalize_down_proj(model, use_align_scale=True) + print(f"Equalized {n_eq} down_proj layers.") + n_eq = equalize_up_gate_with_layernorm(model, use_align_scale=True) + print(f"Equalized {n_eq} up_proj/gate_proj layers with preceding LayerNorm.") model = compress_weights(model, dataset=dataset, **compression_config) save_checkpoint(model, ckpt_file, model_state=not args.basic_init) fq_lr = args.lr / 10 @@ -508,7 +774,7 @@ def main(argv) -> float: num_samples=num_samples, epoch_samples=epoch_samples, microbatches_per_epoch=microbatches_per_epoch, - model_state=not args.basic_init, + model_state=not args.basic_init or args.equalize_mlp, ) # Optional scales-only finetuning phase: LoRA adapters are frozen, only quantizer scales are trained. @@ -531,16 +797,21 @@ def main(argv) -> float: num_samples=num_samples, epoch_samples=epoch_samples, microbatches_per_epoch=microbatches_per_epoch, - model_state=not args.basic_init, + model_state=not args.basic_init or args.equalize_mlp, ) del model + + if args.save_pt: + export_to_dequantized_torch(args.pretrained, ckpt_file, ckpt_file.parent / "dequantized") + print(f"The finetuned model has been exported to OpenVINO and saved to: {ckpt_file.parent / 'dequantized'}\n") + # Export the best tuned model to OpenVINO and evaluate it using LM-Evaluation-Harness. model_for_eval = export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) ov_perplexity = measure_perplexity(model_for_eval, args.eval_seqlen, args.limit) tb.add_scalar("ov_perplexity", ov_perplexity, 0) print( - f"The finetuned model has been exported to OpenVINO and saved to: {last_dir}\n" + f"The finetuned model has been exported to OpenVINO and saved to: {ckpt_file.parent}\n" f"The word perplexity on wikitext (test) = {ov_perplexity:.4f}" ) return ov_perplexity From 8f8cc3505c52594825a39524a4657692b9181d89 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 17 Sep 2026 12:01:08 +0200 Subject: [PATCH 10/27] Avoid awq in the case of weight equalization. --- .../distillation_qat_with_lora_2_3_bit/main.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 8f8dfd71dd4..7d1b5223bb2 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -62,7 +62,7 @@ def get_wikitext2(num_samples: int, seqlen: int, tokenizer: Any, device: torch.d :return: A list of tensors containing the tokenized text samples. """ traindata = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train") - limit = num_samples * seqlen // 4 # ~1k for 128 samples with seqlen=32 to be aligned with optimum + limit = num_samples * seqlen # ~1k for 128 samples with seqlen=32 to be aligned with optimum text = "".join([" \n" if s == "" else s for s in traindata["text"][:limit]]) trainenc = tokenizer(text, return_tensors="pt") trainloader = [] @@ -697,7 +697,7 @@ def main(argv) -> float: compression_config = dict( mode=CompressWeightsMode.INT3_SYM if args.bits == 3 else CompressWeightsMode.INT2_SYM, group_size=64, - awq=not args.basic_init, + awq=not (args.basic_init or args.equalize_mlp), scale_estimation=not args.basic_init, compression_format=CompressionFormat.FQ_LORA, ) @@ -758,6 +758,7 @@ def main(argv) -> float: epoch_samples = num_samples - num_samples % args.microbatch_size microbatches_per_epoch = epoch_samples // args.microbatch_size + save_model_state = not args.basic_init or args.equalize_mlp total_steps = run_training( model=model, train_loader=train_loader, @@ -774,7 +775,7 @@ def main(argv) -> float: num_samples=num_samples, epoch_samples=epoch_samples, microbatches_per_epoch=microbatches_per_epoch, - model_state=not args.basic_init or args.equalize_mlp, + model_state=save_model_state, ) # Optional scales-only finetuning phase: LoRA adapters are frozen, only quantizer scales are trained. @@ -797,14 +798,17 @@ def main(argv) -> float: num_samples=num_samples, epoch_samples=epoch_samples, microbatches_per_epoch=microbatches_per_epoch, - model_state=not args.basic_init or args.equalize_mlp, + model_state=save_model_state, ) + if args.scale_finetune_epochs == 0 and args.epochs == 0: + save_checkpoint(model, ckpt_file, model_state=save_model_state) + del model if args.save_pt: export_to_dequantized_torch(args.pretrained, ckpt_file, ckpt_file.parent / "dequantized") - print(f"The finetuned model has been exported to OpenVINO and saved to: {ckpt_file.parent / 'dequantized'}\n") + print(f"The finetuned model has been exported to torch and saved to: {ckpt_file.parent / 'dequantized'}\n") # Export the best tuned model to OpenVINO and evaluate it using LM-Evaluation-Harness. model_for_eval = export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) From 6ad2ffe246c7188d9e98f207a0b8eb8bb9029045 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 17 Sep 2026 13:29:46 +0200 Subject: [PATCH 11/27] Added pile dataset for train and PTQ. --- .../main.py | 40 ++++++++++++++++++- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 7d1b5223bb2..13314913cac 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -75,6 +75,30 @@ def get_wikitext2(num_samples: int, seqlen: int, tokenizer: Any, device: torch.d return trainloader +def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device): + ds = load_dataset("NeelNanda/pile-10k", split="train") + + trainloader = [] + text = "" + for example in ds: + text += " \n" + example["text"] + trainenc = tokenizer(text, return_tensors="pt") + if trainenc.input_ids.shape[1] < seqlen: + continue + text = "" + if trainenc.input_ids.shape[1] > seqlen + 1: + i = torch.randint(0, trainenc.input_ids.shape[1] - seqlen - 1, (1,)).item() + else: + i = 0 + j = i + seqlen + inp = trainenc.input_ids[:, i:j].to(device) + trainloader.append(inp) + if len(trainloader) >= num_samples: + break + + return trainloader + + def measure_perplexity( optimum_model: OptimizedModel, max_length: int | None = None, @@ -327,6 +351,13 @@ def get_argument_parser() -> argparse.ArgumentParser: action="store_true", help="Whether to save the model with dequantization after training. It is useful for fast evaluation.", ) + parser.add_argument( + "--dataset", + type=str, + default="pile-10k", + choices=["pile-10k", "wikitext-2"], + help="The dataset to use for training and evaluation.", + ) # Data params parser.add_argument("--num_train_samples", type=int, default=2048, help="Number of training samples") @@ -723,14 +754,19 @@ def main(argv) -> float: tokenizer = AutoTokenizer.from_pretrained(args.pretrained) # Prepare training and calibration data - train_loader = get_wikitext2( + dataset_map = { + "pile-10k": get_pile_10k, + "wikitext-2": get_wikitext2, + } + load_fn = dataset_map[args.dataset] + train_loader = load_fn( num_samples=args.num_train_samples, seqlen=args.train_seqlen, tokenizer=tokenizer, device=device ) if args.basic_init: example_input = {k: v.to(device) for k, v in model.dummy_inputs.items()} dataset = Dataset([example_input]) else: - calib_loader = get_wikitext2(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) + calib_loader = load_fn(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) dataset = Dataset(map(get_model_input, calib_loader)) # Pre-compute hiddens of teacher model for distillation loss. From 6d94fee067fa920875b9fe85b70c1e26da3c0e61 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Fri, 18 Sep 2026 16:57:17 +0200 Subject: [PATCH 12/27] Add ultracaht dataset. --- .../main.py | 60 +++++++++++++++++-- 1 file changed, 55 insertions(+), 5 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 13314913cac..6f402895ca8 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -13,6 +13,7 @@ import sys import warnings from datetime import datetime +import os from pathlib import Path from pprint import pprint from typing import Any @@ -21,6 +22,7 @@ import torch.nn.functional as F import transformers from datasets import load_dataset +import json from lm_eval import simple_evaluate from lm_eval.models.optimum_lm import OptimumLM from optimum.exporters.openvino.convert import export_from_model @@ -75,6 +77,46 @@ def get_wikitext2(num_samples: int, seqlen: int, tokenizer: Any, device: torch.d return trainloader +def warmup_triton(): + @torch.compile + def warm_up_compiler(x): + return x * 2 + + dummy_tensor = torch.randn(2, 2, device="cuda") + warm_up_compiler(dummy_tensor) + + +def get_ultrachat_200k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device) -> list[Tensor]: + if not hasattr(tokenizer, "apply_chat_template") or tokenizer.apply_chat_template is None: + raise ValueError("Tokenizer must have an 'apply_chat_template' attribute for ultra chat dataset.") + + trainloader = [] + text = "" + + dataset = load_dataset('HuggingFaceH4/ultrachat_200k', split='train_sft', streaming=True) + + for example in dataset: + text = tokenizer.apply_chat_template(example["messages"], tokenize=False, add_generation_prompt=True) + trainenc = tokenizer(text, return_tensors="pt") + + if trainenc.input_ids.shape[1] < seqlen: + continue + if trainenc.input_ids.shape[1] > seqlen + 1: + i = torch.randint(0, trainenc.input_ids.shape[1] - seqlen - 1, (1,)).item() + else: + i = 0 + j = i + seqlen + inp = trainenc.input_ids[:, i:j].to(device) + trainloader.append(inp) + if len(trainloader) >= num_samples: + break + del dataset + + return trainloader + + + + def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device): ds = load_dataset("NeelNanda/pile-10k", split="train") @@ -83,6 +125,7 @@ def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.de for example in ds: text += " \n" + example["text"] trainenc = tokenizer(text, return_tensors="pt") + if trainenc.input_ids.shape[1] < seqlen: continue text = "" @@ -355,7 +398,7 @@ def get_argument_parser() -> argparse.ArgumentParser: "--dataset", type=str, default="pile-10k", - choices=["pile-10k", "wikitext-2"], + choices=["pile-10k", "wikitext-2", "ultrachat_200k"], help="The dataset to use for training and evaluation.", ) @@ -757,6 +800,7 @@ def main(argv) -> float: dataset_map = { "pile-10k": get_pile_10k, "wikitext-2": get_wikitext2, + "ultrachat_200k": get_ultrachat_200k, } load_fn = dataset_map[args.dataset] train_loader = load_fn( @@ -766,11 +810,17 @@ def main(argv) -> float: example_input = {k: v.to(device) for k, v in model.dummy_inputs.items()} dataset = Dataset([example_input]) else: - calib_loader = load_fn(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) + if args.dataset == "ultrachat_200k": + calib_loader = train_loader[:128] + else: + calib_loader = load_fn(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) dataset = Dataset(map(get_model_input, calib_loader)) - # Pre-compute hiddens of teacher model for distillation loss. - orig_hiddens = calc_hiddens(model, train_loader) + is_ptq = args.scale_finetune_epochs == 0 and args.epochs == 0 + orig_hiddens = None + if not is_ptq: + # Pre-compute hiddens of teacher model for distillation loss. + orig_hiddens = calc_hiddens(model, train_loader) # Create or load model to tune with Fake Quantizers and absorbable LoRA adapters. if args.resume and ckpt_file.exists(): @@ -837,7 +887,7 @@ def main(argv) -> float: model_state=save_model_state, ) - if args.scale_finetune_epochs == 0 and args.epochs == 0: + if is_ptq: save_checkpoint(model, ckpt_file, model_state=save_model_state) del model From 9fdee25b60a62296e64cc52439aaebe97b4ce5a0 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Fri, 18 Sep 2026 17:01:30 +0200 Subject: [PATCH 13/27] Added lr scheduler. --- .../main.py | 63 ++++++++++++++++--- 1 file changed, 54 insertions(+), 9 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 6f402895ca8..c7cefc250db 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -13,7 +13,6 @@ import sys import warnings from datetime import datetime -import os from pathlib import Path from pprint import pprint from typing import Any @@ -22,7 +21,6 @@ import torch.nn.functional as F import transformers from datasets import load_dataset -import json from lm_eval import simple_evaluate from lm_eval.models.optimum_lm import OptimumLM from optimum.exporters.openvino.convert import export_from_model @@ -31,6 +29,8 @@ from torch import Tensor from torch import nn from torch.jit import TracerWarning +from torch.optim.lr_scheduler import LinearLR +from torch.optim.lr_scheduler import LRScheduler from torch.utils.tensorboard import SummaryWriter from transformers import AutoModelForCausalLM from transformers import AutoTokenizer @@ -88,12 +88,13 @@ def warm_up_compiler(x): def get_ultrachat_200k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device) -> list[Tensor]: if not hasattr(tokenizer, "apply_chat_template") or tokenizer.apply_chat_template is None: - raise ValueError("Tokenizer must have an 'apply_chat_template' attribute for ultra chat dataset.") - + msg = "Tokenizer must have an 'apply_chat_template' attribute for ultra chat dataset." + raise ValueError(msg) + trainloader = [] text = "" - - dataset = load_dataset('HuggingFaceH4/ultrachat_200k', split='train_sft', streaming=True) + + dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft", streaming=True) for example in dataset: text = tokenizer.apply_chat_template(example["messages"], tokenize=False, add_generation_prompt=True) @@ -113,9 +114,7 @@ def get_ultrachat_200k(num_samples: int, seqlen: int, tokenizer: Any, device: to del dataset return trainloader - - - + def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device): ds = load_dataset("NeelNanda/pile-10k", split="train") @@ -256,6 +255,31 @@ def set_trainable(model: nn.Module, lora_lr: float, fq_lr: float, scales_only: b return [{"params": adapters_to_train, "lr": lora_lr}, {"params": scales_to_train, "lr": fq_lr}] +def get_linear_lr_scheduler( + optimizer: torch.optim.Optimizer, + enabled: bool, + num_epochs: int, + microbatches_per_epoch: int, + grad_accumulation_steps: int, +) -> LinearLR | None: + """ + Creates a linear learning rate scheduler for the requested training phase. + + :param optimizer: Optimizer to schedule. + :param enabled: Whether to create the scheduler. + :param num_epochs: Number of epochs in the phase. + :param microbatches_per_epoch: Number of microbatch chunks in one epoch. + :param grad_accumulation_steps: Number of microbatches before an optimizer step. + :return: A linear learning rate scheduler or None. + """ + if not enabled: + return None + total_steps = num_epochs * (microbatches_per_epoch // grad_accumulation_steps) + if total_steps == 0: + return None + return LinearLR(optimizer, start_factor=1.0, end_factor=0.0, total_iters=total_steps) + + def save_checkpoint(model: nn.Module, ckpt_file: Path, model_state: bool = True) -> None: """ Stores the current state of a quantized model to a checkpoint file. @@ -430,6 +454,11 @@ def get_argument_parser() -> argparse.ArgumentParser: help="Number of additional epochs after the main training loop during which only the quantizer scales are " "finetuned (LoRA adapters are frozen). Set to 0 to skip this phase.", ) + parser.add_argument( + "--linear_lr_scheduler", + action="store_true", + help="Use a linear learning rate scheduler that decays the learning rate to zero during each training phase.", + ) parser.add_argument("--batch_size", type=int, default=32, help="Size of training batch.") parser.add_argument( "--microbatch_size", @@ -677,6 +706,7 @@ def run_training( train_loader: list[Tensor], orig_hiddens: list[Tensor], optimizer: torch.optim.Optimizer, + scheduler: LRScheduler | None, ckpt_file: Path, tb: SummaryWriter, num_epochs: int, @@ -698,6 +728,7 @@ def run_training( :param train_loader: Training samples used for distillation. :param orig_hiddens: Teacher hidden states computed from the original model. :param optimizer: Optimizer used for updates. + :param scheduler: Optional learning rate scheduler stepped after optimizer updates. :param ckpt_file: Path to save the checkpoint after each epoch. :param tb: TensorBoard writer for loss metrics. :param num_epochs: Number of epochs to run. @@ -746,11 +777,15 @@ def form_batch(inputs: list[Tensor], model_input: bool): (loss / grad_accumulation_steps).backward() if grad_steps == grad_accumulation_steps: optimizer.step() + if scheduler is not None: + scheduler.step() optimizer.zero_grad() aggregated_loss = loss_numerator / grad_steps loss_numerator = grad_steps = 0 total_steps += 1 tb.add_scalar("loss", aggregated_loss, total_steps) + current_lr = scheduler.get_last_lr()[0] if scheduler is not None else optimizer.param_groups[0]["lr"] + tb.add_scalar("lr", current_lr, total_steps) epoch_tracker.update(advance=0, description=f"{phase_desc} {epoch} | loss={aggregated_loss:.4f}") save_checkpoint(model, ckpt_file, model_state=model_state) @@ -850,6 +885,9 @@ def main(argv) -> float: train_loader=train_loader, orig_hiddens=orig_hiddens, optimizer=opt, + scheduler=get_linear_lr_scheduler( + opt, args.linear_lr_scheduler, args.epochs, microbatches_per_epoch, grad_accumulation_steps + ), ckpt_file=ckpt_file, tb=tb, num_epochs=args.epochs, @@ -873,6 +911,13 @@ def main(argv) -> float: train_loader=train_loader, orig_hiddens=orig_hiddens, optimizer=opt, + scheduler=get_linear_lr_scheduler( + opt, + args.linear_lr_scheduler, + args.scale_finetune_epochs, + microbatches_per_epoch, + grad_accumulation_steps, + ), ckpt_file=ckpt_file, tb=tb, num_epochs=args.scale_finetune_epochs, From f418fbbb31ce357ace325b643d194698702fe318 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Mon, 21 Sep 2026 12:53:46 +0200 Subject: [PATCH 14/27] 1) Removed lm_eval validation in the end of training. 2) Added README and requirements.txt --- .../README.md | 158 ++++++++++++++++++ .../main.py | 73 +------- .../requirements.txt | 10 ++ .../weight_compression/algorithm.py | 2 +- .../weight_compression/openvino_backend.py | 10 +- 5 files changed, 181 insertions(+), 72 deletions(-) create mode 100644 examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md create mode 100644 examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md new file mode 100644 index 00000000000..27d0a90a455 --- /dev/null +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md @@ -0,0 +1,158 @@ +# Distillation QAT with absorbable LoRA for 2/3-bit LLM compression + +This example trains a Hugging Face causal LM with NNCF weight compression in `FQ_LORA` mode. The script does not just run a static example; it performs an end-to-end compression + fine-tuning pipeline for 2-bit or 3-bit quantized weights. + +## What the script actually does + +Running `main.py` executes the following sequence: + +1. Loads a pretrained model and tokenizer from Hugging Face. +2. Builds a training dataset from one of: + - `pile-10k` + - `wikitext-2` + - `ultrachat_200k` +3. Optionally equalizes MLP weights before compression to improve low-bit robustness. +4. Applies `compress_weights(...)` with: + - `INT3_SYM` for `--bits 3` + - `INT2_SYM` for `--bits 2` + - `compression_format=CompressionFormat.FQ_LORA` +5. Enables trainable absorbable LoRA adapters and quantizer scales. +6. Runs KL-based distillation training against the original model hidden states / logits. +7. Optionally performs a second scales-only fine-tuning phase. +8. Saves the NNCF checkpoint. +9. Exports the tuned model to OpenVINO IR. +10. Optionally exports a dequantized PyTorch model as well. + +This is a real training pipeline for low-bit LLM optimization, not a simple inference example. + +## Prerequisites + +Before running this example, ensure you have following installed and set up your environment: + +- Python 3.10+ +- CUDA-enabled NVIDIA GPU +- PyTorch with CUDA support +- A working installation of the repository and dependencies + +### 1. Create and activate a virtual environment + +```bash +cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit +python3 -m venv nncf_env +source nncf_env/bin/activate # On Windows: nncf_env\Scripts\activate.bat +``` + +### 2. Install the example dependencies from this folder + +```bash +python -m pip install -r requirements.txt +``` + +### 3. Install NNCF and other dependencies + +```bash +python3 -m pip install ../../../../ -r requirements.txt +``` + +## Usage + +Run the script directly: + +```bash +cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit +python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output +``` + +The default command line is effectively: + +```bash +python main.py \ + --pretrained meta-llama/Llama-3.2-1B-Instruct \ + --bits 3 \ + --output_dir output \ + --dataset pile-10k \ + --num_train_samples 2048 \ + --train_seqlen 1024 \ + --eval_seqlen 2048 \ + --lr 1e-4 \ + --epochs 1 \ + --scale_finetune_epochs 0 \ + --batch_size 32 \ + --microbatch_size 8 \ + --lora_rank 256 +``` + +## Important CLI options + +### Model and compression + +- `--pretrained`: HF model ID or local path. +- `--bits`: `2` or `3` bits per weight. +- `--lora_rank`: rank of absorbable LoRA adapters. +- `--basic_init`: use a simpler initialization without AWQ / scale estimation. +- `--equalize_mlp`: run MLP equalization before compression. +- `--resume`: continue from a saved checkpoint instead of reinitializing. + +### Data + +- `--dataset`: one of `pile-10k`, `wikitext-2`, or `ultrachat_200k` +- `--num_train_samples`: number of training samples +- `--train_seqlen`: context length for training samples +- `--eval_seqlen`: context length for evaluation +- `--limit`: fraction of validation examples to use, in `[0, 1]` + +### Training + +- `--lr`: base learning rate +- `--epochs`: number of main training epochs +- `--scale_finetune_epochs`: extra epochs where only quantizer scales are trained +- `--linear_lr_scheduler`: use a linear learning-rate decay schedule +- `--batch_size`: accumulation target +- `--microbatch_size`: per-step microbatch size + +### Output/export + +- `--output_dir`: root directory for logs, checkpoints, and exported models +- `--save_pt`: export a dequantized PyTorch model after training + +## Output artifacts + +The script writes artifacts under the chosen output directory. For example, with `--output_dir output` and `--bits 3`, it creates: + +```text +output/ + tb/ + YYYY-MM-DD__HH-MM-SS/ + last_bits_3/ + nncf_checkpoint.pth + last_bits_3/ + model.xml + model.bin + ... +``` + +The checkpoint contains the NNCF state and, when appropriate, the model state. A TensorBoard directory is created for loss and LR tracking. When `--save_pt` is enabled, the script also exports a dequantized PyTorch model under: + +```text +output/last_bits_3/dequantized/ +``` + +## Training behavior + +The optimization loop is distillation-based: + +- teacher hidden states are precomputed with the original model +- the compressed model receives the same tokens +- the script minimizes a KL divergence between the student and teacher outputs / hidden states +- optimizer updates are accumulated across microbatches and applied when `batch_size // microbatch_size` is reached + +This makes the example a low-bit QAT + LoRA distillation workflow rather than a plain quantization pass. + +## Notes + +- The script asserts `torch.cuda.is_available()`, so it expects a CUDA-capable environment. +- Default model is `meta-llama/Llama-3.2-1B-Instruct`. +- OpenVINO export is always performed at the end of the run, after checkpoint restoration and stripping. +- `--resume` reuses the checkpoint if present; otherwise the script initializes from scratch. + +For more background on absorbable LoRA and low-bit training-time compression, see the project documentation for QAT LoRA usage and NNCF compression flows. diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index c7cefc250db..2b92b4a62c6 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -21,11 +21,8 @@ import torch.nn.functional as F import transformers from datasets import load_dataset -from lm_eval import simple_evaluate -from lm_eval.models.optimum_lm import OptimumLM from optimum.exporters.openvino.convert import export_from_model from optimum.intel.openvino import OVModelForCausalLM -from optimum.modeling_base import OptimizedModel from torch import Tensor from torch import nn from torch.jit import TracerWarning @@ -77,15 +74,6 @@ def get_wikitext2(num_samples: int, seqlen: int, tokenizer: Any, device: torch.d return trainloader -def warmup_triton(): - @torch.compile - def warm_up_compiler(x): - return x * 2 - - dummy_tensor = torch.randn(2, 2, device="cuda") - warm_up_compiler(dummy_tensor) - - def get_ultrachat_200k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device) -> list[Tensor]: if not hasattr(tokenizer, "apply_chat_template") or tokenizer.apply_chat_template is None: msg = "Tokenizer must have an 'apply_chat_template' attribute for ultra chat dataset." @@ -116,7 +104,7 @@ def get_ultrachat_200k(num_samples: int, seqlen: int, tokenizer: Any, device: to return trainloader -def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device): +def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.device) -> list[Tensor]: ds = load_dataset("NeelNanda/pile-10k", split="train") trainloader = [] @@ -141,27 +129,6 @@ def get_pile_10k(num_samples: int, seqlen: int, tokenizer: Any, device: torch.de return trainloader -def measure_perplexity( - optimum_model: OptimizedModel, - max_length: int | None = None, - limit: int | float | None = None, -) -> float: - """ - Measure perplexity on the Wikitext dataset, via rolling loglikelihoods for a given model. - - :param optimum_model: A model to be evaluated. - :param max_length: The maximum sequence length for evaluation. - :param limit: Limit the number of examples per task (only use this for testing). - If <1, limit is a percentage of the total number of examples. - :return: The similarity score as a float. - """ - task = "wikitext" - print("#" * 50 + " Evaluate via lm-eval-harness " + "#" * 50) - lm_obj = OptimumLM(pretrained=optimum_model, max_length=max_length) - results = simple_evaluate(lm_obj, tasks=[task], limit=limit, log_samples=False) - return results["results"][task]["word_perplexity,none"] - - @torch.no_grad() def calc_hiddens(model: nn.Module, dataloader: list[Tensor]) -> list[Tensor]: """ @@ -316,9 +283,9 @@ def load_checkpoint(model: nn.Module, ckpt_file: Path) -> nn.Module: @torch.no_grad() -def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path) -> OVModelForCausalLM: +def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path): """ - Create a wrapper of OpenVINO model from the checkpoint for evaluation on CPU via WWB. + Export the quantized model to OpenVINO IR format. :param pretrained: The name or path of the pretrained model. :param ckpt_file: The path to the checkpoint file to load the model weights and NNCF configurations. @@ -334,13 +301,6 @@ def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path) -> OVMode model_to_eval.model = repack_weights(model_to_eval.model) model_to_eval.save_pretrained(ir_dir / "repacked") - return OVModelForCausalLM.from_pretrained( - model_id=ir_dir / "repacked", - trust_remote_code=True, - load_in_8bit=False, - compile=True, - ) - @torch.no_grad() def export_to_dequantized_torch(pretrained: str, ckpt_file: Path, pt_dir: Path): @@ -360,14 +320,6 @@ def export_to_dequantized_torch(pretrained: str, ckpt_file: Path, pt_dir: Path): tokenizer.save_pretrained(pt_dir) -def limit_type(astr: str): - value = float(astr) - if value < 0 or value > 1: - msg = "value not in range [0,1]" - raise argparse.ArgumentTypeError(msg) - return value - - def get_argument_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(add_help=True) @@ -429,14 +381,6 @@ def get_argument_parser() -> argparse.ArgumentParser: # Data params parser.add_argument("--num_train_samples", type=int, default=2048, help="Number of training samples") parser.add_argument("--train_seqlen", type=int, default=1024, help="Train data context length.") - parser.add_argument("--eval_seqlen", type=int, default=2048, help="Evaluation data context length.") - parser.add_argument( - "--limit", - type=limit_type, - default=None, - help="A percentage of the total number of examples for evaluation. " - "Should be on the range [0,1]. If None, all samples will be used.", - ) # Training params parser.add_argument( @@ -941,15 +885,8 @@ def main(argv) -> float: export_to_dequantized_torch(args.pretrained, ckpt_file, ckpt_file.parent / "dequantized") print(f"The finetuned model has been exported to torch and saved to: {ckpt_file.parent / 'dequantized'}\n") - # Export the best tuned model to OpenVINO and evaluate it using LM-Evaluation-Harness. - model_for_eval = export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) - ov_perplexity = measure_perplexity(model_for_eval, args.eval_seqlen, args.limit) - tb.add_scalar("ov_perplexity", ov_perplexity, 0) - print( - f"The finetuned model has been exported to OpenVINO and saved to: {ckpt_file.parent}\n" - f"The word perplexity on wikitext (test) = {ov_perplexity:.4f}" - ) - return ov_perplexity + # Export the best tuned model to OpenVINO. + export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) if __name__ == "__main__": diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt new file mode 100644 index 00000000000..17e2e4ad117 --- /dev/null +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt @@ -0,0 +1,10 @@ +tensorboard==2.20.0 +torch==2.13.0 +torchao==0.18.0 +numpy>=1.23.5,<2 +openvino==2026.3.1 +optimum-intel==2.1.0 +optimum==2.3.0 +transformers==5.5.0 +datasets==5.0.1 +accelerate==1.15.0 diff --git a/src/nncf/quantization/algorithms/weight_compression/algorithm.py b/src/nncf/quantization/algorithms/weight_compression/algorithm.py index d8fc424875a..b8ce1766869 100644 --- a/src/nncf/quantization/algorithms/weight_compression/algorithm.py +++ b/src/nncf/quantization/algorithms/weight_compression/algorithm.py @@ -1431,7 +1431,7 @@ def repack_weights(model: TModel, graph: NNCFGraph) -> TModel: Repacks compressed weight constants to a lower-bit representation when possible. Iterates over all nodes in the graph and attempts to replace compressed constants - (i8, u8, i4, u4) with a more compact symmetric representation (i2, i3, i6) if the + (i8, u8, i4, u4) with a more compact symmetric representation (i2, i3) if the actual value range allows it. :param model: Backend-specific model with compressed weights. diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index ae7f144514e..4da6c642c36 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -118,7 +118,11 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: return False ov_node = self.name_to_node_mapping[node.node_name] - bits = 8 if ov_node.tensor_view.element_type in [ov.Type.i8, ov.Type.u8] else 4 + + supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8} + if ov_node.tensor_view.element_type not in supported_src_types: + return False + bits = supported_src_types[ov_node.tensor_view.element_type] asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4] # only repack for sym types are supported for now: i3, i2 @@ -136,14 +140,14 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: max_val = max(int(weight_tensor.max().item()), abs(int(weight_tensor.min().item()))) optimal_bits = int(max_val).bit_length() - supported_optimal_bits = [2, 3] + supported_optimal_bits = {2: ov.Type.u2, 3: ov.Type.u3} # If the optimal bits is equal to the current bits, no need to repack if optimal_bits == bits or optimal_bits not in supported_optimal_bits: return False const_node_name = node.node_name - compression_dtype = ov.Type.u2 if optimal_bits == 2 else ov.Type.u3 + compression_dtype = supported_optimal_bits[optimal_bits] # Repack the weight tensor to the optimal bits and change constant node with new # representation [u2/u3] - zero_point per tensor From 4bb71c78939015ecd58add2d4160e1af29e0c0e3 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Mon, 21 Sep 2026 14:59:18 +0200 Subject: [PATCH 15/27] 1) Added new words to spell check dict. 2) Removed debug prints. --- .ci/cspell_dict.txt | 4 ++ .../README.md | 9 ++--- .../main.py | 40 +++++-------------- .../weight_compression/algorithm.py | 13 +++++- .../weight_compression/openvino_backend.py | 19 +++++---- 5 files changed, 38 insertions(+), 47 deletions(-) diff --git a/.ci/cspell_dict.txt b/.ci/cspell_dict.txt index c8f5d99180b..0c695842814 100644 --- a/.ci/cspell_dict.txt +++ b/.ci/cspell_dict.txt @@ -231,6 +231,7 @@ logicalnot logicalor logicalxor logit +logits loglikelihoods lspec lstmsequence @@ -253,6 +254,7 @@ metatype metatypes metzen mhsa +microbatch microscaling minfo minfos @@ -425,6 +427,7 @@ segm selu senina sess +seqlen sgraph shamporov shapeof @@ -479,6 +482,7 @@ togethercomputer torchdynamo truncnorm tupled +ultrachat unet unflatten unpatch diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md index 27d0a90a455..4f7be2aa938 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md @@ -73,10 +73,9 @@ python main.py \ --dataset pile-10k \ --num_train_samples 2048 \ --train_seqlen 1024 \ - --eval_seqlen 2048 \ --lr 1e-4 \ --epochs 1 \ - --scale_finetune_epochs 0 \ + --scale_epochs 0 \ --batch_size 32 \ --microbatch_size 8 \ --lora_rank 256 @@ -98,14 +97,12 @@ python main.py \ - `--dataset`: one of `pile-10k`, `wikitext-2`, or `ultrachat_200k` - `--num_train_samples`: number of training samples - `--train_seqlen`: context length for training samples -- `--eval_seqlen`: context length for evaluation -- `--limit`: fraction of validation examples to use, in `[0, 1]` ### Training - `--lr`: base learning rate - `--epochs`: number of main training epochs -- `--scale_finetune_epochs`: extra epochs where only quantizer scales are trained +- `--scale_epochs`: extra epochs where only quantizer scales are trained - `--linear_lr_scheduler`: use a linear learning-rate decay schedule - `--batch_size`: accumulation target - `--microbatch_size`: per-step microbatch size @@ -144,7 +141,7 @@ The optimization loop is distillation-based: - teacher hidden states are precomputed with the original model - the compressed model receives the same tokens - the script minimizes a KL divergence between the student and teacher outputs / hidden states -- optimizer updates are accumulated across microbatches and applied when `batch_size // microbatch_size` is reached +- optimizer updates are accumulated across microbatch steps and applied when `batch_size // microbatch_size` is reached This makes the example a low-bit QAT + LoRA distillation workflow rather than a plain quantization pass. diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index 2b92b4a62c6..da1dae9d97f 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -296,10 +296,12 @@ def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path): model_to_eval = load_checkpoint(model_to_eval, ckpt_file) model_to_eval = nncf.strip(model_to_eval, do_copy=False, strip_format=StripFormat.DQ) export_from_model(model_to_eval, ir_dir, device="cpu") + print(f"The OpenVINO model has been exported and saved to: {ir_dir}") model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) model_to_eval.model = repack_weights(model_to_eval.model) model_to_eval.save_pretrained(ir_dir / "repacked") + print(f"The OpenVINO model has been repacked and saved to: {ir_dir / 'repacked'}") @torch.no_grad() @@ -392,7 +394,7 @@ def get_argument_parser() -> argparse.ArgumentParser: ) parser.add_argument("--epochs", type=int, default=1, help="Number of epochs.") parser.add_argument( - "--scale_finetune_epochs", + "--scale_epochs", type=int, default=0, help="Number of additional epochs after the main training loop during which only the quantizer scales are " @@ -492,14 +494,6 @@ def _get_weight_data(module: nn.Module) -> torch.Tensor | None: w = module.weight if w.device.type != "meta": return w.data - hook = getattr(module, "_hf_hook", None) - if hook is not None: - wm = getattr(hook, "weights_map", None) - if wm is not None: - try: - return wm["weight"].cpu() - except Exception: - pass return None @@ -512,16 +506,6 @@ def _set_weight_data(module: nn.Module, data: torch.Tensor) -> None: w = module.weight if w.device.type != "meta": w.data.copy_(data.to(device=w.device, dtype=w.dtype)) - return - hook = getattr(module, "_hf_hook", None) - if hook is not None: - wm = getattr(hook, "weights_map", None) - if wm is not None: - try: - original = wm["weight"] - wm["weight"] = data.to(device=original.device, dtype=original.dtype) - except Exception: - pass # rescale scale to [min, max] to avoid extreme values that cause instability during training or quantization @@ -559,12 +543,8 @@ def equalize_up_gate_with_layernorm(model: nn.Module, eps: float = 1e-5, use_ali if use_align_scale: s = align_scale(s, min=0.1, max=1.0) # Divide gate/up input columns by s. - print("Max val before equalization gate:", w_gate.abs().max().item()) - print("Max val before equalization up:", w_up.abs().max().item()) _set_weight_data(gate, w_gate * (1.0 / s.unsqueeze(0))) _set_weight_data(up, w_up * (1.0 / s.unsqueeze(0))) - print("Max val after equalization gate:", _get_weight_data(gate).abs().max().item()) - print("Max val after equalization up:", _get_weight_data(up).abs().max().item()) # Scale producer (LayerNorm/RMSNorm) so its output is multiplied by s. s_dev = s.to(dtype=w_prod.dtype) @@ -628,9 +608,7 @@ def equalize_down_proj( s = w_down.abs().mean(dim=0).clamp_min(eps).to(dtype=w_down.dtype) if use_align_scale: s = align_scale(s, min=0.1, max=1.0) - print("Max val before equalization:", w_down.abs().max().item()) _set_weight_data(down, w_down * (1.0 / s.unsqueeze(0))) - print("Max val after equalization:", _get_weight_data(down).abs().max().item()) # Scale producer output rows by s. for prod in producers: @@ -790,12 +768,12 @@ def main(argv) -> float: dataset = Dataset([example_input]) else: if args.dataset == "ultrachat_200k": - calib_loader = train_loader[:128] + calib_loader = train_loader[:129] else: - calib_loader = load_fn(num_samples=128, seqlen=128, tokenizer=tokenizer, device=device) + calib_loader = load_fn(num_samples=129, seqlen=128, tokenizer=tokenizer, device=device) dataset = Dataset(map(get_model_input, calib_loader)) - is_ptq = args.scale_finetune_epochs == 0 and args.epochs == 0 + is_ptq = args.scale_epochs == 0 and args.epochs == 0 orig_hiddens = None if not is_ptq: # Pre-compute hiddens of teacher model for distillation loss. @@ -847,7 +825,7 @@ def main(argv) -> float: ) # Optional scales-only finetuning phase: LoRA adapters are frozen, only quantizer scales are trained. - if args.scale_finetune_epochs > 0: + if args.scale_epochs > 0: scale_params = set_trainable(model, lora_lr=args.lr, fq_lr=fq_lr, scales_only=True) opt = torch.optim.AdamW(scale_params, weight_decay=weight_decay) run_training( @@ -858,13 +836,13 @@ def main(argv) -> float: scheduler=get_linear_lr_scheduler( opt, args.linear_lr_scheduler, - args.scale_finetune_epochs, + args.scale_epochs, microbatches_per_epoch, grad_accumulation_steps, ), ckpt_file=ckpt_file, tb=tb, - num_epochs=args.scale_finetune_epochs, + num_epochs=args.scale_epochs, phase_desc="Scales-only epoch", start_total_steps=total_steps, device=device, diff --git a/src/nncf/quantization/algorithms/weight_compression/algorithm.py b/src/nncf/quantization/algorithms/weight_compression/algorithm.py index b8ce1766869..e24ce0a040f 100644 --- a/src/nncf/quantization/algorithms/weight_compression/algorithm.py +++ b/src/nncf/quantization/algorithms/weight_compression/algorithm.py @@ -1447,8 +1447,17 @@ def repack_weights(model: TModel, graph: NNCFGraph) -> TModel: backend_entity = OVWeightCompressionAlgoBackend(model) + n_repacked_tensors = 0 + bits_before = 0 + bits_after = 0 + for node in graph.topological_sort(): - is_repacked = backend_entity.try_repack(node, graph) + is_repacked, bits_before_node, bits_after_node = backend_entity.try_repack(node, graph) if is_repacked: - print(f"Repacked node: {node.node_name}") + n_repacked_tensors += 1 + bits_before += bits_before_node + bits_after += bits_after_node + print(f"Total repacked tensors: {n_repacked_tensors}") + print(f"Total MB before: {bits_before / 8 / 1024 / 1024}") + print(f"Total MB after: {bits_after / 8 / 1024 / 1024}") return model diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index 4da6c642c36..a68599d852f 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -109,25 +109,25 @@ def is_compressed_constant(self, node: NNCFNode) -> bool: res = ov_node.tensor_view.element_type in [ov.Type.i8, ov.Type.u8, ov.Type.i4, ov.Type.u4] return res - def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: + def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> tuple[bool, int, int]: if not self.is_compressed_constant(node): - return False + return False, 0, 0 next_nodes = graph.get_next_nodes(node) if len(next_nodes) != 1 or next_nodes[0].metatype != om.OVConvertMetatype: - return False + return False, 0, 0 ov_node = self.name_to_node_mapping[node.node_name] supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8} if ov_node.tensor_view.element_type not in supported_src_types: - return False + return False, 0, 0 bits = supported_src_types[ov_node.tensor_view.element_type] asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4] # only repack for sym types are supported for now: i3, i2 if asym: - return False + return False, 0, 0 weight_tensor = Tensor(get_const_value_as_ov_tensor(ov_node)) weight_tensor = weight_tensor.as_numpy_tensor().astype(TensorDataType.int32) @@ -135,7 +135,7 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: if len(shape) == 1 or shape[-1] == 1: # Probably zero point - return False + return False, 0, 0 max_val = max(int(weight_tensor.max().item()), abs(int(weight_tensor.min().item()))) optimal_bits = int(max_val).bit_length() @@ -144,7 +144,7 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: # If the optimal bits is equal to the current bits, no need to repack if optimal_bits == bits or optimal_bits not in supported_optimal_bits: - return False + return False, 0, 0 const_node_name = node.node_name compression_dtype = supported_optimal_bits[optimal_bits] @@ -167,7 +167,10 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> bool: ov_node = self.name_to_node_mapping[next_nodes[0].node_name] self._replace_node(ov_node, converted_const) - return True + bits_before = bits * weight_tensor.size + bits_after = optimal_bits * weight_tensor.size + + return True, bits_before, bits_after @staticmethod def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: NNCFGraph) -> tuple[int] | None: From 547582c6792b1461299b12fad2a7b60ca7d3ffb8 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Fri, 25 Sep 2026 09:59:23 +0200 Subject: [PATCH 16/27] 1) Fixed logging of repack_weights. 2) Added arg for deterministic cuda execution. This arg is optional because it can slow down the training. --- .../torch/distillation_qat_with_lora_2_3_bit/main.py | 10 +++++++++- .../algorithms/weight_compression/algorithm.py | 8 +++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py index da1dae9d97f..4eb76635632 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py @@ -412,6 +412,11 @@ def get_argument_parser() -> argparse.ArgumentParser: default=8, help="Size of each training microbatch. Gradients will be accumulated until the batch size is reached.", ) + parser.add_argument( + "--full_determinism", + action="store_true", + help="Use deterministic CUDA kernels for reproducible QAT. This can reduce training throughput.", + ) return parser @@ -722,7 +727,10 @@ def main(argv) -> float: parser = get_argument_parser() args = parser.parse_args(argv) assert torch.cuda.is_available() - transformers.set_seed(42) + if args.full_determinism: + transformers.enable_full_determinism(42) + else: + transformers.set_seed(42) device = "cuda" torch_dtype = torch.bfloat16 compression_config = dict( diff --git a/src/nncf/quantization/algorithms/weight_compression/algorithm.py b/src/nncf/quantization/algorithms/weight_compression/algorithm.py index e24ce0a040f..a8ae0a4fdee 100644 --- a/src/nncf/quantization/algorithms/weight_compression/algorithm.py +++ b/src/nncf/quantization/algorithms/weight_compression/algorithm.py @@ -1457,7 +1457,9 @@ def repack_weights(model: TModel, graph: NNCFGraph) -> TModel: n_repacked_tensors += 1 bits_before += bits_before_node bits_after += bits_after_node - print(f"Total repacked tensors: {n_repacked_tensors}") - print(f"Total MB before: {bits_before / 8 / 1024 / 1024}") - print(f"Total MB after: {bits_after / 8 / 1024 / 1024}") + nncf_logger.info( + f"\nTotal repacked tensors: {n_repacked_tensors}; " + f"\ntotal MB before: {bits_before / 8 / 1024 / 1024}; " + f"\ntotal MB after: {bits_after / 8 / 1024 / 1024}" + ) return model From 8ec441ae2e39e380c9bb8cb4fc6d9211c5478373 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 30 Sep 2026 15:46:51 +0200 Subject: [PATCH 17/27] Rename QAT to QAD. --- .../README.md | 2 +- .../main.py | 0 .../requirements.txt | 2 +- .../algorithms/weight_compression/scale_estimation.py | 1 + 4 files changed, 3 insertions(+), 2 deletions(-) rename examples/llm_compression/torch/{distillation_qat_with_lora_2_3_bit => distillation_qad_with_lora_2_3_bit}/README.md (98%) rename examples/llm_compression/torch/{distillation_qat_with_lora_2_3_bit => distillation_qad_with_lora_2_3_bit}/main.py (100%) rename examples/llm_compression/torch/{distillation_qat_with_lora_2_3_bit => distillation_qad_with_lora_2_3_bit}/requirements.txt (90%) diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md similarity index 98% rename from examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md rename to examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md index 4f7be2aa938..55be4100720 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md @@ -1,4 +1,4 @@ -# Distillation QAT with absorbable LoRA for 2/3-bit LLM compression +# Quantization Aware Discillation with absorbable LoRA for 2/3-bit LLM compression This example trains a Hugging Face causal LM with NNCF weight compression in `FQ_LORA` mode. The script does not just run a static example; it performs an end-to-end compression + fine-tuning pipeline for 2-bit or 3-bit quantized weights. diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py similarity index 100% rename from examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py rename to examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py diff --git a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt similarity index 90% rename from examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt rename to examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt index 17e2e4ad117..987d598ddd1 100644 --- a/examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt @@ -1,7 +1,7 @@ tensorboard==2.20.0 torch==2.13.0 torchao==0.18.0 -numpy>=1.23.5,<2 +numpy==2.2.6 openvino==2026.3.1 optimum-intel==2.1.0 optimum==2.3.0 diff --git a/src/nncf/quantization/algorithms/weight_compression/scale_estimation.py b/src/nncf/quantization/algorithms/weight_compression/scale_estimation.py index 84af6dec42b..1803f243188 100644 --- a/src/nncf/quantization/algorithms/weight_compression/scale_estimation.py +++ b/src/nncf/quantization/algorithms/weight_compression/scale_estimation.py @@ -411,4 +411,5 @@ def estimate_scales(weight: Tensor, target: Tensor, zero_mask: Tensor, importanc ideal_scale = fns.abs(weight) / (fns.abs(target) + zero_mask) weighted_scale = ideal_scale * importance near_to_ideal_scale = fns.sum(weighted_scale, axis=-1, keepdims=True) + return near_to_ideal_scale From 03697553d294f21cfec623b496d5ae980c49e8b0 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 30 Sep 2026 16:19:58 +0200 Subject: [PATCH 18/27] Fixed README. --- .../README.md | 48 ++++--------------- 1 file changed, 8 insertions(+), 40 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md index 55be4100720..faf0584098f 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/README.md @@ -1,4 +1,4 @@ -# Quantization Aware Discillation with absorbable LoRA for 2/3-bit LLM compression +# Quantization Aware Distillation with absorbable LoRA for 2/3-bit LLM compression This example trains a Hugging Face causal LM with NNCF weight compression in `FQ_LORA` mode. The script does not just run a static example; it performs an end-to-end compression + fine-tuning pipeline for 2-bit or 3-bit quantized weights. @@ -27,58 +27,27 @@ This is a real training pipeline for low-bit LLM optimization, not a simple infe ## Prerequisites -Before running this example, ensure you have following installed and set up your environment: - -- Python 3.10+ -- CUDA-enabled NVIDIA GPU -- PyTorch with CUDA support -- A working installation of the repository and dependencies +Before running this example, ensure you have Python 3.10+ installed and set up your environment: ### 1. Create and activate a virtual environment ```bash -cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit python3 -m venv nncf_env source nncf_env/bin/activate # On Windows: nncf_env\Scripts\activate.bat ``` -### 2. Install the example dependencies from this folder - -```bash -python -m pip install -r requirements.txt -``` - -### 3. Install NNCF and other dependencies +### 2. Install NNCF and other dependencies ```bash python3 -m pip install ../../../../ -r requirements.txt ``` -## Usage +## Run Example -Run the script directly: +The example is fully automated. Just run the following command in the prepared Python environment: ```bash -cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit -python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output -``` - -The default command line is effectively: - -```bash -python main.py \ - --pretrained meta-llama/Llama-3.2-1B-Instruct \ - --bits 3 \ - --output_dir output \ - --dataset pile-10k \ - --num_train_samples 2048 \ - --train_seqlen 1024 \ - --lr 1e-4 \ - --epochs 1 \ - --scale_epochs 0 \ - --batch_size 32 \ - --microbatch_size 8 \ - --lora_rank 256 +python main.py ``` ## Important CLI options @@ -106,6 +75,7 @@ python main.py \ - `--linear_lr_scheduler`: use a linear learning-rate decay schedule - `--batch_size`: accumulation target - `--microbatch_size`: per-step microbatch size +- `--full_determinism`: to enable deterministic QAD ### Output/export @@ -143,7 +113,7 @@ The optimization loop is distillation-based: - the script minimizes a KL divergence between the student and teacher outputs / hidden states - optimizer updates are accumulated across microbatch steps and applied when `batch_size // microbatch_size` is reached -This makes the example a low-bit QAT + LoRA distillation workflow rather than a plain quantization pass. +This makes the example a low-bit QAD + LoRA distillation workflow rather than a plain quantization pass. ## Notes @@ -151,5 +121,3 @@ This makes the example a low-bit QAT + LoRA distillation workflow rather than a - Default model is `meta-llama/Llama-3.2-1B-Instruct`. - OpenVINO export is always performed at the end of the run, after checkpoint restoration and stripping. - `--resume` reuses the checkpoint if present; otherwise the script initializes from scratch. - -For more background on absorbable LoRA and low-bit training-time compression, see the project documentation for QAT LoRA usage and NNCF compression flows. From df4691df31972cb9deedb1c76e096956a0410a0d Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 1 Oct 2026 12:01:25 +0200 Subject: [PATCH 19/27] Applied suggestions. --- .../algorithms/weight_compression/openvino_backend.py | 9 ++------- .../algorithms/weight_compression/torch_backend.py | 11 ++++------- src/nncf/torch/function_hook/strip.py | 2 -- 3 files changed, 6 insertions(+), 16 deletions(-) diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index a68599d852f..25556df1d5e 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -119,15 +119,10 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> tuple[bool, int, int]: ov_node = self.name_to_node_mapping[node.node_name] - supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8} + supported_src_types = {ov.Type.i4: 4, ov.Type.i8: 8} if ov_node.tensor_view.element_type not in supported_src_types: return False, 0, 0 - bits = supported_src_types[ov_node.tensor_view.element_type] - asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4] - - # only repack for sym types are supported for now: i3, i2 - if asym: - return False, 0, 0 + bits = ov_node.tensor_view.element_type.get_bitwidth() weight_tensor = Tensor(get_const_value_as_ov_tensor(ov_node)) weight_tensor = weight_tensor.as_numpy_tensor().astype(TensorDataType.int32) diff --git a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py index 510ebe7e01b..9b694dd7593 100644 --- a/src/nncf/quantization/algorithms/weight_compression/torch_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/torch_backend.py @@ -441,13 +441,10 @@ def transform_model( transformation_layout = TransformationLayout() - if hasattr(weight_compression_parameters, "sequence"): - # avoid track - is_all_8bit = all( - wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters.sequence - ) - else: - is_all_8bit = all(wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters) + is_all_8bit = all( + wc_params.compression_config.num_bits == 8 + for wc_params in getattr(weight_compression_parameters, "sequence", weight_compression_parameters) + ) for wc_params in weight_compression_parameters: compression_config = wc_params.compression_config if compression_config.mode in [ diff --git a/src/nncf/torch/function_hook/strip.py b/src/nncf/torch/function_hook/strip.py index 31c87421561..c176a23df34 100644 --- a/src/nncf/torch/function_hook/strip.py +++ b/src/nncf/torch/function_hook/strip.py @@ -126,8 +126,6 @@ def replace_quantizer_to_compressed_weight_with_decompressor(model: TModel) -> T msg = "" if hook_module._qspec.half_range or hook_module._qspec.narrow_range: msg += "Unexpected parameters of quantizers on strip: half_range and narrow_range should be False.\n" - if hook_module.num_bits not in [2, 3, 4, 8]: - msg += f"Unsupported number of bits {hook_module.num_bits} for the quantizer {hook_module}.\n" if msg: raise nncf.ValidationError(msg) From 8c85da3e13ebb526c468fe65b0837db9fb60994f Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 1 Oct 2026 12:29:16 +0200 Subject: [PATCH 20/27] Removed unused comment. --- .../torch/distillation_qad_with_lora_2_3_bit/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py index 4eb76635632..2b83fdc00c5 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py @@ -329,7 +329,6 @@ def get_argument_parser() -> argparse.ArgumentParser: parser.add_argument( "--pretrained", type=str, - # default="Qwen/Qwen3-8B", default="meta-llama/Llama-3.2-1B-Instruct", help="The model id or path of a pretrained HF model configuration.", ) From 2a0130dcab6b952abfedc754eaa97eb26245cfee Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Thu, 1 Oct 2026 15:38:03 +0200 Subject: [PATCH 21/27] Update optimum-intel version. --- .../torch/distillation_qad_with_lora_2_3_bit/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt index 987d598ddd1..361b5cfb000 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/requirements.txt @@ -3,7 +3,7 @@ torch==2.13.0 torchao==0.18.0 numpy==2.2.6 openvino==2026.3.1 -optimum-intel==2.1.0 +optimum-intel==2.2.0 optimum==2.3.0 transformers==5.5.0 datasets==5.0.1 From b414e897bc3b4bf4f48c5253bd996e6849d596d5 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Fri, 2 Oct 2026 10:16:19 +0200 Subject: [PATCH 22/27] 1) Added args check. 2) Fixed convert type after repacking. 3) Added tokenizer to export. --- .../main.py | 19 ++++++++++++++++--- .../weight_compression/openvino_backend.py | 6 ++++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py index 2b83fdc00c5..2f1ec94b513 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py @@ -283,24 +283,26 @@ def load_checkpoint(model: nn.Module, ckpt_file: Path) -> nn.Module: @torch.no_grad() -def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path): +def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path, tokenizer: AutoTokenizer): """ Export the quantized model to OpenVINO IR format. :param pretrained: The name or path of the pretrained model. :param ckpt_file: The path to the checkpoint file to load the model weights and NNCF configurations. :param ir_dir: The directory where the OpenVINO model will be saved. - :return: A wrapper of OpenVINO model ready for evaluation. + :return: None. The OpenVINO model is exported and saved to the specified directory. """ model_to_eval = AutoModelForCausalLM.from_pretrained(pretrained, torch_dtype=torch.float32, device_map="cpu") model_to_eval = load_checkpoint(model_to_eval, ckpt_file) model_to_eval = nncf.strip(model_to_eval, do_copy=False, strip_format=StripFormat.DQ) export_from_model(model_to_eval, ir_dir, device="cpu") print(f"The OpenVINO model has been exported and saved to: {ir_dir}") + tokenizer.save_pretrained(ir_dir) model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) model_to_eval.model = repack_weights(model_to_eval.model) model_to_eval.save_pretrained(ir_dir / "repacked") + tokenizer.save_pretrained(ir_dir / "repacked") print(f"The OpenVINO model has been repacked and saved to: {ir_dir / 'repacked'}") @@ -718,6 +720,15 @@ def form_batch(inputs: list[Tensor], model_input: bool): return total_steps +def check_args(args): + if args.lora_rank <= 0: + msg = f"Invalid LoRA rank: {args.lora_rank}. LoRA rank must be a positive integer." + raise ValueError(msg) + if args.microbatch_size <= 0 or args.batch_size < args.microbatch_size or args.batch_size % args.microbatch_size: + msg = "batch_size must be a positive multiple of microbatch_size" + raise ValueError(msg) + + def main(argv) -> float: """ Fine-tunes the specified model and returns the difference between initial and best validation perplexity in Torch, @@ -725,6 +736,8 @@ def main(argv) -> float: """ parser = get_argument_parser() args = parser.parse_args(argv) + check_args(args) + assert torch.cuda.is_available() if args.full_determinism: transformers.enable_full_determinism(42) @@ -871,7 +884,7 @@ def main(argv) -> float: print(f"The finetuned model has been exported to torch and saved to: {ckpt_file.parent / 'dequantized'}\n") # Export the best tuned model to OpenVINO. - export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent) + export_to_openvino(args.pretrained, ckpt_file, ckpt_file.parent, tokenizer) if __name__ == "__main__": diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index 25556df1d5e..b1b1bfe4bb7 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -148,11 +148,13 @@ def try_repack(self, node: NNCFNode, graph: NNCFGraph) -> tuple[bool, int, int]: # representation [u2/u3] - zero_point per tensor offset = 2 ** (optimal_bits - 1) compressed_tensor = weight_tensor + offset + + target_type = self.name_to_node_mapping[next_nodes[0].node_name].get_element_type() compressed_const = create_ov_const_from_tensor(compressed_tensor, compression_dtype, name=const_node_name) - converted_const = opset.convert(compressed_const, ov.Type.f16) + converted_const = opset.convert(compressed_const, target_type) zero_point_const = opset.constant(offset, dtype=ov.Type.i8, name=f"{const_node_name}/zero_point") - zero_point_const = opset.convert(zero_point_const, ov.Type.f16) + zero_point_const = opset.convert(zero_point_const, target_type) converted_const = opset.subtract( converted_const, zero_point_const, name=f"{const_node_name}/zero_point/subtract" From ecc6f5a8b6f8990ca98fae2c4fe2a553f0eddb1e Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Tue, 6 Oct 2026 17:26:59 +0200 Subject: [PATCH 23/27] Moved repack_weights to other place. --- .../main.py | 3 +- src/nncf/__init__.py | 2 +- .../openvino/quantization/quantize_model.py | 10 ---- src/nncf/quantization/__init__.py | 1 - .../weight_compression/algorithm.py | 39 --------------- .../weight_compression/openvino_backend.py | 47 +++++++++++++++++++ src/nncf/quantization/quantize_model.py | 23 --------- 7 files changed, 49 insertions(+), 76 deletions(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py index 2f1ec94b513..e0238fdce4c 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py @@ -41,7 +41,6 @@ from nncf.quantization.advanced_parameters import AdvancedAWQParameters from nncf.quantization.advanced_parameters import AdvancedCompressionParameters from nncf.quantization.quantize_model import compress_weights -from nncf.quantization.quantize_model import repack_weights from nncf.torch import load_from_config from nncf.torch.function_hook.wrapper import get_hook_storage from nncf.torch.quantization.layers import AsymmetricLoraQuantizer @@ -300,7 +299,7 @@ def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path, tokenizer tokenizer.save_pretrained(ir_dir) model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) - model_to_eval.model = repack_weights(model_to_eval.model) + model_to_eval.model = nncf.repack_weights(model_to_eval.model) model_to_eval.save_pretrained(ir_dir / "repacked") tokenizer.save_pretrained(ir_dir / "repacked") print(f"The OpenVINO model has been repacked and saved to: {ir_dir / 'repacked'}") diff --git a/src/nncf/__init__.py b/src/nncf/__init__.py index 2a1ee31e549..816d2f18ead 100644 --- a/src/nncf/__init__.py +++ b/src/nncf/__init__.py @@ -54,7 +54,6 @@ from nncf.quantization import compress_weights as compress_weights from nncf.quantization import quantize as quantize from nncf.quantization import quantize_with_accuracy_control as quantize_with_accuracy_control -from nncf.quantization import repack_weights as repack_weights from nncf.quantization.advanced_parameters import ( AdvancedAccuracyRestorerParameters as AdvancedAccuracyRestorerParameters, ) @@ -68,6 +67,7 @@ from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters as AdvancedSmoothQuantParameters from nncf.quantization.advanced_parameters import GroupSizeFallbackMode as GroupSizeFallbackMode from nncf.quantization.advanced_parameters import OverflowFix as OverflowFix +from nncf.quantization.algorithms.weight_compression.openvino_backend import repack_weights as repack_weights from nncf.scopes import IgnoredScope as IgnoredScope from nncf.scopes import Subgraph as Subgraph from nncf.version import __version__ as __version__ diff --git a/src/nncf/openvino/quantization/quantize_model.py b/src/nncf/openvino/quantization/quantize_model.py index 39312e0636c..19031fb1674 100644 --- a/src/nncf/openvino/quantization/quantize_model.py +++ b/src/nncf/openvino/quantization/quantize_model.py @@ -422,13 +422,3 @@ def compress_weights_impl( statistics_points = statistics_aggregator.statistic_points return compression_algorithm.apply(model, graph, statistics_points, dataset) - - -def repack_weights_impl(model: ov.Model) -> ov.Model: - """ - Implementation of the `repack_weights()` method for the OpenVINO backend. - """ - model = remove_friendly_name_duplicates(model) - graph = build_graph(model) - - return WeightCompression.repack_weights(model, graph) diff --git a/src/nncf/quantization/__init__.py b/src/nncf/quantization/__init__.py index 0a1e91cabcb..cf5a6c6133e 100644 --- a/src/nncf/quantization/__init__.py +++ b/src/nncf/quantization/__init__.py @@ -14,4 +14,3 @@ from nncf.quantization.quantize_model import compress_weights as compress_weights from nncf.quantization.quantize_model import quantize as quantize from nncf.quantization.quantize_model import quantize_with_accuracy_control as quantize_with_accuracy_control -from nncf.quantization.quantize_model import repack_weights as repack_weights diff --git a/src/nncf/quantization/algorithms/weight_compression/algorithm.py b/src/nncf/quantization/algorithms/weight_compression/algorithm.py index a8ae0a4fdee..9553127d84b 100644 --- a/src/nncf/quantization/algorithms/weight_compression/algorithm.py +++ b/src/nncf/quantization/algorithms/weight_compression/algorithm.py @@ -1424,42 +1424,3 @@ def _get_statistics_for_weights_compression( for node in matmul_nodes: statistics[node.node_name] = copy.deepcopy(stats) return statistics - - @staticmethod - def repack_weights(model: TModel, graph: NNCFGraph) -> TModel: - """ - Repacks compressed weight constants to a lower-bit representation when possible. - - Iterates over all nodes in the graph and attempts to replace compressed constants - (i8, u8, i4, u4) with a more compact symmetric representation (i2, i3) if the - actual value range allows it. - - :param model: Backend-specific model with compressed weights. - :param graph: NNCFGraph instance corresponding to the model. - :return: The model with repacked weight constants. - """ - backend = get_backend(model) - if backend != BackendType.OPENVINO: - msg = f"Unsupported type of backend: {backend}" - raise nncf.UnsupportedBackendError(msg) - - from nncf.quantization.algorithms.weight_compression.openvino_backend import OVWeightCompressionAlgoBackend - - backend_entity = OVWeightCompressionAlgoBackend(model) - - n_repacked_tensors = 0 - bits_before = 0 - bits_after = 0 - - for node in graph.topological_sort(): - is_repacked, bits_before_node, bits_after_node = backend_entity.try_repack(node, graph) - if is_repacked: - n_repacked_tensors += 1 - bits_before += bits_before_node - bits_after += bits_after_node - nncf_logger.info( - f"\nTotal repacked tensors: {n_repacked_tensors}; " - f"\ntotal MB before: {bits_before / 8 / 1024 / 1024}; " - f"\ntotal MB after: {bits_after / 8 / 1024 / 1024}" - ) - return model diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index b1b1bfe4bb7..08516a5a2e2 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -13,12 +13,14 @@ import openvino as ov from openvino import opset13 as opset +from nncf.common.factory import build_graph from nncf.common.graph import NNCFGraph from nncf.common.graph import NNCFNode from nncf.common.graph.operator_metatypes import OperatorMetatype from nncf.common.graph.patterns.patterns import GraphPattern from nncf.common.graph.transformations.commands import TargetType from nncf.common.graph.utils import get_reduction_axes +from nncf.common.logging import nncf_logger from nncf.common.tensor_statistics.collectors import MeanAggregator from nncf.common.tensor_statistics.collectors import NoopAggregator from nncf.common.tensor_statistics.collectors import TensorCollector @@ -27,10 +29,14 @@ from nncf.common.tensor_statistics.statistics import MeanMagnitudeTensorStatistic from nncf.common.tensor_statistics.statistics import MeanVarianceTensorStatistic from nncf.common.tensor_statistics.statistics import WCTensorStatistic +from nncf.common.utils.api_marker import api +from nncf.common.utils.backend import BackendType +from nncf.common.utils.backend import get_backend from nncf.common.utils.caching import disable_results_caching from nncf.openvino.graph.metatypes import openvino_metatypes as om from nncf.openvino.graph.metatypes.groups import ATOMIC_ACTIVATIONS_OPERATIONS from nncf.openvino.graph.model_transformer import OVModelTransformer +from nncf.openvino.graph.model_utils import remove_friendly_name_duplicates from nncf.openvino.graph.node_utils import convert_op from nncf.openvino.graph.node_utils import create_ov_codebook_subgraph from nncf.openvino.graph.node_utils import create_ov_const_from_tensor @@ -540,3 +546,44 @@ def mean_abs_max_statistic_collector(reduction_axes: tuple[int], subset_size: in collector = TensorCollector(MeanMagnitudeTensorStatistic) collector.register_statistic_branch(MeanMagnitudeTensorStatistic.MEAN_MAGNITUDE_STAT, reducer, aggregator) return collector + + +@api(canonical_alias="nncf.repack_weights") +def repack_weights( + model: ov.Model, +) -> ov.Model: + """ + Looking for 4 and 8 bit weights in OV model and repack them if maximal absolute value corresponds + to the supported type with lower bits. + + :param model: A model to be repacked. + :type model: ov.Model + :return: The non-trainable model with repacked weights or the same model. + """ + backend = get_backend(model) + + if backend != BackendType.OPENVINO: + msg = f"Unsupported type of backend: {backend}" + raise Exception(msg) + + model = remove_friendly_name_duplicates(model) + graph = build_graph(model) + + backend_entity = OVWeightCompressionAlgoBackend(model) + + n_repacked_tensors = 0 + bits_before = 0 + bits_after = 0 + + for node in graph.topological_sort(): + is_repacked, bits_before_node, bits_after_node = backend_entity.try_repack(node, graph) + if is_repacked: + n_repacked_tensors += 1 + bits_before += bits_before_node + bits_after += bits_after_node + nncf_logger.info( + f"\nTotal repacked tensors: {n_repacked_tensors}; " + f"\ntotal MB before: {bits_before / 8 / 1024 / 1024}; " + f"\ntotal MB after: {bits_after / 8 / 1024 / 1024}" + ) + return model diff --git a/src/nncf/quantization/quantize_model.py b/src/nncf/quantization/quantize_model.py index 5d2fab5ece4..cdb2e0bdba3 100644 --- a/src/nncf/quantization/quantize_model.py +++ b/src/nncf/quantization/quantize_model.py @@ -782,26 +782,3 @@ def quantize_with_tune_hyperparams( quantized_model = hyperparameter_tuner.apply(model, validation_dataset) return quantized_model - - -@api(canonical_alias="nncf.repack_weights") -def repack_weights( - model: TModel, -) -> TModel: - """ - Looking for 4 and 8 bit weights in OV model and repack them if maximal absolute value corresponds - to the supported type with lower bits. - - :param model: A model to be repacked. - :type model: TModel - :return: The non-trainable model with repacked weights or the same model. - """ - backend = get_backend(model) - - if backend != BackendType.OPENVINO: - msg = f"Unsupported type of backend: {backend}" - raise nncf.UnsupportedBackendError(msg) - - from nncf.openvino.quantization.quantize_model import repack_weights_impl - - return repack_weights_impl(model) # type: ignore From 8101998a69df459d067c5fd7768f41f126930439 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 7 Oct 2026 10:29:00 +0200 Subject: [PATCH 24/27] Moved repack_weights to other place. --- src/nncf/__init__.py | 2 +- .../weight_compression/openvino_backend.py | 18 ++----- src/nncf/quantization/repack_weights.py | 47 +++++++++++++++++++ 3 files changed, 51 insertions(+), 16 deletions(-) create mode 100644 src/nncf/quantization/repack_weights.py diff --git a/src/nncf/__init__.py b/src/nncf/__init__.py index 816d2f18ead..140db9b5d6d 100644 --- a/src/nncf/__init__.py +++ b/src/nncf/__init__.py @@ -67,7 +67,7 @@ from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters as AdvancedSmoothQuantParameters from nncf.quantization.advanced_parameters import GroupSizeFallbackMode as GroupSizeFallbackMode from nncf.quantization.advanced_parameters import OverflowFix as OverflowFix -from nncf.quantization.algorithms.weight_compression.openvino_backend import repack_weights as repack_weights +from nncf.quantization.repack_weights import repack_weights as repack_weights from nncf.scopes import IgnoredScope as IgnoredScope from nncf.scopes import Subgraph as Subgraph from nncf.version import __version__ as __version__ diff --git a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py index 08516a5a2e2..7ec827157fe 100644 --- a/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py +++ b/src/nncf/quantization/algorithms/weight_compression/openvino_backend.py @@ -13,7 +13,6 @@ import openvino as ov from openvino import opset13 as opset -from nncf.common.factory import build_graph from nncf.common.graph import NNCFGraph from nncf.common.graph import NNCFNode from nncf.common.graph.operator_metatypes import OperatorMetatype @@ -29,14 +28,10 @@ from nncf.common.tensor_statistics.statistics import MeanMagnitudeTensorStatistic from nncf.common.tensor_statistics.statistics import MeanVarianceTensorStatistic from nncf.common.tensor_statistics.statistics import WCTensorStatistic -from nncf.common.utils.api_marker import api -from nncf.common.utils.backend import BackendType -from nncf.common.utils.backend import get_backend from nncf.common.utils.caching import disable_results_caching from nncf.openvino.graph.metatypes import openvino_metatypes as om from nncf.openvino.graph.metatypes.groups import ATOMIC_ACTIVATIONS_OPERATIONS from nncf.openvino.graph.model_transformer import OVModelTransformer -from nncf.openvino.graph.model_utils import remove_friendly_name_duplicates from nncf.openvino.graph.node_utils import convert_op from nncf.openvino.graph.node_utils import create_ov_codebook_subgraph from nncf.openvino.graph.node_utils import create_ov_const_from_tensor @@ -548,9 +543,9 @@ def mean_abs_max_statistic_collector(reduction_axes: tuple[int], subset_size: in return collector -@api(canonical_alias="nncf.repack_weights") def repack_weights( model: ov.Model, + graph: NNCFGraph, ) -> ov.Model: """ Looking for 4 and 8 bit weights in OV model and repack them if maximal absolute value corresponds @@ -558,17 +553,10 @@ def repack_weights( :param model: A model to be repacked. :type model: ov.Model + :param graph: The NNCF graph corresponding to the model. + :type graph: NNCFGraph :return: The non-trainable model with repacked weights or the same model. """ - backend = get_backend(model) - - if backend != BackendType.OPENVINO: - msg = f"Unsupported type of backend: {backend}" - raise Exception(msg) - - model = remove_friendly_name_duplicates(model) - graph = build_graph(model) - backend_entity = OVWeightCompressionAlgoBackend(model) n_repacked_tensors = 0 diff --git a/src/nncf/quantization/repack_weights.py b/src/nncf/quantization/repack_weights.py new file mode 100644 index 00000000000..bc67ee423c5 --- /dev/null +++ b/src/nncf/quantization/repack_weights.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 Intel Corporation +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TypeVar + +import nncf +from nncf.common.factory import build_graph +from nncf.common.utils.api_marker import api +from nncf.common.utils.backend import BackendType +from nncf.common.utils.backend import get_backend +from nncf.openvino.graph.model_utils import remove_friendly_name_duplicates + +TModel = TypeVar("TModel") + + +@api(canonical_alias="nncf.repack_weights") +def repack_weights( + model: TModel, +) -> TModel: + """ + Looking for 4 and 8 bit weights in OV model and repack them if maximal absolute value corresponds + to the supported type with lower bits. + + :param model: A model to be repacked. + :type model: ov.Model + :return: The non-trainable model with repacked weights or the same model. + """ + backend = get_backend(model) + + if backend != BackendType.OPENVINO: + msg = f"Unsupported type of backend: {backend}" + raise nncf.UnsupportedBackendError(msg) + + model = remove_friendly_name_duplicates(model) + graph = build_graph(model) + + from nncf.quantization.algorithms.weight_compression.openvino_backend import repack_weights as repack_weights_impl + + return repack_weights_impl(model, graph) From e2df895daa47cadacac949b8a53d76349199a501 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 7 Oct 2026 11:07:32 +0200 Subject: [PATCH 25/27] Added right scale merging for case of Qwen3NextRMSNorm. --- .../torch/distillation_qad_with_lora_2_3_bit/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py index e0238fdce4c..b9dc836c543 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py @@ -559,7 +559,7 @@ def equalize_up_gate_with_layernorm(model: nn.Module, eps: float = 1e-5, use_ali # Standard RMSNorm (e.g. Llama) uses weight*x (weight initialized to ones): # effective multiplier = weight # to scale output by s: w_new = s * w_old - if "gemma" in type(producer).__name__.lower(): + if "gemma" in type(producer).__name__.lower() or "qwen3nextrmsnorm" in type(producer).__name__.lower(): _set_weight_data(producer, s_dev * (1.0 + w_prod) - 1.0) else: _set_weight_data(producer, w_prod * s_dev) From ffc7fa65fb2e52e360a9b7eb058360d3d6f768dc Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 7 Oct 2026 11:24:32 +0200 Subject: [PATCH 26/27] Changes imports order. --- src/nncf/__init__.py | 2 +- src/nncf/quantization/__init__.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/nncf/__init__.py b/src/nncf/__init__.py index 140db9b5d6d..2a1ee31e549 100644 --- a/src/nncf/__init__.py +++ b/src/nncf/__init__.py @@ -54,6 +54,7 @@ from nncf.quantization import compress_weights as compress_weights from nncf.quantization import quantize as quantize from nncf.quantization import quantize_with_accuracy_control as quantize_with_accuracy_control +from nncf.quantization import repack_weights as repack_weights from nncf.quantization.advanced_parameters import ( AdvancedAccuracyRestorerParameters as AdvancedAccuracyRestorerParameters, ) @@ -67,7 +68,6 @@ from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters as AdvancedSmoothQuantParameters from nncf.quantization.advanced_parameters import GroupSizeFallbackMode as GroupSizeFallbackMode from nncf.quantization.advanced_parameters import OverflowFix as OverflowFix -from nncf.quantization.repack_weights import repack_weights as repack_weights from nncf.scopes import IgnoredScope as IgnoredScope from nncf.scopes import Subgraph as Subgraph from nncf.version import __version__ as __version__ diff --git a/src/nncf/quantization/__init__.py b/src/nncf/quantization/__init__.py index cf5a6c6133e..29aa7c41cc0 100644 --- a/src/nncf/quantization/__init__.py +++ b/src/nncf/quantization/__init__.py @@ -14,3 +14,4 @@ from nncf.quantization.quantize_model import compress_weights as compress_weights from nncf.quantization.quantize_model import quantize as quantize from nncf.quantization.quantize_model import quantize_with_accuracy_control as quantize_with_accuracy_control +from nncf.quantization.repack_weights import repack_weights as repack_weights From 1f558e0db1c0b35e73cc7630b9907b2b8b7c4b8b Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 7 Oct 2026 11:57:43 +0200 Subject: [PATCH 27/27] Disable autocast of bfloat16 matmul output to fp32. --- .../torch/distillation_qad_with_lora_2_3_bit/main.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py index b9dc836c543..149e8590cfe 100644 --- a/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py +++ b/examples/llm_compression/torch/distillation_qad_with_lora_2_3_bit/main.py @@ -692,7 +692,8 @@ def form_batch(inputs: list[Tensor], model_input: bool): targets = targets / fls targets = torch.tanh(targets) targets = targets * fls - outputs = model(**inputs).logits + with torch.autocast(device_type=device, dtype=torch_dtype): + outputs = model(**inputs).logits loss = kl_div(outputs, targets.to(dtype=torch_dtype, device=device)) # Perform an optimization step after accumulating gradients over multiple minibatches.