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,128 @@
+"""Smoke-test the model-owned KV offload and long-context compaction paths.
+
+The test intentionally uses the real Qwen package.  It verifies that
+Transformers' CPU-backed DynamicCache can generate through Qwen3.5's hybrid
+linear/full-attention stack and that an over-budget prompt is archived into
+Natural Memory before only the recent hot window is passed to generation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+
+if __package__ in {None, ""}:
+    import sys
+
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from V2_dpskw.qwen_integration import load_qwen_dynamic, load_tokenizer
+from V2_dpskw.stream_chat_qwen_memory import _chat_tensor
+
+
+@torch.inference_mode()
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--model-path",
+        default=r"H:\Memory\V2_dpskw\qwen3_5_4b_natural_memory_v2",
+    )
+    parser.add_argument(
+        "--report",
+        default=r"H:\Memory\V2_dpskw\kv_offload_compaction_test.json",
+    )
+    parser.add_argument("--compact-hot-tokens", type=int, default=32)
+    parser.add_argument("--max-new-tokens", type=int, default=4)
+    parser.add_argument("--no-4bit", action="store_true")
+    args = parser.parse_args()
+
+    model_path = Path(args.model_path)
+    tokenizer = load_tokenizer(model_path)
+    model = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
+    model.eval()
+    model.memory_config.kv_offload = True
+    device = model._find_layer_device()
+
+    answer_inputs = {
+        key: value.to(device)
+        for key, value in _chat_tensor(tokenizer, "请用一句话说明你支持什么。").items()
+    }
+    generated = model.generate(
+        **answer_inputs,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        update_memory=False,
+        use_cache=True,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    answer = tokenizer.decode(generated[0].detach().cpu().tolist(), skip_special_tokens=True).strip()
+
+    if model.memory_os_v2 is None:
+        raise RuntimeError("the selected package does not contain Natural Memory v2")
+    model.memory_os_v2.kv_budget.max_tokens = int(args.compact_hot_tokens)
+    model.memory_os_v2.kv_budget.keep_recent_tokens = min(
+        model.memory_os_v2.kv_budget.keep_recent_tokens,
+        model.memory_os_v2.kv_budget.max_tokens,
+    )
+    model.memory_config.context_chunk_tokens = max(4, min(16, args.compact_hot_tokens // 2))
+    long_text = " ".join(
+        [
+            "历史上下文片段用于验证 Natural Memory 的自动分页压缩。",
+            "这段内容应该被写入长期上下文页面，而不是继续占用当前热 KV。",
+        ]
+        * 12
+    )
+    encoded = tokenizer(long_text, add_special_tokens=False, return_tensors="pt")
+    long_ids = encoded["input_ids"].to(device)
+    long_mask = encoded.get("attention_mask")
+    if long_mask is None:
+        long_mask = torch.ones_like(long_ids)
+    long_mask = long_mask.to(device)
+    before = model.memory_v2_stats()
+    generated_long = model.generate(
+        input_ids=long_ids,
+        attention_mask=long_mask,
+        max_new_tokens=1,
+        do_sample=False,
+        update_memory=False,
+        use_cache=True,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    compaction = dict(model.runtime.context_compaction or {})
+    after = model.memory_v2_stats()
+    compaction_pass = bool(
+        compaction["compacted"]
+        and compaction["archived_records"] > 0
+        and max(compaction["retained_tokens"]) <= args.compact_hot_tokens
+        and int(compaction["retained_tokens"][0]) <= args.compact_hot_tokens
+        and after["active_records"] >= before["active_records"] + compaction["archived_records"]
+    )
+
+    report = {
+        "model_path": str(model_path),
+        "transformers_cache": "DynamicCache(offloading=True)",
+        "kv_offload_pass": True,
+        "answer": answer,
+        "compaction": compaction,
+        "compaction_pass": compaction_pass,
+        "long_generation_tokens": int(generated_long.shape[1]),
+        "memory_before": before,
+        "memory_after": after,
+        "device": str(device),
+        "cuda_memory_allocated_bytes": (
+            int(torch.cuda.memory_allocated(device)) if torch.cuda.is_available() else 0
+        ),
+    }
+    report_path = Path(args.report)
+    report_path.parent.mkdir(parents=True, exist_ok=True)
+    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(json.dumps(report, ensure_ascii=False, indent=2))
+    if not compaction_pass:
+        raise SystemExit("context compaction verification failed")
+
+
+if __name__ == "__main__":
+    main()