[OpenVINO] calibration_device parameter

Author: daniil-lyakhov

src/nncf/openvino/engine.py

@@ -10,6 +10,10 @@
 # limitations under the License.
 
 
+import contextvars
+from collections.abc import Generator
+from contextlib import contextmanager
+
 import numpy as np
 import openvino as ov
 from openvino import Type
@@ -19,6 +23,17 @@
 from nncf.definitions import NNCF_DATASET_RESET_STATE_KEY
 from nncf.openvino.graph.model_utils import model_has_state
 
+_calibration_device: contextvars.ContextVar[str | None] = contextvars.ContextVar("_calibration_device", default=None)
+
+
+@contextmanager
+def calibration_device_context(device: str | None) -> Generator[None, None, None]:
+    token = _calibration_device.set(device)
+    try:
+        yield
+    finally:
+        _calibration_device.reset(token)
+
 
 class OVCompiledModelEngine(Engine):
     """
@@ -79,12 +94,13 @@ def __init__(self, model: ov.Model, use_fp32_precision: bool = True):
         :param use_fp32_precision: A flag that determines whether to force the engine to use FP32
             precision during inference.
         """
+        device_name = _calibration_device.get() or "CPU"
         config = None
         if use_fp32_precision:
             config = {inference_precision: Type.f32}
         ie = ov.Core()
         stateful = model_has_state(model)
-        compiled_model = ie.compile_model(model, device_name="CPU", config=config)
+        compiled_model = ie.compile_model(model, device_name=device_name, config=config)
         self.engine = OVCompiledModelEngine(compiled_model, stateful)
 
     def infer(

src/nncf/openvino/quantization/quantize_model.py

@@ -21,6 +21,7 @@
 from nncf.common.logging import nncf_logger
 from nncf.common.quantization.structs import QuantizationPreset
 from nncf.data import Dataset
+from nncf.openvino.engine import calibration_device_context
 from nncf.openvino.graph.metatypes.groups import OPERATIONS_OUTPUT_HAS_NO_BATCH_AXIS
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVIfMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import get_node_metatype
@@ -119,9 +120,11 @@ def _extract_all_subgraphs(model: ov.Model, current_id: str) -> None:
         f"The model consists of {if_ops_number} If node(-s) with then and else bodies. \
             Main model and all If bodies will be quantized recursively."
     )
-    quantized_model, _ = apply_algorithm_if_bodies(
-        quantization_algorithm, model, graphs, main_model_graph_id, calibration_dataset, subset_size, 1
-    )
+    calibration_device = advanced_parameters.calibration_device if advanced_parameters else None
+    with calibration_device_context(calibration_device):
+        quantized_model, _ = apply_algorithm_if_bodies(
+            quantization_algorithm, model, graphs, main_model_graph_id, calibration_dataset, subset_size, 1
+        )
 
     if is_weight_compression_needed(advanced_parameters):
         compress_quantize_weights_transformation(quantized_model)
@@ -168,7 +171,9 @@ def native_quantize_impl(
     )
     graph = GraphConverter.create_nncf_graph(model)
     warning_model_no_batchwise_support(graph, advanced_parameters, model_type, OPERATIONS_OUTPUT_HAS_NO_BATCH_AXIS)
-    quantized_model = quantization_algorithm.apply(model, graph, dataset=calibration_dataset)
+    calibration_device = advanced_parameters.calibration_device if advanced_parameters else None
+    with calibration_device_context(calibration_device):
+        quantized_model = quantization_algorithm.apply(model, graph, dataset=calibration_dataset)
 
     if is_weight_compression_needed(advanced_parameters):
         compress_quantize_weights_transformation(quantized_model)
@@ -296,15 +301,19 @@ def quantize_with_accuracy_control_impl(
             advanced_accuracy_restorer_parameters.num_ranking_workers,
             advanced_accuracy_restorer_parameters.restore_mode,
         )
