fix: record and check model identity on serialized KV caches - #2302
fix: record and check model identity on serialized KV caches#2302jeojdi1 wants to merge 1 commit into
Conversation
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 MemTensor#2203 / MemTensor#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.
Memtensor-AI
added
area:memory
🤖 Open Code ReviewTarget: PR #2302 🔍 OpenCodeReview found 2 issue(s) in this PR. 1.
|
|
4 participants
Description
A KV cache is the internal activation state of one specific set of weights — it is not portable data.
KVCacheMemory.dump()currently writes only{"kv_cache_memories": ...}, andload()restores it unconditionally. There is nothing recording which model produced a cache and nothing checking it on the way back in.The consequence is silent: a cache dumped under one model and loaded under another is accepted with no error, and the model simply produces different tokens. On a close fine-tune pair I measured a next-token KL shift of 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
RuntimeErrorabout tensor sizes, which does not tell the user what actually went wrong.This PR:
model_identityin the dumped payload (model_name_or_path, best-effort from the configured extractor LLM);load()and warns on mismatch, naming both models and saying what to do about it.It is a warning rather than an exception on purpose: caches dumped before this field existed carry no identity, and refusing to load them would break every existing store.
_model_identity()returnsNonewhen the LLM config exposes no name, so a dump never fails because identity could not be determined, and aNoneon either side skips the check.Related Issue (Required): #2300
Note on overlap: issue #2203 (closed) and PR #2204 cover
pickle.loadin this sameload()path as an unsafe-deserialization sink (CWE-502). This change is orthogonal — it adds a payload field and a check, and does not touch how the payload is deserialized — but it does edit adjacent lines, so it may need a trivial rebase depending on merge order.Type of change
How Has This Been Tested?
Three tests added:
test_dump_records_model_identity— the field is written.test_load_warns_on_model_mismatch— dump underorg/model-a, load underorg/model-b, assert the warning names the original model. Fails on the current code, which emits nothing.test_load_is_quiet_when_model_matches— no warning on the matching case, so this is not a new source of noise.Checklist
A note on the target branch
CONTRIBUTING.mdsays to open PRs againstdev, but nodevbranch exists — onlymainanddev-v2.0.28…dev-v2.0.32. This is againstmain(185ebdb, "Dev v2.0.32"). Happy to retarget.