Key OpenVINO encoder exports by weights and configuration (#20)

Author: aalquwayfili

src/von/backends/option_marker_backend.py

@@ -6,6 +6,7 @@
 - Score: Single-pass ordinal rating over all levels simultaneously.
 """
 
+import hashlib
 import json
 import math
 import os
@@ -248,12 +249,26 @@ def _ov_core_with_cache(target: str):
     return core, cache_dir
 
 
+def _openvino_encoder_cache_key(encoder: torch.nn.Module) -> str:
+    import openvino as ov
+
+    # A release name or checkpoint path does not identify fine-tuned/replaced weights.
+    digest = hashlib.sha256()
+    digest.update(f"{type(encoder).__module__}.{type(encoder).__qualname__}:{torch.__version__}:{ov.__version__}".encode())
+    digest.update(encoder.config.to_json_string(use_diff=False).encode())
+    for name, tensor in sorted(encoder.state_dict().items()):
+        digest.update(f"{name}:{tensor.dtype}:{tuple(tensor.shape)}".encode())
+        # View bytes directly: supports bfloat16 without allocating a second weights blob.
+        digest.update(tensor.detach().cpu().contiguous().reshape(-1).view(torch.uint8).numpy())
+    return digest.hexdigest()
+
+
 def _compile_openvino_encoder(encoder: torch.nn.Module, target: str = "GPU") -> torch.nn.Module:
     """Compiles the PyTorch ModernBERT encoder (plain mode) to OpenVINO with disk caching."""
     import openvino as ov
 
     core, cache_dir = _ov_core_with_cache(target)
-    xml_path = os.path.join(cache_dir, f"encoder_{VON_MODEL_ID}.xml")
+    xml_path = os.path.join(cache_dir, f"encoder_{VON_MODEL_ID}_{_openvino_encoder_cache_key(encoder)}.xml")
     if os.path.exists(xml_path):
         compiled_model = core.compile_model(xml_path, device_name=target)
     else:
@@ -280,7 +295,7 @@ def _compile_openvino_independent_encoder(encoder: torch.nn.Module, target: str
     import openvino as ov
 
     core, cache_dir = _ov_core_with_cache(target)
-    xml_path = os.path.join(cache_dir, f"encoder_indep_{VON_MODEL_ID}.xml")
+    xml_path = os.path.join(cache_dir, f"encoder_indep_{VON_MODEL_ID}_{_openvino_encoder_cache_key(encoder)}.xml")
     if os.path.exists(xml_path):
         compiled_model = core.compile_model(xml_path, device_name=target)
     else:

tests/test_openvino_cache.py

@@ -0,0 +1,74 @@
+"""Exercise cached exports with tiny encoders, without downloading weights."""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput
+
+from von.backends import option_marker_backend as backend
+
+ov = pytest.importorskip("openvino")
+
+
+class TinyEncoder(torch.nn.Module):
+    def __init__(self, weight, scale=1.0):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.tensor([weight], dtype=torch.float32))
+        self.config = PretrainedConfig(scale=scale)
+
+    def forward(self, input_ids, attention_mask, position_ids=None):
+        if isinstance(attention_mask, dict):
+            mask = attention_mask["full_attention"].sum(-1).squeeze(1)
+            mask = mask + attention_mask["sliding_attention"].sum(-1).squeeze(1)
+            values = input_ids.float() + position_ids.float() + mask
+        else:
+            values = input_ids.float() + attention_mask.float()
+        return BaseModelOutput(last_hidden_state=(values * self.weight * self.config.scale).unsqueeze(-1))
+
+
+@pytest.mark.parametrize("independent", [False, True])
+def test_cached_encoder_tracks_weights_and_config(tmp_path, monkeypatch, independent):
+    core = ov.Core()
+    core.set_property("CPU", {"CACHE_DIR": str(tmp_path)})
+    monkeypatch.setattr(backend, "_ov_core_with_cache", lambda target: (core, str(tmp_path)))
+    compile_encoder = (backend._compile_openvino_independent_encoder if independent
+                       else backend._compile_openvino_encoder)
+    inputs = {"input_ids": torch.ones((1, 128), dtype=torch.long)}
+    if independent:
+        mask = torch.ones((1, 1, 128, 128), dtype=torch.bool)
+        inputs.update(attention_mask={"full_attention": mask, "sliding_attention": mask},
+                      position_ids=torch.arange(128).unsqueeze(0))
+    else:
+        inputs["attention_mask"] = torch.ones_like(inputs["input_ids"])
+
+    def check(encoder):
+        compiled = compile_encoder(encoder, target="CPU")
+        with torch.no_grad():
+            expected_inputs = dict(inputs)
+            if independent:
+                expected_inputs["attention_mask"] = {
+                    name: backend._bool_to_additive(mask)
+                    for name, mask in inputs["attention_mask"].items()
+                }
+            expected = encoder(**expected_inputs).last_hidden_state
+            torch.testing.assert_close(compiled(**inputs).last_hidden_state, expected)
+
+    with patch.object(ov, "convert_model", wraps=ov.convert_model) as convert:
+        check(TinyEncoder(2.0).eval())
+        check(TinyEncoder(2.0).eval())
+        assert convert.call_count == 1, "identical encoders should reuse the export"
+
+        encoder = TinyEncoder(3.0).eval()
+        check(encoder)
+        assert convert.call_count == 2, "different checkpoint weights need a new export"
+
+        with torch.no_grad():
+            encoder.weight.fill_(4.0)
+        check(encoder)
+        assert convert.call_count == 3, "replacing weights in place must invalidate the export"
+
+        encoder.config.scale = 2.0
+        check(encoder)
+        assert convert.call_count == 4, "configuration changes must invalidate the export"