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,162 @@
+"""Stress the durable Natural Memory v2 page tier without loading Qwen."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import torch
+
+if __package__ in {None, ""}:
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from V2_dpskw.memory_os_v2 import MemoryRouterV2, PagedMemoryBankV2
+from V2_dpskw.tiered_memory_store_v2 import TieredMemoryStoreV2
+
+
+def run(args: argparse.Namespace) -> dict[str, object]:
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    report_path = Path(args.output)
+    report_path.parent.mkdir(parents=True, exist_ok=True)
+    with TemporaryDirectory(prefix="natural-memory-v2-tiered-", dir=str(report_path.parent)) as temp_dir:
+        store_path = Path(temp_dir) / "memory.sqlite"
+        router = MemoryRouterV2(args.hidden_size, router_dim=args.key_dim, num_heads=args.heads)
+        store = TieredMemoryStoreV2(
+            store_path,
+            key_dim=args.key_dim,
+            page_capacity=args.page_capacity,
+        )
+        bank = PagedMemoryBankV2(
+            args.hidden_size,
+            page_capacity=args.page_capacity,
+            max_pages=max(1, (args.records + args.page_capacity - 1) // args.page_capacity + 8),
+            hot_pages=args.hot_pages,
+            top_k_pages=args.top_k_pages,
+            top_k_records=args.top_k_records,
+            router=router,
+            key_dim=args.key_dim,
+            coarse_index_bits=args.coarse_index_bits,
+            tier_store=store,
+            max_resident_pages=args.resident_pages,
+        )
+        target_key = None
+        target_text = ""
+        started = time.perf_counter()
+        for start in range(0, args.records, args.batch_size):
+            batch: list[dict[str, object]] = []
+            for index in range(start, min(args.records, start + args.batch_size)):
+                key = torch.randn(args.hidden_size)
+                if index == args.target_index:
+                    target_key = key.clone()
+                    target_text = f"tiered-record-{index}"
+                batch.append(
+                    {
+                        "text": f"tiered-record-{index}",
+                        "key": key,
+                        "summary": key,
+                        "entity": "benchmark",
+                        "attribute": f"attribute-{index}",
+                        "value": f"value-{index}",
+                        "importance": 0.2 if index != args.target_index else 1.0,
+                        "confidence": 0.95,
+                        "source": "tiered-benchmark",
+                        "trusted": True,
+                    }
+                )
+            bank.write_batch(batch)
+        write_seconds = time.perf_counter() - started
+        before = bank.stats()
+        store.close()
+
+        reopen_started = time.perf_counter()
+        reopened_store = TieredMemoryStoreV2(
+            store_path,
+            key_dim=args.key_dim,
+            page_capacity=args.page_capacity,
+        )
+        reopened = PagedMemoryBankV2(
+            args.hidden_size,
+            page_capacity=args.page_capacity,
+            max_pages=max(1, (args.records + args.page_capacity - 1) // args.page_capacity + 8),
+            hot_pages=args.hot_pages,
+            top_k_pages=args.top_k_pages,
+            top_k_records=args.top_k_records,
+            router=router,
+            key_dim=args.key_dim,
+            coarse_index_bits=args.coarse_index_bits,
+            tier_store=reopened_store,
+            max_resident_pages=args.resident_pages,
+        )
+        reopen_seconds = time.perf_counter() - reopen_started
+        if target_key is None:
+            raise RuntimeError("target index was not generated")
+        records, decision = reopened.query(
+            query_key=target_key,
+            query_text=target_text,
+            top_k_pages=args.top_k_pages,
+            top_k_records=args.top_k_records,
+        )
+        after = reopened.stats()
+        found = any(record.text == target_text for record in records)
+        target_row = reopened_store.find_by_text(target_text, active_status="active")
+        target_page_id = target_row["page_id"] if target_row is not None else None
+        candidate_pages = reopened._candidate_page_ids(target_key)
+        reopened_store.close()
+
+    report = {
+        "format_version": 2,
+        "records_requested": args.records,
+        "target_index": args.target_index,
+        "page_capacity": args.page_capacity,
+        "coarse_index_bits": args.coarse_index_bits,
+        "resident_pages": args.resident_pages,
+        "write_seconds": write_seconds,
+        "reopen_seconds": reopen_seconds,
+        "before_close": before,
+        "after_reopen": after,
+        "target_recalled_after_restart": found,
+        "target_page_id": target_page_id,
+        "target_page_in_coarse_candidates": target_page_id in candidate_pages if target_page_id else False,
+        "coarse_candidate_count": len(candidate_pages),
+        "decision": {
+            "page_ids": decision.page_ids,
+            "record_ids": decision.record_ids,
+            "hop_count": decision.hop_count,
+            "confidence": decision.confidence,
+            "stop_reason": decision.stop_reason,
+        },
+    }
+    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
+    return report
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--records", type=int, default=100_000)
+    parser.add_argument("--target-index", type=int, default=99_999)
+    parser.add_argument("--hidden-size", type=int, default=32)
+    parser.add_argument("--key-dim", type=int, default=16)
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--page-capacity", type=int, default=32)
+    parser.add_argument("--resident-pages", type=int, default=64)
+    parser.add_argument("--hot-pages", type=int, default=8)
+    parser.add_argument("--top-k-pages", type=int, default=4)
+    parser.add_argument("--top-k-records", type=int, default=8)
+    parser.add_argument("--coarse-index-bits", type=int, default=12)
+    parser.add_argument("--batch-size", type=int, default=2_000)
+    parser.add_argument("--seed", type=int, default=20260904)
+    parser.add_argument(
+        "--output",
+        default="W:/Flash/model/V2_dpskw/tiered_memory_v2_benchmark.json",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    print(json.dumps(run(parse_args()), ensure_ascii=False, indent=2))