Fix cross-process recall: MEMB trailer + raw eval/sample in chat()

Two bugs were blocking memba's main promise (load .memb in a fresh
process → model continues with full recalled context):

1. llama_state_set_data() restores the C-level KV-cache + SSM hidden
   state, but llama-cpp-python's Python wrapper still reports
   n_tokens=0. The next eval() then decodes new tokens at offset 0
   and overwrites the loaded state.

   Fix: extend MEMB format with an optional 12-byte trailer appended
   after the CRC32. It carries the wrapper's n_tokens. The C library
   reads up to CRC and ignores anything past it, so files stay
   backward-compatible with libmemba; only the Python loader uses it.

2. Llama.__call__ / create_chat_completion / generate all re-tokenise
   the prompt on every call and clear the KV-cache when the new
   tokens don't prefix-match input_ids. That destroys any state we
   just loaded.

   Fix: rewrite Session.chat() to use raw tokenize → eval → sample.
   eval() appends tokens to the live state without resetting, and we
   handle stop-token detection ourselves.

Verified end-to-end on Nemotron-3-Nano-4B (hybrid 21x Mamba-2 + 4x
attention) — see experiments/README.md for the full findings log.
diag_session_nemotron.py, mood_batch_poc.py, mood_stream_poc.py and
recall_poc.py now all pass their cross-process tests; Falcon-Mamba
still fails because the trained model itself can't do cross-turn
recall — that was the original misdiagnosis.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
emil
2026-05-16 13:18:31 +03:00
co-authored by Claude Opus 4.7
parent 75c9ee4576
commit ee8a9fc5bc
10 changed files with 424 additions and 34 deletions
+44 -19
View File
@@ -1,32 +1,57 @@
# experiments/
Throwaway scripts used to probe capabilities of SSM models with memba.
Throwaway scripts used to probe SSM/hybrid model capabilities with memba.
Not part of the library API — kept in the repo as reference and reproducible
evidence for product decisions.
Each script is self-contained and prints what it finds; read the source for
the test's claim and run it yourself if you want to verify on different
models or hardware.
Each script is self-contained. Read the source for what it claims to test
and run it yourself if you want to verify on different models or hardware.
## Scripts
| File | What it measures |
|------|------------------|
| `recall_poc.py` | Can a memba state, built from N days of git activity, answer "what did I work on last month" in a fresh process? |
| `mood_poc.py` | Batch sentiment-trajectory test (single prompt with full chat log). Compares Falcon-Mamba vs a Transformer. |
| `mood_stream_poc.py` | The same trajectory but fed turn-by-turn through `Session.chat()`, then queried cross-process. |
| `mood_batch_poc.py` | Batch ingest in build process, save, then query in a fresh process. |
| `diag_saveload.py` | Minimal diagnostic: tell the model one fact, ask it back before save, after save, after cross-process load. |
| `recall_poc.py` | git log → state → cross-process "what did I work on last month" |
| `mood_poc.py` | Batch sentiment trajectory in one prompt — Falcon-Mamba vs Gemma |
| `mood_batch_poc.py` | Batch ingest, save, query in fresh process |
| `mood_stream_poc.py` | 15 separate observation turns, save, query in fresh process |
| `diag_saveload.py` | Minimal hamster recall test on Falcon-Mamba |
| `diag_nemotron.py` | Same hamster test on Nemotron 4B hybrid |
| `diag_nemotron2.py` | Cross-process recall via raw `generate()` (bypass `create_chat_completion`) |
| `diag_nemotron3.py` | Same with llama-cpp-python's native `save_state()`/`load_state()` |
| `diag_session_nemotron.py` | Full hamster recall through the rewritten `Session.chat()` |
## Headline finding (2026-05-16, Falcon-Mamba-7B-Instruct Q4_K_M)
## Findings log
- **Batch single-prompt analysis** (all input + question in one call): works
for both sentiment and recall.
- **Multi-turn fact recall** (ingest in turn 1, ask in turn 2): fails even
in the *same process*. The model does not preserve specific facts in its
hidden state across conversation turns.
- **Save/load roundtrip**: byte-identical, no information loss attributable
to memba's file format. The persistence layer works correctly; the
trained model just doesn't use the state for cross-turn recall.
### 2026-05-16 (initial, Falcon-Mamba-7B-Instruct Q4_K_M)
See the script outputs (or rerun) for the raw evidence.
- Batch single-prompt analysis works for both sentiment and recall.
- Multi-turn fact recall failed even in-process. *Misdiagnosed at first as
a model capability issue.*
### 2026-05-16 (revised, after Nemotron-3-Nano-4B Q4_K_M test)
- Real root cause: `Llama.create_chat_completion()` and `Llama.generate()`
retokenise the entire prompt each call and reset KV-cache when the new
tokens don't prefix-match `input_ids`. Loaded memba state was being wiped.
- Secondary issue: `llama_state_set_data()` restores the C-level cache but
llama-cpp-python's wrapper still reports `n_tokens=0`, so the next
`eval()` writes new tokens at offset 0 and overwrites the loaded state.
- Two fixes in memba:
1. MEMB file now carries a Python-only trailer with `n_tokens` so the
wrapper position is restored after `load_state` (12 extra bytes,
backward-compatible — the C library ignores anything after the CRC).
2. `Session.chat()` rewritten to use raw `tokenize → eval → sample`
instead of `Llama.__call__`, avoiding the prefix-matching reset.
- Results on Nemotron-3-Nano-4B (hybrid: 21 Mamba-2 + 4 attention layers)
through the new `Session.chat()`:
- `diag_session_nemotron.py` — cross-process hamster recall: ✅
- `mood_batch_poc.py` — cross-process mood trajectory: ✅
- `mood_stream_poc.py` — 15 streaming turns + cross-process: ✅
- `recall_poc.py` — per-project digest from 30-day git log: ✅ accurate,
no hallucinated projects.
Falcon-Mamba was rerun through the new `Session.chat()` but still fails on
the same tests — the trained model genuinely doesn't have the cross-turn
recall capability that the hybrid Nemotron does. The 4 attention layers
make the difference.
+95
View File
@@ -0,0 +1,95 @@
"""
diag_nemotron.py — hamster recall test on Nemotron 3 Nano 4B (hybrid Mamba-Transformer).
The Falcon-Mamba 0/4 hamster failure was the killshot for several product ideas.
This rerun tests whether the hybrid architecture (21 Mamba-2 layers + 4 attention)
fixes cross-turn recall.
Three scenarios are measured:
1. In-process multi-turn (tell fact, ask next turn)
2. Same process: save then ask after save
3. Cross process: build (ingest, save, exit), then query (load, ask)
Uses create_chat_completion which applies the GGUF's own chat template.
"""
from __future__ import annotations
import sys, argparse, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from llama_cpp import Llama
from memba import core
MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE = "/tmp/diag_nemotron.memb"
INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named "
"Bartholomew. He is 4 years old. Reply with just 'noted'.")
QUERY = "What is the name of my pet?"
def make_llama():
return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
def chat_continued(m, messages):
"""Send accumulated chat history, return assistant text + cleaned (strip reasoning)."""
out = m.create_chat_completion(
messages=messages,
max_tokens=200,
temperature=0.1,
)
full = out["choices"][0]["message"]["content"].strip()
# Nemotron leaks reasoning — try to extract the final answer if present
short = full[-300:] if len(full) > 300 else full
return full, short
def build():
print(f"[build] loading Nemotron 4B…")
t0 = time.time()
m = make_llama()
print(f"[build] loaded in {time.time()-t0:.1f}s")
messages = [{"role": "user", "content": INGEST}]
full, _ = chat_continued(m, messages)
print(f"[build] ack (full):\n{full!r}\n")
messages.append({"role": "assistant", "content": full})
# Test 1: in-process recall WITHIN the same chat
messages.append({"role": "user", "content": QUERY})
full, short = chat_continued(m, messages)
print(f"[build] in-process query BEFORE save:\n{full}\n")
messages.append({"role": "assistant", "content": full})
# Save state at this point
core.save_state(m, MODEL, STATE)
print(f"[build] saved state ({Path(STATE).stat().st_size:,} B)")
# Test 2: in-process query AFTER save — should still work
messages.append({"role": "user", "content": QUERY})
full, _ = chat_continued(m, messages)
print(f"[build] in-process query AFTER save:\n{full}\n")
def query():
if not Path(STATE).exists():
print("[query] no state — run build first"); return 1
print(f"[query] loading model + state ({Path(STATE).stat().st_size:,} B)…")
t0 = time.time()
m = make_llama()
core.load_state(m, MODEL, STATE)
print(f"[query] loaded in {time.time()-t0:.1f}s")
# Cross-process: send a fresh user turn with only the question
# The state should already encode the prior conversation
out = m.create_chat_completion(
messages=[{"role": "user", "content": QUERY}],
max_tokens=200,
temperature=0.1,
)
print(f"[query] cross-process answer:\n{out['choices'][0]['message']['content']}")
if __name__ == "__main__":
cmd = sys.argv[1] if len(sys.argv) > 1 else "build"
{"build": build, "query": query}[cmd]()
+79
View File
@@ -0,0 +1,79 @@
"""
diag_nemotron2.py — same test, but in cross-process query phase we use
raw __call__ with reset=False to NOT clear KV-cache before generation.
Hypothesis: create_chat_completion() resets KV-cache on each call, which
defeats memba's loaded state. Bypassing that should let state work.
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from llama_cpp import Llama
from memba import core
MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE = "/tmp/diag_nemotron.memb" # reuse state from previous test
def make_llama():
return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
def query_raw():
"""Cross-process query, bypassing create_chat_completion's reset."""
if not Path(STATE).exists():
print("[query] no state — run diag_nemotron.py build first")
return
print(f"[query] loading state ({Path(STATE).stat().st_size:,} B)…")
t0 = time.time()
m = make_llama()
print(f"[query] model loaded in {time.time()-t0:.1f}s")
core.load_state(m, MODEL, STATE)
print(f"[query] state loaded in {time.time()-t0:.1f}s total")
# Continue the conversation: just append a new user turn manually.
# Nemotron uses ChatML-like markers based on what we've seen:
# <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n
# If wrong, model will still answer something interpretable.
continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n"
# Tokenize continuation WITHOUT BOS — we're mid-sequence
tokens = m.tokenize(continuation.encode("utf-8"), add_bos=False, special=True)
print(f"[query] continuation tokenized to {len(tokens)} tokens")
# Use generate() which feeds via eval (does NOT reset KV-cache)
print(f"[query] generating tokens (preserving loaded state)…")
out_tokens = []
eos = m.token_eos()
for token in m.generate(tokens, top_k=1, temp=0.0): # greedy
if token == eos:
break
out_tokens.append(token)
if len(out_tokens) >= 200:
break
# Also stop at <|im_end|>
snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
if "<|im_end|>" in snippet:
break
text = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
print(f"\n[query] answer:\n{text}")
def query_reset_default():
"""Control: same continuation but with default reset=True (clears KV)."""
if not Path(STATE).exists():
return
print(f"\n[control] same continuation, default reset=True (should fail):")
m = make_llama()
core.load_state(m, MODEL, STATE)
out = m(
"<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n",
max_tokens=120, temperature=0.1, echo=False, # reset=True is default
)
print(out["choices"][0]["text"])
if __name__ == "__main__":
query_raw()
# query_reset_default() # optionally enable for direct comparison
+91
View File
@@ -0,0 +1,91 @@
"""
diag_nemotron3.py — test hypothesis that llama-cpp-python's built-in
save_state/load_state preserves Python-side trackers (n_tokens, input_ids)
which memba's raw C-level save/load is missing.
If built-in works → memba's MEMB format needs to be extended to include
those trackers. If built-in also fails → the issue is elsewhere.
"""
import sys, time, pickle
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from llama_cpp import Llama
MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE = "/tmp/diag_nemotron_native.pickle"
INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named "
"Bartholomew. He is 4 years old. Reply with just 'noted'.")
QUERY = "What is the name of my pet?"
def make_llama():
return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
def build():
m = make_llama()
print(f"[build] initial n_tokens={m.n_tokens}")
msgs = [{"role": "user", "content": INGEST}]
out = m.create_chat_completion(messages=msgs, max_tokens=80, temperature=0.1)
ack = out["choices"][0]["message"]["content"]
msgs.append({"role": "assistant", "content": ack})
print(f"[build] after ingest: n_tokens={m.n_tokens}")
print(f"[build] ack snippet: {ack[-100:]!r}")
# Use llama-cpp-python's NATIVE save_state — captures Python trackers too
state = m.save_state()
with open(STATE, "wb") as f:
pickle.dump(state, f)
print(f"[build] saved native state to {STATE} ({Path(STATE).stat().st_size:,} B)")
def query():
if not Path(STATE).exists():
print("[query] no state file"); return
m = make_llama()
print(f"[query] before load: n_tokens={m.n_tokens}")
with open(STATE, "rb") as f:
state = pickle.load(f)
m.load_state(state)
print(f"[query] after load: n_tokens={m.n_tokens}")
# Now ask via create_chat_completion. Pass ONLY the new question
# (history is already in KV-cache+n_tokens).
# The template will format this as if it's turn 1 — see what happens.
out = m.create_chat_completion(
messages=[{"role": "user", "content": QUERY}],
max_tokens=150, temperature=0.1,
)
print(f"\n[query] answer via chat_completion (resets context):\n{out['choices'][0]['message']['content']}")
# Alternative: try to continue manually, using raw eval
print(f"\n[query] re-loading state for raw continuation test")
with open(STATE, "rb") as f:
state = pickle.load(f)
m.load_state(state)
print(f"[query] re-loaded: n_tokens={m.n_tokens}")
# Append a new user turn via raw tokens, then generate
continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n"
tokens = m.tokenize(continuation.encode(), add_bos=False, special=True)
out_tokens = []
eos = m.token_eos()
for tok in m.generate(tokens, top_k=1, temp=0.0):
if tok == eos or len(out_tokens) >= 150:
break
out_tokens.append(tok)
snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
if "<|im_end|>" in snippet:
break
text = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
print(f"\n[query] raw continuation answer:\n{text}")
if __name__ == "__main__":
cmd = sys.argv[1] if len(sys.argv) > 1 else "build"
{"build": build, "query": query}[cmd]()
+32
View File
@@ -0,0 +1,32 @@
"""Same hamster test, but through the new Session.chat() (which uses eval+sample)."""
import sys, argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from memba import Session
MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/diag_session_nemotron"
SESSION = "hamster"
def build():
s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
# Wipe any previous state file so it's a real fresh run
p = Path(STATE_DIR) / f"{SESSION}.memb"
if p.exists(): p.unlink()
print(f"[build] ack: {s.chat('My pet hamster is named Bartholomew. He is 4 years old. Reply noted.', max_tokens=40)!r}")
print(f"[build] in-process query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
s.save()
print(f"[build] saved")
def query():
s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
print(f"[query] loaded state, n_tokens={s._llama.n_tokens}")
print(f"[query] CROSS-PROCESS query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
if __name__ == "__main__":
{"build": build, "query": query}[sys.argv[1]]()
+1 -1
View File
@@ -11,7 +11,7 @@ from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from memba import Session
MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/mood_batch_test"
SESSION = "mood_batch"
+1 -1
View File
@@ -14,7 +14,7 @@ from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from memba import Session
MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/mood_stream_test"
SESSION = "mood_stream"
+1 -1
View File
@@ -34,7 +34,7 @@ SCAN_ROOTS = [
]
DAYS_BACK = 30
STATE_DIR = Path.home() / ".recall/states"
SESSION_ID = "recall_poc"
SESSION_ID = "recall_poc_nemotron"
DEFAULT_Q = (
"Summarise what I've been working on over the last month. "
"Group by project. Mention the main themes per project. "
+28 -1
View File
@@ -25,6 +25,16 @@ MODEL_ID_LEN = 64
_HEADER_FMT = "<4sI64sIIQ"
_HEADER_SIZE = struct.calcsize(_HEADER_FMT) # == 88
# Optional Python trailer (appended after the V1 CRC32). The C library reads
# only up to data_size+CRC and ignores anything past it, so adding a trailer
# keeps the file C-compatible. The trailer carries llama-cpp-python's
# wrapper state (n_tokens) which is needed for `Llama.eval()` to append new
# tokens at the correct position instead of overwriting from offset 0.
_TRAILER_MAGIC = b"MTRL"
_TRAILER_FMT = "<4sII" # magic | version | n_tokens
_TRAILER_SIZE = struct.calcsize(_TRAILER_FMT) # == 12
_TRAILER_VERSION = 1
# ── Helpers ───────────────────────────────────────────────────────
@@ -140,9 +150,14 @@ def save_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
len(data),
)
# Python-only trailer with Llama wrapper position so eval() resumes
# at the right offset instead of overwriting from 0 after load_state.
n_tokens = getattr(llama_model, "n_tokens", 0)
trailer = struct.pack(_TRAILER_FMT, _TRAILER_MAGIC, _TRAILER_VERSION, n_tokens)
out = Path(file_path)
out.parent.mkdir(parents=True, exist_ok=True)
out.write_bytes(header + data + struct.pack("<I", crc))
out.write_bytes(header + data + struct.pack("<I", crc) + trailer)
def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
@@ -187,3 +202,15 @@ def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
restored = _state_set_data(lib, ctx, buf, len(data))
if restored == 0:
raise RuntimeError("llama_state_set_data returned 0 — state restore failed")
# Optional trailer: restore llama-cpp-python wrapper position
trailer_off = offset + data_size + 4
if len(raw) >= trailer_off + _TRAILER_SIZE:
magic, version, n_tokens = struct.unpack_from(_TRAILER_FMT, raw, trailer_off)
if magic == _TRAILER_MAGIC and version == _TRAILER_VERSION:
# Restore Llama wrapper state so eval() appends at the right
# offset instead of overwriting the loaded KV-cache from 0.
try:
llama_model.n_tokens = n_tokens
except AttributeError:
pass # wrapper version doesn't expose .n_tokens — eval may misbehave
+52 -11
View File
@@ -83,25 +83,66 @@ class Session:
# ── Public methods ─────────────────────────────────────────────
def chat(self, prompt: str, max_tokens: int = 512) -> str:
def chat(self, prompt: str, max_tokens: int = 512,
temperature: float = 0.1, top_k: int = 1) -> str:
"""
Feed *prompt* to the model (wrapped in the active chat format) and
return the generated text.
For SSM models the hidden state accumulates in llama_context across
calls there is no explicit message history list, the recurrent
state IS the memory. Call save() at any checkpoint you want to
resume from later.
This uses raw tokenize+eval+sample rather than ``Llama.__call__`` or
``create_chat_completion`` because those high-level helpers re-tokenise
the entire prompt and then reset the KV-cache when the prefix does not
match which destroys any state loaded from a .memb file. Going
through ``llama.eval()`` directly *appends* new tokens to the live
state, which is exactly what memba needs.
For SSM/hybrid models the hidden state accumulates in llama_context
across calls there is no explicit message history list, the recurrent
state IS the memory. Call save() at any checkpoint you want to resume
from later.
"""
template, stops = CHAT_FORMATS[self._chat_format]
wrapped = template.format(prompt=prompt)
result = self._llama(
wrapped,
max_tokens=max_tokens,
echo=False,
stop=stops,
# Tokenise the wrapped turn. Add BOS only when the context is fresh
# (no prior tokens in either the live conversation or a loaded state).
is_fresh = self._llama.n_tokens == 0
new_tokens = self._llama.tokenize(
wrapped.encode("utf-8"),
add_bos=is_fresh,
special=True,
)
return result["choices"][0]["text"].strip()
# Append to the live state. eval() does NOT reset the KV-cache.
self._llama.eval(new_tokens)
# Build the set of single-token stops + add EOS.
stop_token_ids = {self._llama.token_eos()}
for s in stops:
for t in self._llama.tokenize(s.encode("utf-8"), add_bos=False, special=True):
stop_token_ids.add(t)
out_tokens: list[int] = []
for _ in range(max_tokens):
tok = self._llama.sample(top_k=top_k, temp=temperature)
out_tokens.append(tok)
# Always eval the sampled token so it lives in the state too —
# that way the next chat() turn sees the assistant reply as
# part of the conversation.
self._llama.eval([tok])
if tok in stop_token_ids:
break
# Multi-token stop-string check (some stops span several BPE pieces)
if stops:
snippet = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore")
if any(s in snippet for s in stops):
break
text = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore")
for s in stops:
if s in text:
text = text.split(s)[0]
return text.strip()
def save(self, session_id: Optional[str] = None) -> Path:
"""Persist the current state. Returns the path written."""