Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据

- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线 - 修复 prepare_memory_router_dataset 候选池重建缺陷（mega 家族 3568x 加速，输出逐字节相同） - 修复 v5 被破坏的拒答与多跳标签（train 未知样本 319 -> 16319，multi_hop 平均正例 1.00 -> 2.00） - 同存储预算下 V2-128 v6 逐轴 22/22 通过：Top-1 41.12% -> 94.62%，未知拒答 0.00% -> 100.00% - 记录三条被实测推翻的显然优化（logits_to_keep=1 反而慢 55%、XL 容量未带来收益） - 记忆手术跨架构可移植性 14/14，读写关闭时与原生模型逐位相同
2026-09-19 11:11:31 +08:00
commit 643e22ecb9
484 changed files with 306821 additions and 0 deletions
@@ -0,0 +1,206 @@
+"""General-capability regression with the memory read path ACTIVE.
+
+``comprehensive_benchmark.py`` runs its adapted model without
+``memory_query_*``, so the router never executes: it verifies that the memory
+*architecture* does not disturb generation, not that the router behaves on
+unrelated questions.  In production the router runs on every turn, so the real
+risk is a populated memory bank polluting a general question.
+
+This harness therefore measures the harder condition:
+
+1. a plain Qwen baseline with no memory at all;
+2. the shipped package with a **deliberately populated** bank of irrelevant
+   personal facts, the read path active, for each router under test;
+3. per-category and overall scores plus how often the router chose to read.
+
+A correct router abstains on arithmetic/logic/code prompts, so the memory-enabled
+scores must match the no-memory baseline.
+
+Usage::
+
+    python -m V2_dpskw.eval_general_capability ^
+        --data <comprehensive_general.jsonl> ^
+        --router "deployed=" ^
+        --router "V2-128-v6=checkpoints/router_v6_v2_128/router_best.pt" ^
+        --output router_general_capability.json --markdown router_general_capability.md
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+import time
+import unicodedata
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+
+if __package__ in {None, ""}:
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from V2_dpskw.eval_end_to_end_memory import _chat_tensor, answer, write_fact
+from V2_dpskw.qwen_integration import (
+    load_memory_config,
+    load_qwen_base,
+    load_qwen_dynamic,
+    load_tokenizer,
+)
+
+#: Unrelated personal facts used to populate the bank before general questions.
+DISTRACTOR_FACTS = [
+    "我的工作地点是 V-8QPA0VMT。",
+    "我的档案代号是 V-2RPEJ33G。",
+    "我住在杭州，通勤方式是地铁。",
+    "我的项目编号是 PROJECT-M3GWABOJ。",
+    "我的默认语言是中文。",
+    "我常用的编辑器是 VS Code。",
+    "我的工位在 3 号楼 4 层。",
+    "我的邮箱后缀是 example.com。",
+    "我的常用编程语言是 Python。",
+    "我的生日是 3 月 12 日。",
+    "我的主管是李工。",
+    "我的团队名称是记忆系统组。",
+]
+
+
+def normalize(value: str) -> str:
+    value = unicodedata.normalize("NFKC", str(value)).lower()
+    return re.sub(r"\s+", "", value)
+
+
+def contains_answer(text: str, acceptable: list[str]) -> bool:
+    """Same scoring rule as comprehensive_benchmark.py."""
+
+    normalized = normalize(text)
+    for expected in acceptable:
+        expected = normalize(str(expected))
+        if not expected:
+            continue
+        if expected.isdigit() and len(expected) == 1:
+            if re.search(rf"(?<!\d){re.escape(expected)}(?!\d)", normalized):
+                return True
+        elif expected in normalized:
+            return True
+    return False
+
+
+def load_cases(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+@torch.inference_mode()
+def run_cases(model, tokenizer, cases, device, *, max_new_tokens: int, use_memory: bool) -> dict:
+    per_category: dict[str, list[float]] = defaultdict(list)
+    read_decisions = 0
+    rows = []
+    started = time.perf_counter()
+    for case in cases:
+        encoded = {key: value.to(device) for key, value in _chat_tensor(tokenizer, str(case["prompt"])).items()}
+        kwargs: dict = {}
+        if use_memory:
+            query_tokens = tokenizer(str(case["prompt"]), add_special_tokens=False, return_tensors="pt")
+            kwargs = {
+                "update_memory": False,
+                "memory_query_input_ids": query_tokens["input_ids"].to(device),
+                "memory_query_attention_mask": query_tokens.get(
+                    "attention_mask", torch.ones_like(query_tokens["input_ids"])
+                ).to(device),
+                "memory_query_text": str(case["prompt"]),
+            }
+        output = model.generate(
+            **encoded,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            use_cache=True,
+            pad_token_id=tokenizer.pad_token_id,
+            **kwargs,
+        )
+        generated = output[0][encoded["input_ids"].shape[1]:] if isinstance(output, torch.Tensor) else output
+        text = tokenizer.decode(generated, skip_special_tokens=True).strip()
+        passed = contains_answer(text, list(case["acceptable"]))
+        per_category[str(case["category"])].append(float(passed))
+        decided = False
+        if use_memory and model.runtime.v2_last_decisions:
+            decided = bool(model.runtime.v2_last_decisions[-1].get("need_memory"))
+            read_decisions += int(decided)
+        rows.append({"id": case["id"], "category": case["category"], "passed": passed,
+                     "read_memory": decided, "reply": text[:160]})
+    summary = {
+        "cases": len(rows),
+        "overall_pct": 100 * sum(r["passed"] for r in rows) / max(1, len(rows)),
+        "per_category_pct": {name: 100 * sum(values) / max(1, len(values)) for name, values in sorted(per_category.items())},
+        "read_decisions": read_decisions,
+        "seconds": round(time.perf_counter() - started, 1),
+    }
+    return {"summary": summary, "rows": rows}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--package", default="qwen3_5_4b_natural_memory_v2")
+    parser.add_argument("--data", required=True, help="comprehensive_general.jsonl")
+    parser.add_argument("--router", action="append", required=True, help="LABEL=PATH ('LABEL=' keeps the deployed router)")
+    parser.add_argument("--max-new-tokens", type=int, default=32)
+    parser.add_argument("--facts", type=int, default=len(DISTRACTOR_FACTS))
+    parser.add_argument("--output", default="router_general_capability.json")
+    parser.add_argument("--markdown", default="")
+    args = parser.parse_args()
+
+    cases = load_cases(Path(args.data))
+    print(json.dumps({"cases": len(cases), "categories": len({c['category'] for c in cases}),
+                      "distractor_facts": args.facts}, ensure_ascii=False), flush=True)
+
+    model_path = Path(args.package)
+    tokenizer = load_tokenizer(model_path)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    results: dict = {}
+
+    baseline = load_qwen_base(model_path, load_in_4bit=True, max_memory={0: "10.5GiB", "cpu": "48GiB"})
+    baseline.eval()
+    results["Qwen3.5-4B baseline (no memory)"] = run_cases(
+        baseline, tokenizer, cases, device, max_new_tokens=args.max_new_tokens, use_memory=False)
+    print(json.dumps({"baseline": results["Qwen3.5-4B baseline (no memory)"]["summary"]}, ensure_ascii=False), flush=True)
+    del baseline
+    torch.cuda.empty_cache()
+
+    memory_config = load_memory_config(model_path)
+    model = load_qwen_dynamic(model_path, memory_config=memory_config, load_in_4bit=True,
+                             max_memory={0: "10.5GiB", "cpu": "48GiB"})
+    model.eval()
+    for spec in args.router:
+        label, _, path_value = spec.partition("=")
+        if path_value.strip():
+            state = torch.load(path_value.strip(), map_location="cpu", weights_only=True)
+            model.memory_router_v2.load_state_dict(state.get("router_state_dict", state), strict=True)
+            model.memory_router_v2.to(device).eval()
+        model.reset_memory(batch_size=1, device=device)
+        written = sum(int(write_fact(model, tokenizer, text, device)) for text in DISTRACTOR_FACTS[: args.facts])
+        entry = run_cases(model, tokenizer, cases, device, max_new_tokens=args.max_new_tokens, use_memory=True)
+        entry["summary"]["facts_written"] = written
+        entry["summary"]["router"] = label
+        results[label] = entry
+        print(json.dumps(entry["summary"], ensure_ascii=False), flush=True)
+
+    Path(args.output).write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
+    base = results["Qwen3.5-4B baseline (no memory)"]["summary"]
+    lines = ["| 条件 | 总体 | " + " | ".join(base["per_category_pct"]) + " | 触发读取次数 |",
+             "|---" * (len(base["per_category_pct"]) + 3) + "|"]
+    for label, block in results.items():
+        s = block["summary"]
+        cells = " | ".join(f"{s['per_category_pct'].get(name, float('nan')):.2f}%" for name in base["per_category_pct"])
+        delta = s["overall_pct"] - base["overall_pct"]
+        lines.append(f"| {label} | {s['overall_pct']:.2f}% ({delta:+.2f}pp) | {cells} | {s.get('read_decisions', 0)}/{s['cases']} |")
+    table = "\n".join(lines)
+    print("\n" + table, flush=True)
+    if args.markdown:
+        Path(args.markdown).write_text(table + "\n", encoding="utf-8")
+        print(f"wrote {args.markdown}", flush=True)
+    print(f"wrote {args.output}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())