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，读写关闭时与原生模型逐位相同
This commit is contained in:
WpyQwq
@@ -0,0 +1,178 @@
|
||||
"""Build the NM2.1 package: the delivered router merged into the Qwen3.5 memory model.
|
||||
|
||||
"Merge" here means producing a **self-contained package** rather than a sidecar file: the
|
||||
base Qwen3.5-4B shards are copied unchanged and only the 16
|
||||
``dynamic_memory.memory_router_v2.*`` tensors inside the memory shard are replaced with the
|
||||
delivered router. Everything else -- text retriever, persistent slots, memory policy,
|
||||
configs, tokenizer -- is left byte-identical, so anything the package did before it still
|
||||
does, with a different router.
|
||||
|
||||
Verification is not "the script exited zero": the script reloads the written shard and
|
||||
compares all 16 tensors bit-exactly against the source checkpoint, and separately reports
|
||||
every other tensor as unchanged (same bytes as the source shard).
|
||||
|
||||
Usage::
|
||||
|
||||
python -m V2_dpskw.build_nm2_1_package ^
|
||||
--source-package H:\\Memory\\dynamic_memory_lab\\qwen3_5_4b_natural_memory_v2 ^
|
||||
--output-package H:\\Memory\\dynamic_memory_lab\\qwen3_5_4b_natural_memory_v2_1 ^
|
||||
--router-checkpoint checkpoints/router_replay_v7_v2_128/memory_router_v2.pt ^
|
||||
--label NM2.1
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
ROUTER_PREFIX = "dynamic_memory.memory_router_v2."
|
||||
INDEX_NAME = "model.safetensors.index.json"
|
||||
MERGE_NAME = "memory_merge.json"
|
||||
|
||||
|
||||
def sha256_of(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1 << 22), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--source-package", required=True)
|
||||
parser.add_argument("--output-package", required=True)
|
||||
parser.add_argument("--router-checkpoint", required=True)
|
||||
parser.add_argument("--label", default="NM2.1")
|
||||
parser.add_argument("--report", default="nm2_1_build_report.json")
|
||||
args = parser.parse_args()
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
source = Path(args.source_package)
|
||||
output = Path(args.output_package)
|
||||
router_path = Path(args.router_checkpoint)
|
||||
if not source.exists():
|
||||
raise SystemExit(f"source package not found: {source}")
|
||||
if not router_path.exists():
|
||||
raise SystemExit(f"router checkpoint not found: {router_path}")
|
||||
if output.exists() and any(output.iterdir()):
|
||||
raise SystemExit(f"output package already exists and is not empty: {output}")
|
||||
|
||||
router = torch.load(router_path, map_location="cpu", weights_only=True)
|
||||
if isinstance(router, dict) and "router_state_dict" in router:
|
||||
router = router["router_state_dict"]
|
||||
|
||||
# --- which shard holds the router tensors? -------------------------------------
|
||||
index = json.loads((source / INDEX_NAME).read_text(encoding="utf-8"))
|
||||
weight_map = index["weight_map"]
|
||||
router_keys = sorted(key for key in weight_map if key.startswith(ROUTER_PREFIX))
|
||||
shards = sorted({weight_map[key] for key in router_keys})
|
||||
if len(shards) != 1:
|
||||
raise SystemExit(f"router tensors span several shards: {shards}")
|
||||
print(json.dumps({"phase": "locate", "router_tensors": len(router_keys),
|
||||
"shard": shards[0],
|
||||
"expected_from_checkpoint": len(router),
|
||||
"router_sha256": sha256_of(router_path)}), flush=True)
|
||||
|
||||
missing = [key for key in router if ROUTER_PREFIX + key not in weight_map]
|
||||
if missing:
|
||||
raise SystemExit(f"checkpoint keys not present in the package: {missing}")
|
||||
if len(router_keys) != len(router):
|
||||
raise SystemExit(f"package has {len(router_keys)} router tensors, checkpoint has {len(router)}")
|
||||
|
||||
# --- copy the package -----------------------------------------------------------
|
||||
print(json.dumps({"phase": "copy", "from": str(source), "to": str(output)}), flush=True)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
result = subprocess.run(["robocopy", str(source), str(output), "/E", "/NFL", "/NDL",
|
||||
"/NJH", "/NJS", "/NP", "/R:2", "/W:2"],
|
||||
capture_output=True, text=True)
|
||||
if result.returncode > 7:
|
||||
raise SystemExit(f"robocopy failed ({result.returncode}): {result.stdout[-800:]}")
|
||||
copied = sorted(p.name for p in output.iterdir())
|
||||
expected = sorted(p.name for p in source.iterdir())
|
||||
if copied != expected:
|
||||
raise SystemExit(f"copy mismatch: missing={set(expected) - set(copied)} extra={set(copied) - set(expected)}")
|
||||
|
||||
# --- rewrite the memory shard ---------------------------------------------------
|
||||
shard_path = output / shards[0]
|
||||
source_shard_path = source / shards[0]
|
||||
with safe_open(str(source_shard_path), framework="pt") as handle:
|
||||
metadata = handle.metadata() or {}
|
||||
tensors = {key: handle.get_tensor(key) for key in handle.keys()}
|
||||
untouched = 0
|
||||
replaced = []
|
||||
for key, value in router.items():
|
||||
full = ROUTER_PREFIX + key
|
||||
original = tensors[full]
|
||||
if tuple(original.shape) != tuple(value.shape):
|
||||
raise SystemExit(f"shape mismatch for {full}: {tuple(original.shape)} vs {tuple(value.shape)}")
|
||||
tensors[full] = value.to(dtype=original.dtype).contiguous()
|
||||
replaced.append(full)
|
||||
for key in tensors:
|
||||
if key not in replaced:
|
||||
untouched += 1
|
||||
save_file(tensors, str(shard_path), metadata=metadata)
|
||||
print(json.dumps({"phase": "shard_rewritten", "shard": shards[0],
|
||||
"replaced": len(replaced), "untouched": untouched}), flush=True)
|
||||
|
||||
# --- provenance in the merge manifest -------------------------------------------
|
||||
merge_path = output / MERGE_NAME
|
||||
merge = json.loads(merge_path.read_text(encoding="utf-8"))
|
||||
merge["package_label"] = args.label
|
||||
merge["router_swap"] = {
|
||||
"built_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"built_from_package": str(source),
|
||||
"router_checkpoint": str(router_path),
|
||||
"router_sha256": sha256_of(router_path),
|
||||
"replaced_tensors": len(replaced),
|
||||
"shard": shards[0],
|
||||
}
|
||||
merge_path.write_text(json.dumps(merge, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
# --- verify: reload what was written, compare bit-exactly ------------------------
|
||||
with safe_open(str(shard_path), framework="pt") as handle:
|
||||
written = {key: handle.get_tensor(key) for key in handle.keys()}
|
||||
mismatches = []
|
||||
for key, value in router.items():
|
||||
full = ROUTER_PREFIX + key
|
||||
if not torch.equal(written[full], value.to(dtype=written[full].dtype)):
|
||||
mismatches.append(full)
|
||||
with safe_open(str(source_shard_path), framework="pt") as handle:
|
||||
originals = {key: handle.get_tensor(key) for key in handle.keys()}
|
||||
changed_others = [key for key in originals
|
||||
if key not in replaced and not torch.equal(originals[key], written.get(key))]
|
||||
|
||||
report = {
|
||||
"package_label": args.label,
|
||||
"output_package": str(output),
|
||||
"source_package": str(source),
|
||||
"router_checkpoint": str(router_path),
|
||||
"router_sha256": sha256_of(router_path),
|
||||
"shard": shards[0],
|
||||
"files_copied": len(copied),
|
||||
"router_tensors_replaced": len(replaced),
|
||||
"tensors_left_untouched": untouched,
|
||||
"verification": {
|
||||
"router_tensors_bit_exact": not mismatches,
|
||||
"router_mismatches": mismatches,
|
||||
"non_router_tensors_unchanged": not changed_others,
|
||||
"changed_non_router": changed_others,
|
||||
},
|
||||
}
|
||||
report["passed"] = (not mismatches) and (not changed_others)
|
||||
Path(args.report).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2), flush=True)
|
||||
return 0 if report["passed"] else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user