Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,293 @@
|
||||
"""Build an embedded Natural Memory v2 package from the current v1 package.
|
||||
|
||||
Unchanged Qwen shards are hard-linked when the filesystem permits it. The
|
||||
custom memory shard is rewritten once to include the trained V2 router and a
|
||||
compact V2 page payload, while the official model shards remain untouched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
if __package__ in {None, ""}:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from dynamic_memory_lab.memory_os_v2 import MemoryOSV2, MemoryRouterV2
|
||||
|
||||
|
||||
def _link_or_copy(source: Path, target: Path, *, allow_copy: bool) -> None:
|
||||
try:
|
||||
os.link(source, target)
|
||||
except OSError:
|
||||
if not allow_copy:
|
||||
raise RuntimeError(
|
||||
"the destination filesystem does not support hard links; "
|
||||
"refusing to duplicate multi-gigabyte Qwen shards. "
|
||||
"Re-run with --allow-copy-base only when enough disk space "
|
||||
"has been explicitly reserved"
|
||||
)
|
||||
shutil.copy2(source, target)
|
||||
|
||||
|
||||
def _pack_v2_payload(payload: dict[str, Any]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
metadata = dict(payload)
|
||||
|
||||
def pack_record(item: dict[str, Any], prefix: str) -> dict[str, Any]:
|
||||
item = dict(item)
|
||||
for field in ("key", "summary", "token_ids", "token_mask"):
|
||||
value = item.pop(field, None)
|
||||
if isinstance(value, torch.Tensor):
|
||||
name = f"dynamic_memory.v2.{prefix}.{field}"
|
||||
tensors[name] = value.detach().cpu().contiguous()
|
||||
item[f"{field}_ref"] = name
|
||||
return item
|
||||
|
||||
metadata["records"] = [
|
||||
pack_record(item, f"records.{index}")
|
||||
for index, item in enumerate(payload.get("records", []))
|
||||
]
|
||||
metadata["quarantine"] = [
|
||||
pack_record(item, f"quarantine.{index}")
|
||||
for index, item in enumerate(payload.get("quarantine", []))
|
||||
]
|
||||
metadata["pages"] = []
|
||||
for index, item in enumerate(payload.get("pages", [])):
|
||||
item = dict(item)
|
||||
for field in ("key", "summary"):
|
||||
value = item.pop(field, None)
|
||||
if isinstance(value, torch.Tensor):
|
||||
name = f"dynamic_memory.v2.pages.{index}.{field}"
|
||||
tensors[name] = value.detach().cpu().contiguous()
|
||||
item[f"{field}_ref"] = name
|
||||
metadata["pages"].append(item)
|
||||
return tensors, metadata
|
||||
|
||||
|
||||
def build(args: argparse.Namespace) -> dict[str, Any]:
|
||||
base_dir = Path(args.base_package)
|
||||
output_dir = Path(args.output_dir)
|
||||
if output_dir.exists():
|
||||
raise FileExistsError(f"refusing to overwrite existing output: {output_dir}")
|
||||
output_dir.mkdir(parents=True)
|
||||
|
||||
base_manifest = json.loads((base_dir / "memory_merge.json").read_text(encoding="utf-8"))
|
||||
base_config = json.loads((base_dir / "memory_config.json").read_text(encoding="utf-8"))
|
||||
memory_name = str(base_manifest.get("memory_weights", "model.safetensors-00003-of-00003.safetensors"))
|
||||
source_memory = base_dir / memory_name
|
||||
if not source_memory.exists():
|
||||
raise FileNotFoundError(source_memory)
|
||||
|
||||
sources = [
|
||||
source
|
||||
for source in base_dir.iterdir()
|
||||
if source.is_file()
|
||||
and source.name not in {
|
||||
memory_name,
|
||||
"memory_config.json",
|
||||
"memory_merge.json",
|
||||
"model.safetensors.index.json",
|
||||
}
|
||||
]
|
||||
if not args.allow_copy_base:
|
||||
probe_source = next(
|
||||
(source for source in sources if source.name.endswith(".safetensors")),
|
||||
None,
|
||||
)
|
||||
if probe_source is not None:
|
||||
probe_target = output_dir / ".hardlink-probe"
|
||||
try:
|
||||
os.link(probe_source, probe_target)
|
||||
except OSError as error:
|
||||
raise RuntimeError(
|
||||
"base package cannot be assembled without copying its large "
|
||||
"shards on this filesystem; pass --allow-copy-base after "
|
||||
"checking free space"
|
||||
) from error
|
||||
finally:
|
||||
if probe_target.exists():
|
||||
probe_target.unlink()
|
||||
|
||||
for source in sources:
|
||||
_link_or_copy(source, output_dir / source.name, allow_copy=args.allow_copy_base)
|
||||
|
||||
with safe_open(str(source_memory), framework="pt", device="cpu") as handle:
|
||||
base_tensors = {key: handle.get_tensor(key) for key in handle.keys()}
|
||||
|
||||
router = MemoryRouterV2(
|
||||
int(base_config["hidden_size"]),
|
||||
router_dim=args.router_dim,
|
||||
num_heads=args.num_heads,
|
||||
max_hops=args.max_hops,
|
||||
)
|
||||
router_state = torch.load(args.router_checkpoint, map_location="cpu", weights_only=True)
|
||||
router.load_state_dict(router_state, strict=True)
|
||||
for key, value in router.state_dict().items():
|
||||
base_tensors[f"dynamic_memory.memory_router_v2.{key}"] = value.detach().cpu().contiguous()
|
||||
|
||||
memory_os = MemoryOSV2(
|
||||
int(base_config["hidden_size"]),
|
||||
router=router,
|
||||
)
|
||||
# Migrate the existing model-owned hot facts into V2 address pages. The
|
||||
# old bank remains intact; this is only a compatibility seed for the new
|
||||
# hierarchical route.
|
||||
legacy_prefix = "dynamic_memory.persistent."
|
||||
token_ids = base_tensors.get(f"{legacy_prefix}text_token_ids")
|
||||
token_mask = base_tensors.get(f"{legacy_prefix}text_token_mask")
|
||||
slot_valid = base_tensors.get(f"{legacy_prefix}text_slot_valid")
|
||||
slot_keys = base_tensors.get(f"{legacy_prefix}text_slot_keys")
|
||||
if all(isinstance(value, torch.Tensor) for value in (token_ids, token_mask, slot_valid, slot_keys)):
|
||||
for batch_index in range(slot_valid.shape[0]):
|
||||
for slot in range(slot_valid.shape[1]):
|
||||
if not bool(slot_valid[batch_index, slot].item()):
|
||||
continue
|
||||
ids = token_ids[batch_index, slot][token_mask[batch_index, slot]]
|
||||
memory_os.write(
|
||||
text=f"legacy_hot_slot:{batch_index}:{slot}",
|
||||
key=slot_keys[batch_index, slot],
|
||||
summary=slot_keys[batch_index, slot],
|
||||
memory_type="legacy_hot_text",
|
||||
importance=0.95,
|
||||
confidence=0.95,
|
||||
source="v1_migration",
|
||||
slot_index=slot,
|
||||
token_ids=ids,
|
||||
token_mask=torch.ones_like(ids, dtype=torch.bool),
|
||||
trusted=True,
|
||||
)
|
||||
|
||||
v2_tensors, v2_metadata = _pack_v2_payload(memory_os.export_payload())
|
||||
base_tensors.update(v2_tensors)
|
||||
save_file(
|
||||
base_tensors,
|
||||
str(output_dir / memory_name),
|
||||
metadata={
|
||||
"format": "qwen_dynamic_memory_embedded_v2",
|
||||
"memory_os_v2_payload": json.dumps(v2_metadata, ensure_ascii=False, separators=(",", ":")),
|
||||
},
|
||||
)
|
||||
|
||||
memory_config = dict(base_config)
|
||||
saved = dict(memory_config.get("memory_config", {}))
|
||||
saved.update(
|
||||
{
|
||||
"memory_version": 2,
|
||||
"hierarchical_memory": True,
|
||||
"memory_router_dim": args.router_dim,
|
||||
"memory_router_heads": args.num_heads,
|
||||
"memory_page_capacity": args.page_capacity,
|
||||
"memory_max_pages": args.max_pages,
|
||||
"memory_hot_pages": args.hot_pages,
|
||||
"memory_top_k_pages": args.top_k_pages,
|
||||
"memory_top_k_records": args.top_k_records,
|
||||
"memory_max_hops": args.max_hops,
|
||||
"memory_coarse_index_bits": args.coarse_index_bits,
|
||||
"memory_v2_read_threshold": args.read_threshold,
|
||||
"memory_v2_write_threshold": args.write_threshold,
|
||||
"memory_storage_mode": args.memory_storage_mode,
|
||||
"memory_storage_path": args.memory_storage_path,
|
||||
"memory_resident_pages": args.memory_resident_pages,
|
||||
"memory_gpu_cache_records": args.memory_gpu_cache_records,
|
||||
"memory_gpu_cache_tokens": args.memory_gpu_cache_tokens,
|
||||
"memory_gpu_cache_reserve_mb": args.memory_gpu_cache_reserve_mb,
|
||||
"memory_gpu_cache_adaptive": args.memory_gpu_cache_adaptive,
|
||||
"kv_budget_tokens": args.kv_budget,
|
||||
"kv_hard_max_tokens": args.kv_hard_max,
|
||||
"kv_compaction_trigger": args.kv_trigger,
|
||||
"kv_keep_recent_tokens": args.kv_keep_recent,
|
||||
"persistent_memory": True,
|
||||
}
|
||||
)
|
||||
memory_config["memory_config"] = saved
|
||||
memory_config["router_v2_ready"] = True
|
||||
memory_config["checkpoint_contains_user_memory"] = True
|
||||
(output_dir / "memory_config.json").write_text(
|
||||
json.dumps(memory_config, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
|
||||
index = json.loads((base_dir / "model.safetensors.index.json").read_text(encoding="utf-8"))
|
||||
weight_map = index.setdefault("weight_map", {})
|
||||
for key in base_tensors:
|
||||
weight_map[key] = memory_name
|
||||
index.setdefault("metadata", {})["total_size"] = int(
|
||||
sum(value.numel() * value.element_size() for value in base_tensors.values())
|
||||
)
|
||||
(output_dir / "model.safetensors.index.json").write_text(
|
||||
json.dumps(index, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
|
||||
)
|
||||
manifest = dict(base_manifest)
|
||||
manifest.update(
|
||||
{
|
||||
"format_version": 2,
|
||||
"format": "qwen_dynamic_memory_v2_embedded",
|
||||
"base_model": str(base_dir),
|
||||
"source_router_checkpoint": str(args.router_checkpoint),
|
||||
"checkpoint_contains_user_memory": True,
|
||||
"memory_config": saved,
|
||||
}
|
||||
)
|
||||
(output_dir / "memory_merge.json").write_text(
|
||||
json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
return {
|
||||
"output_dir": str(output_dir),
|
||||
"memory_shard": str(output_dir / memory_name),
|
||||
"memory_tensor_count": len(base_tensors),
|
||||
"v2_records": memory_os.stats()["records"],
|
||||
"v2_pages": memory_os.stats()["pages"],
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--base-package", default="W:/Flash/model/dynamic_memory_lab/qwen3_5_4b_memory_merged_v13")
|
||||
parser.add_argument("--output-dir", default="W:/Flash/model/dynamic_memory_lab/qwen3_5_4b_memory_merged_v2")
|
||||
parser.add_argument("--router-checkpoint", default="W:/Flash/model/dynamic_memory_lab/checkpoints/natural_memory_v2_router/memory_router_v2.pt")
|
||||
parser.add_argument("--router-dim", type=int, default=128)
|
||||
parser.add_argument("--num-heads", type=int, default=8)
|
||||
parser.add_argument("--max-hops", type=int, default=3)
|
||||
parser.add_argument("--page-capacity", type=int, default=32)
|
||||
parser.add_argument("--max-pages", type=int, default=32768)
|
||||
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=20)
|
||||
parser.add_argument("--read-threshold", type=float, default=0.65)
|
||||
parser.add_argument("--write-threshold", type=float, default=0.50)
|
||||
parser.add_argument("--memory-storage-mode", choices=("embedded", "tiered"), default="embedded")
|
||||
parser.add_argument("--memory-storage-path", default=None)
|
||||
parser.add_argument("--memory-resident-pages", type=int, default=256)
|
||||
parser.add_argument("--memory-gpu-cache-records", type=int, default=256)
|
||||
parser.add_argument("--memory-gpu-cache-tokens", type=int, default=131072)
|
||||
parser.add_argument("--memory-gpu-cache-reserve-mb", type=int, default=2048)
|
||||
parser.add_argument(
|
||||
"--memory-gpu-cache-adaptive",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="keep the VRAM cache below the reserve line and fall back to system RAM when needed",
|
||||
)
|
||||
parser.add_argument("--kv-budget", type=int, default=32768)
|
||||
parser.add_argument("--kv-hard-max", type=int, default=131072)
|
||||
parser.add_argument("--kv-trigger", type=float, default=0.90)
|
||||
parser.add_argument("--kv-keep-recent", type=int, default=8192)
|
||||
parser.add_argument(
|
||||
"--allow-copy-base",
|
||||
action="store_true",
|
||||
help="allow copying the frozen Qwen shards when hard links are unavailable; requires substantial free disk space",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(json.dumps(build(parse_args()), ensure_ascii=False, indent=2))
|
||||
Reference in New Issue
Block a user