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 dc58f78c59e..fe50a862e1a 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 @@ -23,13 +23,16 @@ from datasets import load_dataset from optimum.exporters.openvino.convert import export_from_model from optimum.intel.openvino import OVModelForCausalLM +from optimum.intel.openvino import OVModelForVisualCausalLM 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 AutoConfig from transformers import AutoModelForCausalLM +from transformers import AutoModelForMultimodalLM from transformers import AutoTokenizer import nncf @@ -281,29 +284,88 @@ def load_checkpoint(model: nn.Module, ckpt_file: Path) -> nn.Module: return model +def extract_causal_lm_from_visual(model: nn.Module) -> nn.Module: + """ + Extracts the causal language model part from a visual model. + + :param model: The visual model containing a causal language model. + :returns: The extracted causal language model. + """ + model_container = getattr(model, "model", None) + language_model = getattr(model_container, "language_model", None) + lm_head = getattr(model, "lm_head", None) + text_config = getattr(getattr(model, "config", None), "text_config", None) + if not isinstance(language_model, nn.Module) or not isinstance(lm_head, nn.Module) or text_config is None: + msg = "Expected a multimodal model with model.language_model, lm_head, and config.text_config." + raise ValueError(msg) + + with torch.device("meta"): + causal_lm = AutoModelForCausalLM.from_config(text_config) + + causal_lm.model = language_model + causal_lm.lm_head = lm_head + return causal_lm + + @torch.no_grad() -def export_to_openvino(pretrained: str, ckpt_file: Path, ir_dir: Path, tokenizer: AutoTokenizer, bits: int): +def export_to_openvino( + pretrained: str, + ckpt_file: Path, + ir_dir: Path, + tokenizer: AutoTokenizer, + bits: int, +) -> None: """ 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: None. The OpenVINO model is exported and saved to the specified directory. + :param tokenizer: The tokenizer to save alongside the repacked model. + :param bits: Number of bits used for quantization. """ - 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) - - if bits < 4: - model_to_eval = OVModelForCausalLM.from_pretrained(ir_dir) - 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'}") + config = AutoConfig.from_pretrained(pretrained) + is_visual = hasattr(config, "vision_config") and hasattr(config, "text_config") + + if is_visual: + model = AutoModelForMultimodalLM.from_pretrained(pretrained, device_map="cpu", torch_dtype=torch.float32) + if "gemma" in pretrained.lower(): + model = load_checkpoint(model, ckpt_file) + nncf.strip(model, do_copy=False, strip_format=StripFormat.RP) + else: + language_model = extract_causal_lm_from_visual(model) + language_model = load_checkpoint(language_model, ckpt_file) + nncf.strip(language_model, do_copy=False, strip_format=StripFormat.RP) + else: + model = AutoModelForCausalLM.from_pretrained(pretrained, device_map="cpu", torch_dtype=torch.float32) + model = load_checkpoint(model, ckpt_file) + model = nncf.strip(model, do_copy=False, strip_format=StripFormat.DQ) + + model.eval() + try: + export_from_model(model, ir_dir, device="cpu") + except Exception as e: + print(f"Failed to export the model or part of the model, but repacking will still proceed: {e}") + + if is_visual: + ov_model = OVModelForVisualCausalLM.from_pretrained(ir_dir) + if bits < 4: + ov_model.language_model.model = nncf.repack_weights(ov_model.language_model.model) + for component in ov_model.components.values(): + if component is ov_model.language_model: + continue + component.model = nncf.compress_weights( + component.model, + mode=nncf.CompressWeightsMode.INT8_ASYM, + ) + else: + ov_model = OVModelForCausalLM.from_pretrained(ir_dir) + if bits < 4: + ov_model.model = nncf.repack_weights(ov_model.model) + + repacked_dir = ir_dir / "repacked" + ov_model.save_pretrained(repacked_dir) + tokenizer.save_pretrained(repacked_dir) + print(f"The OpenVINO model has been repacked and saved to: {repacked_dir}") @torch.no_grad() diff --git a/src/nncf/parameters.py b/src/nncf/parameters.py index 6b7592dea8c..5d81e8730e5 100644 --- a/src/nncf/parameters.py +++ b/src/nncf/parameters.py @@ -150,11 +150,14 @@ class StripFormat(StrEnum): in low-bit precision using fake quantize parameters. This is the default format for deploying models with compressed weights. :param IN_PLACE: Directly applies NNCF operations to the weights, replacing the original weights. + :param RP: Registers parametrization for the compressed weights, allowing to handle decompression graph + without hooks then changing the model infernce graph in optimum-intel. """ NATIVE = "native" DQ = "dequantize" IN_PLACE = "in_place" + RP = "register_parametrization" @api(canonical_alias="nncf.BackupMode") diff --git a/src/nncf/torch/function_hook/strip.py b/src/nncf/torch/function_hook/strip.py index c176a23df34..73659dc0c4f 100644 --- a/src/nncf/torch/function_hook/strip.py +++ b/src/nncf/torch/function_hook/strip.py @@ -59,6 +59,8 @@ def strip_model(model: TModel, example_input: Any = None, strip_format: StripFor model = replace_quantizer_to_compressed_weight_with_decompressor(model) elif strip_format == StripFormat.IN_PLACE: model = apply_compression_in_place(model) + elif strip_format == StripFormat.RP: + model = register_parametrization_with_decompressor(model) else: msg = f"Unsupported strip format: {strip_format}" raise nncf.ParameterNotSupportedError(msg) @@ -199,3 +201,61 @@ def apply_compression_in_place(model: TModel) -> TModel: for hook_name in hooks_to_delete: hook_storage.delete_hook(hook_name) return model + + +@torch.no_grad() +def register_parametrization_with_decompressor(model: TModel) -> TModel: + """ + Performs transformation from fake quantize format (FQ) to dequantization one (DQ): + (weights + FQ) -> (compressed_weights + DQ) + + The decompressor is registered as a parametrization of every module attribute that holds the compressed + weight, so tied weights (e.g. shared between lm_head and embeddings) are decompressed correctly + for any number of bits. + + :param model: Compressed model + :return: The modified NNCF network. + """ + hook_storage = get_hook_storage(model) + hooks_to_delete = [] + decompressor_per_param = {} + + for hook_name, hook_module in hook_storage.named_hooks(): + if not isinstance(hook_module, (SymmetricQuantizer, AsymmetricQuantizer)): + continue + 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" + raise nncf.ValidationError(msg) + + _, op_name, _ = decode_hook_name(hook_name) + module_name, weight_attr_name = split_const_name(op_name) + weight_param = getattr(get_module_by_name(module_name, model), weight_attr_name) + hooks_to_delete.append(hook_name) + + if weight_param in decompressor_per_param: + # The weight is already packed through another hook + continue + + if isinstance(hook_module, AsymmetricQuantizer): + decompressor, q_weight = asym_fq_to_decompressor(hook_module, weight_param) + else: + decompressor, q_weight = sym_fq_to_decompressor(hook_module, weight_param) # type: ignore[assignment] + + weight_param.requires_grad = False + weight_param.data = decompressor.pack_weight(q_weight) + decompressor_per_param[weight_param] = decompressor + + # A parameter can be shared by several modules, all of them have to decompress it. + # Targets are collected before registration, because registration changes the module parameters. + targets = [ + (module, attr_name, decompressor_per_param[param]) + for module in model.modules() + for attr_name, param in module.named_parameters(recurse=False) + if param in decompressor_per_param + ] + for module, attr_name, decompressor in targets: + nn.utils.parametrize.register_parametrization(module, attr_name, decompressor, unsafe=True) + + for hook_name in hooks_to_delete: + hook_storage.delete_hook(hook_name) + return model diff --git a/tests/torch/function_hook/quantization/strip/test_strip_parametrization.py b/tests/torch/function_hook/quantization/strip/test_strip_parametrization.py new file mode 100644 index 00000000000..bba37727f6d --- /dev/null +++ b/tests/torch/function_hook/quantization/strip/test_strip_parametrization.py @@ -0,0 +1,114 @@ +# 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 pytest +import torch +from torch import nn +from torch.nn.utils import parametrize + +import nncf +from nncf.parameters import CompressWeightsMode +from nncf.parameters import StripFormat +from nncf.torch.function_hook.wrapper import get_hook_storage +from nncf.torch.function_hook.wrapper import wrap_model +from nncf.torch.quantization.layers import BaseWeightsDecompressor +from tests.torch.helpers import LinearModel + +INPUT_SHAPE = [1, 16] + + +class TiedModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = nn.Embedding(16, 16) + self.head = nn.Linear(16, 16, bias=False) + self.head.weight = self.embedding.weight + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.head(self.embedding(x)) + + +def _compress(model: nn.Module, example_input: torch.Tensor, mode: CompressWeightsMode) -> nn.Module: + extra_args = {"ratio": 1, "group_size": 4, "all_layers": True} if mode.value.startswith("int4") else {} + return nncf.compress_weights( + model, + mode=mode, + dataset=nncf.Dataset([example_input]), + compression_format=nncf.CompressionFormat.FQ_LORA, + advanced_parameters=nncf.AdvancedCompressionParameters(lora_adapter_rank=1), + **extra_args, + ) + + +@pytest.mark.parametrize( + "mode", + ( + CompressWeightsMode.INT4_SYM, + CompressWeightsMode.INT4_ASYM, + CompressWeightsMode.INT8_SYM, + CompressWeightsMode.INT8_ASYM, + ), + ids=lambda mode: mode.value, +) +def test_strip_rp_linear(mode: CompressWeightsMode) -> None: + example_input = torch.ones(INPUT_SHAPE) + compressed_model = _compress(LinearModel(input_shape=INPUT_SHAPE), example_input, mode) + + with torch.no_grad(): + ref_output = compressed_model(example_input) + stripped_model = nncf.strip(compressed_model, do_copy=True, strip_format=StripFormat.RP) + output = stripped_model(example_input) + + # Quantizers are removed from the hook storage and replaced by a parametrization + assert not list(get_hook_storage(stripped_model).named_hooks()) + assert parametrize.is_parametrized(stripped_model.linear, "weight") + decompressors = stripped_model.linear.parametrizations.weight + assert len(decompressors) == 1 + assert isinstance(decompressors[0], BaseWeightsDecompressor) + + original_weight = stripped_model.linear.parametrizations.weight.original + assert original_weight.dtype in (torch.uint8, torch.int8) + assert not original_weight.requires_grad + assert torch.allclose(ref_output, output, atol=1e-3) + + +def test_strip_rp_in_place() -> None: + example_input = torch.ones(INPUT_SHAPE) + compressed_model = _compress(LinearModel(input_shape=INPUT_SHAPE), example_input, CompressWeightsMode.INT8_SYM) + + stripped_model = nncf.strip(compressed_model, do_copy=False, strip_format=StripFormat.RP) + + assert stripped_model is compressed_model + + +def test_strip_rp_tied_weights() -> None: + example_input = torch.arange(4) + model = TiedModel() + compressed_model = _compress(model, example_input, CompressWeightsMode.INT8_SYM) + with torch.no_grad(): + ref_output = compressed_model(example_input) + stripped_model = nncf.strip(compressed_model, do_copy=False, strip_format=StripFormat.RP) + output = stripped_model(example_input) + + # All modules sharing the packed weight must decompress it + for module in (stripped_model.embedding, stripped_model.head): + assert parametrize.is_parametrized(module, "weight") + assert torch.allclose(ref_output, output, atol=5e-2) + + +def test_strip_rp_without_compression_hooks() -> None: + model = LinearModel(input_shape=INPUT_SHAPE) + ref_weight = model.linear.weight.clone() + + stripped_model = nncf.strip(wrap_model(model), strip_format=StripFormat.RP) + + assert not parametrize.is_parametrized(stripped_model.linear) + assert torch.equal(stripped_model.linear.weight, ref_weight)