-        quantized_model = accuracy_restorer.apply(
-            model,
-            initial_metric_results,
-            quantized_model,
-            quantized_metric_results,
-            validation_dataset,
-            validation_dataset_size,
-            evaluator,
+        calibration_device = (
+            advanced_quantization_parameters.calibration_device if advanced_quantization_parameters else None
         )
+        with calibration_device_context(calibration_device):
+            quantized_model = accuracy_restorer.apply(
+                model,
+                initial_metric_results,
+                quantized_model,
+                quantized_metric_results,
+                validation_dataset,
+                validation_dataset_size,
+                evaluator,
+            )
 
     if compress_weights:
         compress_quantize_weights_transformation(quantized_model)
@@ -402,12 +411,15 @@ def compress_weights_impl(
         advanced_parameters,
     )
 
+    calibration_device = advanced_parameters.calibration_device if advanced_parameters else None
+
     statistics_points = None
     if advanced_parameters and advanced_parameters.statistics_path:
         # If there is no such directory, then caches statistics
         statistics_path = Path(advanced_parameters.statistics_path)
         if not statistics_path.exists():
-            cache_weight_compression_statistics(model, graph, dataset, subset_size, statistics_path)
+            with calibration_device_context(calibration_device):
+                cache_weight_compression_statistics(model, graph, dataset, subset_size, statistics_path)
         statistics_aggregator = StatisticsAggregatorFactory.create(model, dataset)
         compression_algorithm.set_backend_entity(model)
         _, matmul_input_to_output_nodes_map = compression_algorithm.get_compression_nodes_info(graph)
@@ -421,4 +433,5 @@ def compress_weights_impl(
         statistics_aggregator.load_statistics_from_dir(statistics_path)
         statistics_points = statistics_aggregator.statistic_points
 
-    return compression_algorithm.apply(model, graph, statistics_points, dataset)
+    with calibration_device_context(calibration_device):
+        return compression_algorithm.apply(model, graph, statistics_points, dataset)

src/nncf/quantization/advanced_parameters.py

@@ -252,6 +252,10 @@ class AdvancedQuantizationParameters:
     :type smooth_quant_alpha: float
     :param backend_params: Backend-specific parameters.
     :type backend_params: dict[str, Any]
+    :param calibration_device: OpenVINO device name to use for calibration inference
+        (e.g. "CPU", "GPU", "GPU.0", "AUTO:GPU,CPU"). If None, defaults to "CPU".
+        Only applicable to the OpenVINO backend.
+    :type calibration_device: Optional[str]
     """
 
     # General parameters
@@ -282,6 +286,9 @@ class AdvancedQuantizationParameters:
     # Backend specific parameters
     backend_params: dict[str, Any] = field(default_factory=dict)
 
+    # Calibration device
+    calibration_device: str | None = None
+
 
 @api()
 @dataclass
@@ -427,6 +434,10 @@ class AdvancedCompressionParameters:
     :type lora_correction_params: AdvancedLoraCorrectionParameters
     :param backend_params: Backend-specific parameters.
     :type backend_params: dict[str, Any]
+    :param calibration_device: OpenVINO device name to use for calibration inference
+        (e.g. "CPU", "GPU", "GPU.0", "AUTO:GPU,CPU"). If None, defaults to "CPU".
+        Only applicable to the OpenVINO backend.
+    :type calibration_device: Optional[str]
     :param codebook: The codebook (LUT) for the weight compression.
         Applicable for vector quantization. Must be a numpy array or ov Tensor.
     :type codebook: TTensor
@@ -445,6 +456,7 @@ class AdvancedCompressionParameters:
     gptq_params: AdvancedGPTQParameters = field(default_factory=AdvancedGPTQParameters)
     lora_correction_params: AdvancedLoraCorrectionParameters = field(default_factory=AdvancedLoraCorrectionParameters)
     backend_params: dict[str, Any] = field(default_factory=dict)
+    calibration_device: str | None = None
     codebook: TTensor | None = None
     adaptive_codebook_params: AdvancedAdaptiveCodebookParameters = field(
         default_factory=AdvancedAdaptiveCodebookParameters

src/nncf/version.py

@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "3.2.0"
+__version__ = "3.2.0.dev0+00576e031"

tests/openvino/native/quantization/test_quantize_api.py

@@ -35,3 +35,30 @@ def test_non_positive_subset_size():
     with pytest.raises(nncf.ValidationError) as e:
         nncf.quantize(model_to_test, Dataset(MockDataset(INPUT_SHAPE)), subset_size=0)
         assert "Subset size must be positive." in e.info
+
+
+def test_quantize_calibration_device(monkeypatch):
+    import numpy as np
+    import openvino as ov
+
+    from nncf.quantization.advanced_parameters import AdvancedQuantizationParameters
+    from tests.openvino.native.models import LinearModel
+
+    model_to_test = LinearModel().ov_model
+    input_shape = [inp.shape for inp in model_to_test.inputs][0]
+    dataset = Dataset([np.random.rand(*input_shape).astype(np.float32) for _ in range(2)])
+    captured_devices = []
+
+    original_compile = ov.Core.compile_model
+
+    def mock_compile(self, model, device_name="CPU", config=None):
+        captured_devices.append(device_name)
+        return original_compile(self, model, device_name="CPU", config=config)
+
+    monkeypatch.setattr(ov.Core, "compile_model", mock_compile)
+    nncf.quantize(
+        model_to_test,
+        dataset,
+        advanced_parameters=AdvancedQuantizationParameters(calibration_device="GPU"),
+    )
+    assert any(d == "GPU" for d in captured_devices)

tests/openvino/native/quantization/test_weights_compression.py

@@ -2767,3 +2767,27 @@ def test_phi_rope_model(self):
             group_size=-1,
         )
         assert self.get_num_int8_nodes(compressed_model) == 0
+
+
+def test_compress_weights_calibration_device(monkeypatch):
+    model = AWQMatmulModel().ov_model
+    dataset = Dataset([np.ones([2, 8, 8])])
+    captured_devices = []
+
+    original_compile = ov.Core.compile_model
+
+    def mock_compile(self, model, device_name="CPU", config=None):
+        captured_devices.append(device_name)
+        return original_compile(self, model, device_name="CPU", config=config)
+
+    monkeypatch.setattr(ov.Core, "compile_model", mock_compile)
+    compress_weights(
+        model,
+        mode=CompressWeightsMode.INT4_SYM,
+        ratio=1.0,
+        group_size=2,
+        dataset=dataset,
+        awq=True,
+        advanced_parameters=AdvancedCompressionParameters(calibration_device="GPU"),
+    )
+    assert any(d == "GPU" for d in captured_devices)

tests/openvino/native/test_engine.py

@@ -15,6 +15,7 @@
 
 from nncf.definitions import NNCF_DATASET_RESET_STATE_KEY
 from nncf.openvino.engine import OVNativeEngine
+from nncf.openvino.engine import calibration_device_context
 from tests.openvino.native.models import ConvModel
 from tests.openvino.native.models import LinearModel
 from tests.openvino.native.models import QuantizedModel
@@ -123,3 +124,57 @@ def _reset_state():
         "infer",
         "infer",
     ]
+
+
+def test_calibration_device_default(monkeypatch):
+    model = LinearModel().ov_model
+    captured_device = {}
+
+    import openvino as ov
+
+    original_compile = ov.Core.compile_model
+
+    def mock_compile(self, model, device_name="CPU", config=None):
+        captured_device["device_name"] = device_name
+        return original_compile(self, model, device_name="CPU", config=config)
+
+    monkeypatch.setattr(ov.Core, "compile_model", mock_compile)
+    OVNativeEngine(model)
+    assert captured_device["device_name"] == "CPU"
+
+
+def test_calibration_device_context(monkeypatch):
+    model = LinearModel().ov_model
+    captured_device = {}
+
+    import openvino as ov
+
+    original_compile = ov.Core.compile_model
+
+    def mock_compile(self, model, device_name="CPU", config=None):
+        captured_device["device_name"] = device_name
+        return original_compile(self, model, device_name="CPU", config=config)
+
+    monkeypatch.setattr(ov.Core, "compile_model", mock_compile)
+    with calibration_device_context("GPU"):
+        OVNativeEngine(model)
+    assert captured_device["device_name"] == "GPU"
+
+
+def test_calibration_device_context_resets(monkeypatch):
+    model = LinearModel().ov_model
+    captured_devices = []
+
+    import openvino as ov
+
+    original_compile = ov.Core.compile_model
+
+    def mock_compile(self, model, device_name="CPU", config=None):
+        captured_devices.append(device_name)
+        return original_compile(self, model, device_name="CPU", config=config)
+
+    monkeypatch.setattr(ov.Core, "compile_model", mock_compile)
+    with calibration_device_context("GPU"):
+        OVNativeEngine(model)
+    OVNativeEngine(model)
+    assert captured_devices == ["GPU", "CPU"]