[OpenVINO] Introduce a way to control resetting of the state of stateful OV models (#3714)

### Changes

Added `nncf.definitions.NNCF_DATASET_RESET_STATE_KEY` constant to
specify when to reset model state. This constant is used by OpenVINO
backend to control resetting of internal model state between model
inferences. This key can be added to a dataset sample input dictionary
with either `True` or `False` value. With `True` value, the model state
will be reset before inference on the corresponding sample, and with
`False` the state will not be reset.

For an example of usage please see
huggingface/optimum-intel#1505.

### Reason for changes

Without this logic static quantization quality of stateful Whisper
models is poor because a state of a stateful model must be cleared with
the same schedule as it is done during calibration input data
collection.

### Related tickets

172705

### Tests

Added
`tests/openvino/native/test_engine.py::test_stateful_model_inference_with_controlled_resetting`.

Author: nikita-savelyevv

src/nncf/data/dataset.py

@@ -30,6 +30,11 @@ class Dataset:
     usually contains both examples and labels. So transformation function should extract
     the examples from the data item.
 
+    A special input key nncf.definitions.NNCF_DATASET_RESET_STATE_KEY can be used by OpenVINO backend to control
+    resetting of internal model state between model inferences. This key can be added to a dataset sample input
+    dictionary with either `True` or `False` value. With `True` value, the model state will be reset before inference
+    on the corresponding sample, and with `False` the state will not be reset.
+
     :param data_source: The iterable object serving as the source of data items.
     :param transform_func: The function that is used to extract the model's input
         from the data item. The data item here is the data item that is returned from

src/nncf/definitions.py

@@ -23,3 +23,9 @@
 # debug dumps, are performed or not performed
 NNCF_CI_ENV_VAR_NAME = "NNCF_CI"  # Must be set in CI environments
 NNCF_DEV_ENV_VAR_NAME = "NNCF_DEV"  # Must be set in environments of the NNCF dev team machines
+
+# This is a special input key used by OpenVINO backend to control resetting of internal model state
+# between model inferences. This key can be added to a dataset sample input dictionary with either
+# `True` or `False` value. With `True` value, the model state will be reset before inference on the corresponding
+# sample, and with `False` the state will not be reset.
+NNCF_DATASET_RESET_STATE_KEY = "nncf_reset_state"

src/nncf/openvino/engine.py

@@ -17,6 +17,7 @@
 from openvino.properties.hint import inference_precision
 
 from nncf.common.engine import Engine
+from nncf.definitions import NNCF_DATASET_RESET_STATE_KEY
 from nncf.openvino.graph.model_utils import model_has_state
 
 
@@ -43,7 +44,16 @@ def infer(
         :param input_data: Inputs for the model.
         :return output_data: Model's output.
         """
-        if self.reset_state:
+        if isinstance(input_data, dict) and NNCF_DATASET_RESET_STATE_KEY in input_data:
+            # In this case state resetting is controlled by the input data flag
+            input_data = input_data.copy()
+            if input_data.pop(NNCF_DATASET_RESET_STATE_KEY):
+                if self.reset_state:
+                    self.infer_request.reset_state()
+                else:
+                    msg = "Cannot reset state of a stateless model."
+                    raise RuntimeError(msg)
+        elif self.reset_state:
             self.infer_request.reset_state()
 
         model_outputs = self.infer_request.infer(input_data, share_inputs=True)

tests/openvino/native/test_engine.py

@@ -8,10 +8,12 @@
 # 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 functools import wraps
 
 import numpy as np
 import pytest
 
+from nncf.definitions import NNCF_DATASET_RESET_STATE_KEY
 from nncf.openvino.engine import OVNativeEngine
 from tests.openvino.native.models import ConvModel
 from tests.openvino.native.models import LinearModel
@@ -77,3 +79,47 @@ def test_compiled_model_engine_inference_stateful(stateful):
     out = out["Result"]
 
     assert np.array_equal(out[0], input_data[0])
+
+
+def test_stateful_model_inference_with_controlled_resetting():
+    def wrap_reset_state(infer_request):
+        nonlocal reset_order
+        original_reset_state = infer_request.reset_state
+
+        @wraps(infer_request.reset_state)
+        def _reset_state():
+            reset_order.append("reset")
+            original_reset_state()
+
+        infer_request.reset_state = _reset_state
+
+    model = StatefulModel(True).ov_model
+    inp = model.get_parameters()[0]
+    input_data = [{"input_data": np.ones(inp.shape), NNCF_DATASET_RESET_STATE_KEY: False} for _ in range(10)]
+    reset_ind = [2, 5, 7]
+    for ind in reset_ind:
+        input_data[ind][NNCF_DATASET_RESET_STATE_KEY] = True
+
+    engine = OVNativeEngine(model)
+    reset_order = []
+    wrap_reset_state(engine.engine.infer_request)
+
+    for inp_data in input_data:
+        engine.infer(inp_data)
+        reset_order.append("infer")
+
+    assert reset_order == [
+        "infer",
+        "infer",
+        "reset",
+        "infer",
+        "infer",
+        "infer",
+        "reset",
+        "infer",
+        "infer",
+        "reset",
+        "infer",
+        "infer",
+        "infer",
+    ]