no wa for device in fq, test for lora and chkp load

Author: ljaljushkin

nncf/quantization/algorithms/weight_compression/algorithm.py

@@ -160,9 +160,21 @@ def check_user_compression_configuration(
         msg = f"The ratio should be between 0 and 1, but ratio={ratio} is specified."
         raise nncf.ValidationError(msg)
 
-    if subset_size <= 0:
-        msg = f"The subset_size value should be positive, but subset_size={subset_size} is given."
-        raise nncf.ValidationError(msg)
+    for size in [
+        subset_size,
+        advanced_parameters.awq_params.subset_size,
+        advanced_parameters.scale_estimation_params.subset_size,
+        advanced_parameters.gptq_params.subset_size,
+        advanced_parameters.lora_correction_params.subset_size,
+    ]:
+        if size <= 0:
+            msg = f"The subset_size value should be positive, but subset_size={size} is given."
+            raise nncf.ValidationError(msg)
+
+    for rank in [advanced_parameters.lora_adapter_rank, advanced_parameters.lora_correction_params.adapter_rank]:
+        if rank <= 0:
+            msg = f"The lora adapter rank should be positive, but rank={rank} is given."
+            raise nncf.ValidationError(msg)
 
     if (
         ratio
@@ -656,6 +668,7 @@ def apply(
             zero_points,
             lora_correction_algo,
             self._compression_format,
+            self._advanced_parameters,
         )
 
         self._backend_entity.dump_parameters(

nncf/torch/model_transformer.py

@@ -29,6 +29,8 @@
 from nncf.torch.graph.transformations.commands import PTTargetPoint
 from nncf.torch.graph.transformations.commands import PTWeightUpdateCommand
 from nncf.torch.graph.transformations.layout import PTTransformationLayout
+from nncf.torch.model_graph_manager import get_const_data
+from nncf.torch.model_graph_manager import get_const_node
 from nncf.torch.model_graph_manager import update_fused_bias
 from nncf.torch.module_operations import UpdateWeight
 from nncf.torch.nncf_network import NNCFNetwork
@@ -74,6 +76,19 @@ def transform(self, transformation_layout: PTTransformationLayout) -> NNCFNetwor
                 model = transformation_fn(model, transformations)
 
         if requires_graph_rebuild:
+            graph = model.nncf.get_original_graph()
+            for command in transformation_layout.transformations:
+                compression_module = command.fn
+                if isinstance(compression_module, nn.Module):
+                    target_point = command.target_points[0]
+                    node_with_weight = graph.get_node_by_name(target_point.target_node_name)
+                    weight_node = get_const_node(node_with_weight, target_point.input_port_id, graph)
+                    if weight_node is None:
+                        weight_node = node_with_weight  # Decompression in DQ compression format is applied to const.
+                    const_data = get_const_data(weight_node, model)
+                    # Compression module and the corresponding layer may have a different device in multi-device setup
+                    # (e.g. when HF model was loaded with device_map='auto'). Need to align devices.
+                    compression_module.to(const_data.device)
             model.nncf.rebuild_graph()
 
         return model

nncf/torch/quantization/layers.py

@@ -1067,9 +1067,14 @@ class LoraMixin:
 
     def __init__(self, lspec: PTLoraSpec):
         self._lspec = lspec
+        default_lora_dtype = torch.bfloat16
         out_features, in_features = lspec.orig_weight_shape
-        self._lora_A = torch.nn.Parameter(torch.ones((lspec.lora_rank, in_features), dtype=torch.bfloat16))
-        self._lora_B = torch.nn.Parameter(torch.zeros((out_features, lspec.lora_rank), dtype=torch.bfloat16))
+        rank = lspec.lora_rank
+        if rank > out_features or rank > in_features:
+            msg = f"Specified LoRA rank={rank} cannot exceed any dimension of the weight tensor"
+            raise nncf.ValidationError(msg)
+        self._lora_A = torch.nn.Parameter(torch.ones((rank, in_features), dtype=default_lora_dtype))
+        self._lora_B = torch.nn.Parameter(torch.zeros((out_features, rank), dtype=default_lora_dtype))
 
     def enable_gradients(self):
         self._lora_A.requires_grad = True

tests/torch/ptq/test_fq_lora.py

@@ -15,9 +15,18 @@
 from transformers import AutoTokenizer
 
 import nncf
+from nncf.data.dataset import Dataset
+from nncf.errors import ValidationError
+from nncf.parameters import CompressionFormat
+from nncf.parameters import CompressWeightsMode
+from nncf.quantization.advanced_parameters import AdvancedCompressionParameters
+from nncf.quantization.quantize_model import compress_weights
+from nncf.scopes import IgnoredScope
+from nncf.torch import load_from_config
 from nncf.torch.quantization.layers import AsymmetricQuantizer as AQ
 from nncf.torch.quantization.layers import LoraMixin
 from nncf.torch.quantization.layers import SymmetricQuantizer as SQ
+from tests.torch.test_models.synthetic import LinearModel
 
 
 @pytest.mark.parametrize(
@@ -80,3 +89,70 @@ def test_fq_lora_tuning(mode, backup_mode, compression_kwargs, ref_num_trainable
 
     assert first_loss > 8
     assert float(loss) < 1
+
+
+def test_checkpoint_loading(tmp_path):
+    model_id = "hf-internal-testing/tiny-random-GPTNeoXForCausalLM"
+    if not torch.cuda.is_available():
+        pytest.skip("Skipping CUDA test case for CPU only setups.")
+    device = "cuda"
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    example_input = tokenizer("dummy", return_tensors="pt").to(device)
+    ref_output = tokenizer.decode(
+        model.generate(**example_input, do_sample=False, max_new_tokens=20)[0], skip_special_tokens=True
+    )
+    except_lm_head_and_5th_vproj = (
+        r"^(?!.*(GPTNeoXLayer\[2\]/GPTNeoXSdpaAttention\[attention\]/Linear\[query_key_value\]/l|embed_out).*$).*$"
+    )
+    model = compress_weights(
+        model,
+        group_size=32,
+        mode=CompressWeightsMode.INT4_ASYM,
+        backup_mode=CompressWeightsMode.INT8_ASYM,
+        dataset=Dataset([dict(example_input)]),
+        compression_format=CompressionFormat.FQ_LORA,
+        ignored_scope=IgnoredScope(patterns=[except_lm_head_and_5th_vproj]),
+        advanced_parameters=AdvancedCompressionParameters(lora_adapter_rank=2),
+    )
+    ref_output = tokenizer.decode(
+        model.generate(**example_input, do_sample=False, max_new_tokens=20)[0], skip_special_tokens=True
+    )
+
+    # save checkpoint
+    ckpt_path = tmp_path / "nncf_ckpt.pth"
+    torch.save(
+        {
+            "nncf_state_dict": model.nncf.state_dict(),
+            "nncf_config": model.nncf.get_config(),
+        },
+        ckpt_path,
+    )
+    del model
+
+    # load checkpoint
+    nncf_ckpt = torch.load(ckpt_path, weights_only=False)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+    model = load_from_config(model, nncf_ckpt["nncf_config"], example_input=dict(example_input))
+    model.nncf.load_state_dict(nncf_ckpt["nncf_state_dict"])
+
+    actual_output = tokenizer.decode(
+        model.generate(**example_input, do_sample=False, max_new_tokens=20)[0],
+        skip_special_tokens=True,
+    )
+    assert actual_output == ref_output
+
+
+def test_invalid_lora_rank():
+    too_big_rank = 4
+    model = LinearModel(torch.ones(2, 2))
+    with pytest.raises(ValidationError):
+        compress_weights(
+            model,
+            mode=CompressWeightsMode.INT4_ASYM,
+            group_size=2,
+            all_layers=True,
+            dataset=Dataset([torch.ones(2, 2)]),
+            compression_format=CompressionFormat.FQ_LORA,
+            advanced_parameters=AdvancedCompressionParameters(lora_adapter_rank=too_big_rank),
+        )

tests/torch/ptq/test_weights_compression.py

@@ -39,6 +39,7 @@
 from nncf.torch.quantization.quantize_functions import unpack_int4
 from nncf.torch.quantization.quantize_functions import unpack_uint4
 from tests.cross_fw.test_templates.template_test_weights_compression import TemplateWeightCompression
+from tests.torch.test_models.synthetic import LinearModel
 from tests.torch.test_models.synthetic import ShortTransformer
 from tests.torch.test_tensor import cast_to
 
@@ -82,16 +83,6 @@ def forward(self, input):
         return input @ self.w
 
 
-class LinearModel(torch.nn.Module):
-    def __init__(self, weight: torch.Tensor = torch.ones(size=(256, 256), dtype=torch.float32)):
-        super().__init__()
-        self.linear = torch.nn.Linear(weight.shape[0], weight.shape[1], False)
-        self.linear.weight = torch.nn.Parameter(weight)
-
-    def forward(self, input):
-        return self.linear(input)
-
-
 class AWQActLinearModel(nn.Module):
     def __init__(self, with_multiply=False, n_layers=8):
         super().__init__()

tests/torch/test_models/synthetic.py

@@ -660,6 +660,16 @@ def forward(self, input_ids):
         return res
 
 
+class LinearModel(torch.nn.Module):
+    def __init__(self, weight: torch.Tensor = torch.ones(size=(256, 256), dtype=torch.float32)):
+        super().__init__()
+        self.linear = torch.nn.Linear(weight.shape[0], weight.shape[1], False)
+        self.linear.weight = torch.nn.Parameter(weight)
+
+    def forward(self, input):
+        return self.linear(input)
+
+
 class YOLO11N_SDPABlock(torch.nn.Module):
     INPUT_SIZE = (1, 2, 4)