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

merge_memory_weights.py

@@ -0,0 +1,307 @@
+"""Merge dynamic-memory tensors into a copy of a local safetensors shard.
+
+The base Qwen shards are never modified.  The output directory contains
+hardlinks for unchanged files and a newly written second shard with the
+memory tensors appended to its safetensors payload.  The normal HF loader
+continues to see the original Qwen keys; ``qwen_integration`` loads the
+embedded ``dynamic_memory.*`` keys when it sees ``memory_merge.json``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import struct
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+
+
+def _read_header(path: Path) -> tuple[dict[str, Any], int, int]:
+    with path.open("rb") as handle:
+        raw_length = handle.read(8)
+        if len(raw_length) != 8:
+            raise ValueError(f"not a safetensors file: {path}")
+        header_length = struct.unpack("<Q", raw_length)[0]
+        header_bytes = handle.read(header_length)
+        if len(header_bytes) != header_length:
+            raise ValueError(f"truncated safetensors header: {path}")
+    return json.loads(header_bytes.decode("utf-8")), int(header_length), 8 + int(header_length)
+
+
+def _load_adapter_tensors(adapter_dir: Path, memory_state: Path | None) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+    tensors: dict[str, torch.Tensor] = {}
+    memory = torch.load(adapter_dir / "memory.pt", map_location="cpu", weights_only=True)
+    tensors.update({f"dynamic_memory.memory.{key}": value.detach().cpu().contiguous() for key, value in memory.items()})
+
+    retriever_path = adapter_dir / "text_retriever.pt"
+    if retriever_path.exists():
+        retriever = torch.load(retriever_path, map_location="cpu", weights_only=True)
+        tensors.update(
+            {f"dynamic_memory.text_retriever.{key}": value.detach().cpu().contiguous() for key, value in retriever.items()}
+        )
+
+    policy_path = adapter_dir / "memory_policy.pt"
+    if policy_path.exists():
+        policy = torch.load(policy_path, map_location="cpu", weights_only=True)
+        tensors.update(
+            {f"dynamic_memory.memory_policy.{key}": value.detach().cpu().contiguous() for key, value in policy.items()}
+        )
+
+    surgery_path = adapter_dir / "surgery.pt"
+    if surgery_path.exists():
+        surgery = torch.load(surgery_path, map_location="cpu", weights_only=True)
+        for layer, value in surgery.get("blend_logits", {}).items():
+            tensors[f"dynamic_memory.blend_logits.{layer}"] = value.detach().cpu().contiguous()
+
+    metadata = json.loads((adapter_dir / "memory_config.json").read_text(encoding="utf-8"))
+    memory_config = dict(metadata.get("memory_config", {}))
+
+    persistent_source = memory_state
+    if persistent_source is None and (adapter_dir / "persistent_memory.pt").exists():
+        persistent_source = adapter_dir / "persistent_memory.pt"
+    if persistent_source is not None:
+        payload = torch.load(persistent_source, map_location="cpu", weights_only=True)
+        if not isinstance(payload, dict) or not isinstance(payload.get("memory_state"), torch.Tensor):
+            raise ValueError(f"memory state file has no tensor memory_state: {persistent_source}")
+        for key, value in payload.items():
+            if isinstance(value, torch.Tensor):
+                tensors[f"dynamic_memory.persistent.{key}"] = value.detach().cpu().contiguous()
+        memory_config["persistent_memory"] = True
+
+    if not tensors:
+        raise ValueError("no tensors were found to merge")
+    metadata["memory_config"] = memory_config
+    metadata["checkpoint_contains_user_memory"] = persistent_source is not None
+    return tensors, metadata
+
+
+def _write_merged_shard(base_shard: Path, extra_shard: Path, output_shard: Path) -> None:
+    base_header, base_header_length, base_data_start = _read_header(base_shard)
+    extra_header, _, extra_data_start = _read_header(extra_shard)
+    if "__metadata__" in extra_header:
+        extra_metadata = extra_header.pop("__metadata__")
+    else:
+        extra_metadata = {}
+    if "__metadata__" in base_header:
+        base_metadata = base_header.pop("__metadata__")
+    else:
+        base_metadata = {}
+    collisions = set(base_header).intersection(extra_header)
+    if collisions:
+        raise ValueError(f"safetensors key collision while merging: {sorted(collisions)[:4]}")
+
+    base_payload_size = base_shard.stat().st_size - base_data_start
+    merged_header: dict[str, Any] = {}
+    new_metadata = dict(base_metadata)
+    new_metadata.update({str(key): str(value) for key, value in extra_metadata.items()})
+    merged_header["__metadata__"] = new_metadata
+
+    for key, entry in base_header.items():
+        item = dict(entry)
+        start, end = item["data_offsets"]
+        item["data_offsets"] = [int(start), int(end)]
+        merged_header[key] = item
+    # Copy the entries: the offset-fixup loop mutates merged_header and must
+    # not mutate extra_header, otherwise the second loop iteration applies the
+    # base offset repeatedly.
+    merged_header.update({key: dict(entry) for key, entry in extra_header.items()})
+
+    # Offsets are relative to the beginning of the data section, so changing
+    # the header length does not move the base tensors in that coordinate
+    # system.  Recompute the header until its padded size is stable.
+    while True:
+        encoded = json.dumps(merged_header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+        padded_length = (len(encoded) + 7) // 8 * 8
+        for key, entry in base_header.items():
+            start, end = entry["data_offsets"]
+            merged_header[key]["data_offsets"] = [int(start), int(end)]
+        # ``data_offsets`` are relative to the beginning of the data section,
+        # not absolute file offsets.  The extra payload follows the base
+        # payload inside that merged data section.
+        extra_start = base_payload_size
+        for key, entry in extra_header.items():
+            start, end = entry["data_offsets"]
+            merged_header[key]["data_offsets"] = [int(extra_start + start), int(extra_start + end)]
+        encoded_next = json.dumps(merged_header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+        next_length = (len(encoded_next) + 7) // 8 * 8
+        if next_length == padded_length:
+            encoded = encoded_next + b" " * (next_length - len(encoded_next))
+            break
+
+    output_shard.parent.mkdir(parents=True, exist_ok=True)
+    with output_shard.open("wb") as output, base_shard.open("rb") as base, extra_shard.open("rb") as extra:
+        output.write(struct.pack("<Q", len(encoded)))
+        output.write(encoded)
+        base.seek(base_data_start)
+        shutil.copyfileobj(base, output, length=16 * 1024 * 1024)
+        extra.seek(extra_data_start)
+        shutil.copyfileobj(extra, output, length=16 * 1024 * 1024)
+
+
+def _hardlink_base_files(base_dir: Path, output_dir: Path, merged_shard_name: str) -> None:
+    for source in base_dir.iterdir():
+        if not source.is_file() or source.name in {"memory_config.json", "memory_merge.json", merged_shard_name}:
+            continue
+        destination = output_dir / source.name
+        try:
+            os.link(source, destination)
+        except OSError:
+            try:
+                # Some Windows volumes reject hardlinks but allow symlinks;
+                # this keeps the unchanged 5GB shard from being duplicated.
+                os.symlink(source, destination)
+            except OSError:
+                # Last-resort fallback for filesystems that allow neither.
+                # The large first shard is still copied only when necessary.
+                shutil.copy2(source, destination)
+
+
+def merge_package(
+    base_dir: Path,
+    adapter_dir: Path,
+    output_dir: Path,
+    memory_state: Path | None,
+) -> dict[str, Any]:
+    if output_dir.exists():
+        raise FileExistsError(f"refusing to overwrite existing output: {output_dir}")
+    base_index = base_dir / "model.safetensors.index.json"
+    index = json.loads(base_index.read_text(encoding="utf-8"))
+    shard_names = sorted(set(index["weight_map"].values()))
+    if len(shard_names) != 2 or "model.safetensors-00002-of-00002.safetensors" not in shard_names:
+        raise ValueError(f"expected the local two-shard Qwen layout, got {shard_names}")
+
+    output_dir.mkdir(parents=True)
+    merged_shard_name = "model.safetensors-00002-of-00002.safetensors"
+    _hardlink_base_files(base_dir, output_dir, merged_shard_name)
+
+    tensors, adapter_metadata = _load_adapter_tensors(adapter_dir, memory_state)
+    extra_shard = output_dir / "memory_extra.safetensors"
+    save_file(tensors, str(extra_shard), metadata={"format": "qwen_dynamic_memory_embedded_v1"})
+    _write_merged_shard(
+        base_dir / merged_shard_name,
+        extra_shard,
+        output_dir / merged_shard_name,
+    )
+    # This is a self-owned staging file; the merged shard is the only copy
+    # that belongs in the output package.
+    extra_shard.unlink()
+
+    (output_dir / "memory_config.json").write_text(
+        json.dumps(adapter_metadata, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+    manifest = {
+        "format_version": 1,
+        "format": "qwen_dynamic_memory_merged_shard",
+        "shard_mode": "combined",
+        "memory_weights": merged_shard_name,
+        "base_model": str(base_dir),
+        "source_adapter": str(adapter_dir),
+        "source_memory_state": str(memory_state) if memory_state is not None else None,
+        "tensor_prefix": "dynamic_memory.",
+        "checkpoint_contains_user_memory": bool(adapter_metadata.get("checkpoint_contains_user_memory")),
+        "memory_config": adapter_metadata.get("memory_config", {}),
+    }
+    (output_dir / "memory_merge.json").write_text(
+        json.dumps(manifest, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+    return manifest
+
+
+def add_memory_shard(
+    base_dir: Path,
+    adapter_dir: Path,
+    shard_path: Path,
+    memory_state: Path | None,
+) -> dict[str, Any]:
+    """Add a small extra-only memory shard beside an existing Qwen model."""
+
+    if shard_path.exists():
+        raise FileExistsError(f"refusing to overwrite existing shard: {shard_path}")
+    if not (base_dir / "config.json").exists() or not (base_dir / "model.safetensors.index.json").exists():
+        raise ValueError(f"not a complete local Qwen model directory: {base_dir}")
+
+    tensors, adapter_metadata = _load_adapter_tensors(adapter_dir, memory_state)
+    shard_path.parent.mkdir(parents=True, exist_ok=True)
+    save_file(
+        tensors,
+        str(shard_path),
+        metadata={"format": "qwen_dynamic_memory_extra_shard_v1"},
+    )
+    index_path = base_dir / "model.safetensors.index.json"
+    index = json.loads(index_path.read_text(encoding="utf-8"))
+    weight_map = index.setdefault("weight_map", {})
+    collisions = sorted(set(weight_map).intersection(tensors))
+    if collisions:
+        raise ValueError(f"weight-map key collision while adding memory shard: {collisions[:4]}")
+    weight_map.update({key: shard_path.name for key in tensors})
+    tensor_bytes = sum(value.numel() * value.element_size() for value in tensors.values())
+    index.setdefault("metadata", {})["total_size"] = int(
+        index.get("metadata", {}).get("total_size", 0) + tensor_bytes
+    )
+    index_path.write_text(json.dumps(index, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
+    memory_config_path = base_dir / "memory_config.json"
+    memory_config_path.write_text(
+        json.dumps(adapter_metadata, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+    manifest = {
+        "format_version": 1,
+        "format": "qwen_dynamic_memory_extra_shard",
+        "shard_mode": "extra_only",
+        "memory_weights": shard_path.name,
+        "base_model": str(base_dir),
+        "source_adapter": str(adapter_dir),
+        "source_memory_state": str(memory_state) if memory_state is not None else None,
+        "tensor_prefix": "dynamic_memory.",
+        "checkpoint_contains_user_memory": bool(adapter_metadata.get("checkpoint_contains_user_memory")),
+        "memory_config": adapter_metadata.get("memory_config", {}),
+    }
+    (base_dir / "memory_merge.json").write_text(
+        json.dumps(manifest, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+    return manifest
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base-model", required=True)
+    parser.add_argument("--adapter", required=True)
+    parser.add_argument("--output", default=None)
+    parser.add_argument("--add-shard", default=None, help="add a small extra-only shard beside the base model")
+    parser.add_argument("--memory-state", default=None)
+    args = parser.parse_args()
+    if args.add_shard:
+        manifest = add_memory_shard(
+            Path(args.base_model),
+            Path(args.adapter),
+            Path(args.add_shard),
+            Path(args.memory_state) if args.memory_state else None,
+        )
+        output_shard = Path(args.add_shard)
+        print(f"merged_package={Path(args.base_model).resolve()}")
+    else:
+        if not args.output:
+            parser.error("--output is required unless --add-shard is used")
+        manifest = merge_package(
+            Path(args.base_model),
+            Path(args.adapter),
+            Path(args.output),
+            Path(args.memory_state) if args.memory_state else None,
+        )
+        output_shard = Path(args.output) / str(manifest["memory_weights"])
+        print(f"merged_package={Path(args.output).resolve()}")
+    print(f"merged_shard={output_shard.resolve()}")
+    print(f"merged_size_gb={output_shard.stat().st_size / (1024**3):.3f}")
+    print(f"contains_user_memory={manifest['checkpoint_contains_user_memory']}")
+
+
+if __name__ == "__main__":
+    main()