Torch strip register parametrization and export image-text-to-text models after training.

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()

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")

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

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)