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>
80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
"""
|
|
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
|