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，读写关闭时与原生模型逐位相同

probe_encoder_precision.py

@@ -0,0 +1,209 @@
+"""Probe whether a faster encoder precision reproduces the frozen 4-bit features.
+
+The v6 feature encode needs ~3 hours because the 4B model runs in 4-bit NF4,
+which dequantises every weight on every forward.  A bf16 forward is usually much
+faster, but the deployed routers were trained on 4-bit features, so switching
+precision is only acceptable if it does not move the *downstream* numbers.
+
+This tool answers that empirically instead of by assumption:
+
+1. rebuild the exact text list of an existing cache deterministically from the
+   dataset (and verify the ``sha1`` set matches the reference ``index.json``);
+2. encode it with a chosen precision into a second cache;
+3. report feature-level agreement (cosine) against the reference bank.
+
+Feeding both caches to ``eval_router_scorecard`` then shows whether the router
+metrics - and therefore any conclusion drawn from them - survive the change.
+
+Usage::
+
+    python -m V2_dpskw.probe_encoder_precision ^
+        --train-file data/router_training_v3/train.jsonl ^
+        --eval-file data/router_training_v3/eval.jsonl ^
+        --reference-cache checkpoints/router_shared/feature_cache ^
+        --output-dir H:\\Memory\\nm_cache\\nm_probe_bf16 --precision bf16
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import hashlib
+import json
+import sys
+import time
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+
+if __package__ in {None, ""}:
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from V2_dpskw.qwen_integration import load_qwen_dynamic, load_tokenizer
+
+
+def _resolve(value: str | Path) -> Path:
+    path = Path(value)
+    if path.is_absolute() or path.exists():
+        return path
+    return Path(__file__).resolve().parent / path
+
+
+def _text_key(text: str) -> str:
+    return hashlib.sha1(str(text).encode("utf-8", errors="replace")).hexdigest()
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def collect_texts(paths: list[Path]) -> tuple[list[str], dict[str, int]]:
+    """Reproduce the trainer's text order: first appearance across train+eval."""
+
+    texts: list[str] = []
+    lookup: dict[str, int] = {}
+    for path in paths:
+        with path.open("r", encoding="utf-8") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                row = json.loads(line)
+                values = [str(row.get("query", ""))]
+                for candidate in row.get("candidates") or []:
+                    if isinstance(candidate, dict):
+                        values.append(str(candidate.get("text", "")))
+                for text in values:
+                    text = text.strip()
+                    if not text:
+                        continue
+                    key = _text_key(text)
+                    if key not in lookup:
+                        lookup[key] = len(texts)
+                        texts.append(text)
+    return texts, lookup
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train-file", default="data/router_training_v3/train.jsonl")
+    parser.add_argument("--eval-file", default="data/router_training_v3/eval.jsonl")
+    parser.add_argument("--reference-cache", default="checkpoints/router_shared/feature_cache")
+    parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--precision", choices=("bf16", "4bit"), default="bf16")
+    parser.add_argument("--max-key-tokens", type=int, default=256)
+    parser.add_argument("--max-batch", type=int, default=128)
+    parser.add_argument("--token-budget", type=int, default=8192)
+    parser.add_argument("--gpu-memory-gb", type=float, default=10.5)
+    args = parser.parse_args()
+
+    train_path = _resolve(args.train_file)
+    eval_path = _resolve(args.eval_file)
+    reference_dir = _resolve(args.reference_cache)
+    output_dir = _resolve(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    reference_index = json.loads((reference_dir / "index.json").read_text(encoding="utf-8"))
+    texts, lookup = collect_texts([train_path, eval_path])
+    print(json.dumps({
+        "phase": "texts",
+        "collected": len(texts),
+        "reference": len(reference_index),
+        "sha_set_identical": set(lookup) == set(reference_index),
+        "row_order_identical": lookup == reference_index,
+    }, ensure_ascii=False), flush=True)
+
+    if args.gpu_memory_gb > 0 and torch.cuda.is_available():
+        total = torch.cuda.get_device_properties(0).total_memory
+        torch.cuda.set_per_process_memory_fraction(min(0.95, args.gpu_memory_gb * 1024**3 / total), device=0)
+
+    tokenizer = load_tokenizer(_resolve(args.model_path))
+    encoded = tokenizer(texts, add_special_tokens=False, truncation=True, max_length=args.max_key_tokens, padding=False)["input_ids"]
+    lengths = np.array([max(1, min(len(ids), args.max_key_tokens)) for ids in encoded], dtype=np.int64)
+    order = np.argsort(lengths, kind="stable")
+    print(json.dumps({"phase": "tokenized", "texts": len(texts), "mean_len": float(lengths.mean())}, ensure_ascii=False), flush=True)
+
+    four_bit = args.precision == "4bit"
+    model = load_qwen_dynamic(
+        _resolve(args.model_path), load_in_4bit=four_bit,
+        max_memory={0: f"{args.gpu_memory_gb:.1f}GiB", "cpu": "48GiB"} if torch.cuda.is_available() else None,
+    )
+    model.eval()
+    hidden = int(model.memory.hidden_size)
+    device = next(model.base_model.parameters()).device
+
+    bank = np.lib.format.open_memmap(
+        output_dir / "features.f16.npy", mode="w+", dtype=np.float16, shape=(len(texts), hidden)
+    )
+    boundaries = np.flatnonzero(np.diff(lengths[order])) + 1
+    starts = np.concatenate(([0], boundaries))
+    ends = np.concatenate((boundaries, [len(order)]))
+    started = time.perf_counter()
+    done = 0
+    for start, end in zip(starts.tolist(), ends.tolist()):
+        rows = order[start:end]
+        length = int(lengths[start])
+        capacity = max(1, min(args.max_batch, args.token_budget // max(1, length)))
+        for offset in range(0, len(rows), capacity):
+            block = rows[offset : offset + capacity]
+            ids = torch.tensor([encoded[i][:length] for i in block], dtype=torch.long, device=device)
+            mask = torch.ones_like(ids)
+            with torch.inference_mode():
+                vectors = model._encode_model_key(ids, mask).detach().to(torch.float16).cpu().numpy()
+            bank[block] = vectors
+            done += len(block)
+        if done % 5000 < capacity:
+            rate = done / max(1e-9, time.perf_counter() - started)
+            print(json.dumps({"phase": "encode", "precision": args.precision, "done": done,
+                              "total": len(texts), "texts_per_second": round(rate, 1)}), flush=True)
+    bank.flush()
+    seconds = time.perf_counter() - started
+    tokens = int(lengths.sum())
+    print(json.dumps({
+        "phase": "encode_done", "precision": args.precision, "texts": len(texts),
+        "seconds": round(seconds, 1), "texts_per_second": round(len(texts) / max(1e-9, seconds), 1),
+        "tokens_per_second": round(tokens / max(1e-9, seconds), 1),
+    }, ensure_ascii=False), flush=True)
+
+    (output_dir / "index.json").write_text(json.dumps(lookup, ensure_ascii=False), encoding="utf-8")
+    manifest = json.loads((reference_dir / "manifest.json").read_text(encoding="utf-8"))
+    manifest.update({"dtype": f"float16_{args.precision}", "bank": "features.f16.npy",
+                     "text_count": len(texts), "encoded_by": "probe_encoder_precision.py",
+                     "precision": args.precision})
+    (output_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+
+    # Feature-level agreement against the reference bank.
+    reference_path = reference_dir / ("features.pt" if (reference_dir / "features.pt").exists() else "features.f16.npy")
+    if reference_path.suffix == ".pt":
+        reference = torch.load(reference_path, map_location="cpu", weights_only=True).float()
+    else:
+        reference = torch.from_numpy(np.load(reference_path, mmap_mode="r")).float()
+    mine = torch.from_numpy(np.asarray(bank, dtype=np.float32))
+    common = [key for key in lookup if key in reference_index]
+    rows_mine = torch.tensor([lookup[key] for key in common])
+    rows_ref = torch.tensor([reference_index[key] for key in common])
+    a = reference[rows_ref]
+    b = mine[rows_mine]
+    cosine = torch.nn.functional.cosine_similarity(a, b, dim=-1)
+    print(json.dumps({
+        "phase": "compare", "precision": args.precision, "shared_texts": len(common),
+        "cosine_min": float(cosine.min()), "cosine_mean": float(cosine.mean()),
+        "cosine_p01": float(cosine.kthvalue(max(1, int(len(cosine) * 0.01))).values),
+        "max_abs_diff": float((a - b).abs().max()),
+    }, ensure_ascii=False), flush=True)
+    del model, reference, mine
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())