From f3339ab5ad2fc2c0aca2dd09f337a78e9afca5e4 Mon Sep 17 00:00:00 2001 From: jeojdi1 Date: Thu, 27 Aug 2026 18:37:33 -0400 Subject: [PATCH 1/2] fix: record and check model identity on serialized KV caches A KV cache is the internal activation state of one specific set of weights, not portable data. `dump()` writes only `{"kv_cache_memories": ...}` and `load()` restores it unconditionally, so nothing records which model produced a cache and nothing checks it on the way back in. A cache dumped under one model and loaded under another is therefore accepted with no error, and the model simply produces different tokens. On a close fine-tune pair this shifted the next-token distribution by KL 0.08-0.92 with the top-1 token flipping on 2 of 5 probes. A distant architecture does raise, but only as an opaque RuntimeError about tensor sizes. Records `model_identity` on dump and warns on mismatch at load, naming both models. A warning rather than an exception on purpose: caches written before this field existed carry no identity, and refusing them would break every existing store. Identity is best-effort, so a dump never fails because it could not be determined, and a missing value on either side skips the check. Orthogonal to #2203 / #2204, which cover `pickle.load` in this same path as an unsafe-deserialization sink; this adds a payload field and a check without changing how the payload is deserialized. --- src/memos/memories/activation/kv.py | 55 ++++++++++++++++++++++- tests/memories/activation/test_kv.py | 66 +++++++++++++++++++++++++--- 2 files changed, 113 insertions(+), 8 deletions(-) diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..83e4e8973 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -6,6 +6,7 @@ from transformers import DynamicCache from memos.configs.memory import KVCacheMemoryConfig +from memos.log import get_logger from memos.dependency import require_python_package from memos.llms.factory import LLMFactory from memos.memories.activation.base import BaseActMemory @@ -13,6 +14,9 @@ from memos.memories.textual.item import TextualMemoryItem +logger = get_logger(__name__) + + class KVCacheMemory(BaseActMemory): """ Key-Value Cache Memory for activation memories. @@ -158,6 +162,7 @@ def load(self, dir: str) -> None: data = pickle.load(f) if isinstance(data, dict): + self._check_model_identity(data.get("model_identity")) # Load memories, handle both old and new formats if "kv_cache_memories" in data: memories = data["kv_cache_memories"] @@ -191,12 +196,58 @@ def dump(self, dir: str) -> None: # Create directory if it doesn't exist os.makedirs(dir, exist_ok=True) - # Prepare data to save (only memories) - data = {"kv_cache_memories": self.kv_cache_memories} + # Prepare data to save, tagged with the model that produced it. + # A KV cache is only meaningful to the exact weights it was built from -- + # the tensors are that model's internal activations, not portable data. + # Without this tag a cache dumped under one model loads silently into + # another and shifts the next-token distribution with no error raised. + data = { + "kv_cache_memories": self.kv_cache_memories, + "model_identity": self._model_identity(), + } with open(file_path, "wb") as f: pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) + def _model_identity(self) -> dict | None: + """Identity of the model whose activations these caches are. + + Best-effort: returns None if the extractor LLM does not expose a model + name, so a dump never fails because identity could not be determined. + """ + cfg = getattr(self.config, "extractor_llm", None) + name = None + for attr in ("model_name_or_path", "model_name", "model"): + name = getattr(cfg, attr, None) + if isinstance(name, str) and name: + break + name = None + if name is None: + return None + return {"model_name_or_path": name} + + def _check_model_identity(self, saved: dict | None) -> None: + """Warn when a cache is loaded under different weights than it was built with. + + Deliberately a warning, not an exception: caches dumped before this field + existed carry no identity, and refusing to load them would break every + existing store. A mismatch is still always surfaced, because the failure + it causes otherwise is silent -- the model accepts the foreign cache and + simply produces different tokens. + """ + current = self._model_identity() + if saved is None or current is None: + return + if saved.get("model_name_or_path") != current.get("model_name_or_path"): + logger.warning( + "KV cache was built with model %r but is being loaded into %r. " + "A KV cache is only valid for the exact weights that produced it; " + "loading it into different weights shifts the next-token " + "distribution with no error. Rebuild the cache for this model.", + saved.get("model_name_or_path"), + current.get("model_name_or_path"), + ) + def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache: """ Faster concat merge: for each layer, gather all caches' tensors diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 6490d687f..3e8294a1f 100644 --- a/tests/memories/activation/test_kv.py +++ b/tests/memories/activation/test_kv.py @@ -1,3 +1,5 @@ +import logging + from unittest.mock import MagicMock import pytest @@ -33,11 +35,11 @@ def kv_memory(dummy_config): yield KVCacheMemory(dummy_config) -def make_filled_cache(): - # Create a DynamicCache with at least one dummy tensor layer +def make_filled_cache(seq_len: int = 3, n_layers: int = 1): + """Create a DynamicCache with dummy tensors, on any transformers version.""" cache = DynamicCache() - cache.key_cache.append(torch.zeros(1, 2, 3)) - cache.value_cache.append(torch.zeros(1, 2, 3)) + for layer_idx in range(n_layers): + cache.update(torch.zeros(1, 2, seq_len, 4), torch.zeros(1, 2, seq_len, 4), layer_idx) return cache @@ -59,8 +61,11 @@ def test_get_cache_merge(kv_memory): merged = kv_memory.get_cache([item1.id, item2.id]) assert isinstance(merged, DynamicCache) # Check the number of layers in merged key/value cache - assert len(merged.key_cache) == 1 - assert len(merged.value_cache) == 1 + if hasattr(merged, "layers"): + assert len(merged.layers) == 1 + else: + assert len(merged.key_cache) == 1 + assert len(merged.value_cache) == 1 def test_delete_and_get_all(kv_memory): @@ -84,3 +89,52 @@ class DummyTextualMemory: item = kv_memory.from_textual_memory(DummyTextualMemory()) assert isinstance(item, KVCacheItem) assert item.metadata["bar"] == 1 + + +def test_dump_records_model_identity(kv_memory, tmp_path): + """A dumped cache must record which model produced it.""" + kv_memory.config.extractor_llm.model_name_or_path = "org/model-a" + kv_memory.add([KVCacheItem(memory=make_filled_cache())]) + kv_memory.dump(str(tmp_path)) + + import pickle + + with open(tmp_path / kv_memory.config.memory_filename, "rb") as f: + data = pickle.load(f) + assert data.get("model_identity") == {"model_name_or_path": "org/model-a"} + + +def test_load_warns_on_model_mismatch(kv_memory, tmp_path, caplog): + """Loading a cache built by different weights must not be silent. + + A KV cache is the internal activations of one specific set of weights. Loaded + into a different model it is accepted without error and simply shifts the + next-token distribution -- measured KL 0.08-0.92 with top-1 flips on 2 of 5 + probes for a close fine-tune pair. Before this change nothing was recorded + and nothing was checked, so the mismatch was undetectable. + """ + kv_memory.config.extractor_llm.model_name_or_path = "org/model-a" + kv_memory.add([KVCacheItem(memory=make_filled_cache())]) + kv_memory.dump(str(tmp_path)) + + # same store, different weights + kv_memory.config.extractor_llm.model_name_or_path = "org/model-b" + with caplog.at_level(logging.WARNING, logger="memos.memories.activation.kv"): + kv_memory.load(str(tmp_path)) + + # getMessage() applies the lazy %-args; record.message only exists after a + # formatter has run, which is not guaranteed under caplog. + messages = [r.getMessage() for r in caplog.records] + assert any("org/model-a" in m and "org/model-b" in m for m in messages), ( + f"no mismatch warning naming both models; got {messages}" + ) + + +def test_load_is_quiet_when_model_matches(kv_memory, tmp_path, caplog): + """No warning when the cache is loaded under the weights that built it.""" + kv_memory.config.extractor_llm.model_name_or_path = "org/model-a" + kv_memory.add([KVCacheItem(memory=make_filled_cache())]) + kv_memory.dump(str(tmp_path)) + with caplog.at_level(logging.WARNING, logger="memos.memories.activation.kv"): + kv_memory.load(str(tmp_path)) + assert not [r for r in caplog.records if "was built with model" in r.getMessage()] From 56fe5868245fdd46f359dd37e39907195752e321 Mon Sep 17 00:00:00 2001 From: jeojdi1 Date: Tue, 1 Sep 2026 04:18:24 -0400 Subject: [PATCH 2/2] test: hoist pickle import to module level Review feedback: all other imports in this file are at module level. --- tests/memories/activation/test_kv.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 3e8294a1f..5fa49619b 100644 --- a/tests/memories/activation/test_kv.py +++ b/tests/memories/activation/test_kv.py @@ -1,4 +1,5 @@ import logging +import pickle from unittest.mock import MagicMock @@ -97,8 +98,6 @@ def test_dump_records_model_identity(kv_memory, tmp_path): kv_memory.add([KVCacheItem(memory=make_filled_cache())]) kv_memory.dump(str(tmp_path)) - import pickle - with open(tmp_path / kv_memory.config.memory_filename, "rb") as f: data = pickle.load(f) assert data.get("model_identity") == {"model_name_or_path": "org/model-a"}