Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,198 @@
|
||||
"""Run a bounded, restart-like Natural Memory v2 endurance test.
|
||||
|
||||
The test deliberately clears only the in-process working copy. It never
|
||||
calls ``save_embedded_memory_weights`` and therefore does not modify the
|
||||
shipped model package. Defaults are conservative for a 12 GB GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import load_qwen_dynamic, load_tokenizer
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _gpu_stats() -> dict[str, float | int | None]:
|
||||
if not torch.cuda.is_available():
|
||||
return {"allocated_mb": None, "reserved_mb": None, "free_mb": None, "total_mb": None}
|
||||
allocated = torch.cuda.memory_allocated() / (1024 * 1024)
|
||||
reserved = torch.cuda.memory_reserved() / (1024 * 1024)
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
return {
|
||||
"allocated_mb": round(allocated, 2),
|
||||
"reserved_mb": round(reserved, 2),
|
||||
"free_mb": round(free / (1024 * 1024), 2),
|
||||
"total_mb": round(total / (1024 * 1024), 2),
|
||||
}
|
||||
|
||||
|
||||
def _long_input(tokenizer, target_tokens: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
seed_text = (
|
||||
"这是 Natural Memory v2 的长上下文压力测试。它应当把旧上下文归档到有界的记忆记录，"
|
||||
"保留最近工作窗口，并且不让当前 token 对全部页面做注意力。"
|
||||
)
|
||||
seed = tokenizer(seed_text, add_special_tokens=False, return_tensors="pt")["input_ids"][0]
|
||||
repeats = max(1, (target_tokens + seed.numel() - 1) // seed.numel())
|
||||
ids = seed.repeat(repeats)[:target_tokens].unsqueeze(0)
|
||||
return ids, torch.ones_like(ids)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
||||
parser.add_argument("--rounds", type=int, default=40)
|
||||
parser.add_argument("--records-per-round", type=int, default=8)
|
||||
parser.add_argument("--long-context-every", type=int, default=10)
|
||||
parser.add_argument("--long-context-tokens", type=int, default=2048)
|
||||
parser.add_argument("--kv-budget", type=int, default=1024)
|
||||
parser.add_argument("--chunk-tokens", type=int, default=256)
|
||||
parser.add_argument("--duration-minutes", type=float, default=0.0)
|
||||
parser.add_argument("--output", default="natural_memory_v2_stress_report.json")
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
args = parser.parse_args()
|
||||
if args.rounds < 1 or args.records_per_round < 1:
|
||||
raise SystemExit("--rounds and --records-per-round must be positive")
|
||||
if args.kv_budget < 1 or args.long_context_tokens < args.kv_budget:
|
||||
raise SystemExit("--long-context-tokens must be >= --kv-budget >= 1")
|
||||
|
||||
model_path = _project_path(args.model_path)
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
model = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit)
|
||||
model.eval()
|
||||
device = model._find_layer_device()
|
||||
if model.memory_os_v2 is None:
|
||||
raise RuntimeError("the selected model does not contain hierarchical memory")
|
||||
|
||||
# Work on an empty ephemeral memory instance. The embedded package on
|
||||
# disk is never changed by this script.
|
||||
model.clear_hierarchical_memory()
|
||||
model.memory_os_v2.read_threshold = 0.0
|
||||
model.memory_os_v2.kv_budget.max_tokens = min(
|
||||
int(args.kv_budget), int(model.memory_os_v2.kv_budget.hard_max_tokens)
|
||||
)
|
||||
model.reset_runtime_memory(batch_size=1, device=device)
|
||||
|
||||
latencies: list[float] = []
|
||||
writes = 0
|
||||
read_hits = 0
|
||||
read_attempts = 0
|
||||
compactions = 0
|
||||
archived_records = 0
|
||||
errors: list[str] = []
|
||||
start = time.perf_counter()
|
||||
next_record = 0
|
||||
try:
|
||||
for round_index in range(args.rounds):
|
||||
if args.duration_minutes > 0 and (time.perf_counter() - start) >= args.duration_minutes * 60:
|
||||
break
|
||||
round_start = time.perf_counter()
|
||||
try:
|
||||
for _ in range(args.records_per_round):
|
||||
record_index = next_record
|
||||
next_record += 1
|
||||
key = torch.randn(model.memory.hidden_size)
|
||||
token_ids = torch.tensor([1000 + (record_index % 10000), 2000 + (record_index % 10000)])
|
||||
text = f"stress record {record_index} belongs to Natural Memory endurance test"
|
||||
record, _ = model.memory_os_v2.write(
|
||||
text=text,
|
||||
key=key,
|
||||
summary=key,
|
||||
token_ids=token_ids,
|
||||
token_mask=torch.ones_like(token_ids, dtype=torch.bool),
|
||||
memory_type="stress_test",
|
||||
entity=f"stress-user-{record_index % 7}",
|
||||
attribute=f"attribute-{record_index}",
|
||||
value=str(record_index),
|
||||
importance=0.9,
|
||||
confidence=0.99,
|
||||
source="stress_test",
|
||||
trusted=True,
|
||||
force=True,
|
||||
)
|
||||
writes += 1
|
||||
read_attempts += 1
|
||||
found, _ = model.memory_os_v2.read(
|
||||
query_key=key,
|
||||
query_text=text,
|
||||
query_token_ids=token_ids,
|
||||
top_k_pages=4,
|
||||
top_k_records=8,
|
||||
max_hops=3,
|
||||
)
|
||||
if any(item.record_id == record.record_id for item in found):
|
||||
read_hits += 1
|
||||
|
||||
if args.long_context_every > 0 and (round_index + 1) % args.long_context_every == 0:
|
||||
ids, mask = _long_input(tokenizer, args.long_context_tokens)
|
||||
_, _, plan = model.compact_context_for_kv(
|
||||
ids.to(device),
|
||||
mask.to(device),
|
||||
archive=True,
|
||||
chunk_tokens=args.chunk_tokens,
|
||||
)
|
||||
if plan.get("compacted"):
|
||||
compactions += 1
|
||||
archived_records += int(plan.get("archived_records", 0))
|
||||
except Exception as error: # keep the report useful after one bad round
|
||||
errors.append(f"round {round_index}: {type(error).__name__}: {error}")
|
||||
latencies.append(time.perf_counter() - round_start)
|
||||
finally:
|
||||
final_stats = model.memory_v2_stats()
|
||||
audit = model.audit_memory()
|
||||
gpu_peak = _gpu_stats()
|
||||
model.close_memory_storage()
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
latencies_sorted = sorted(latencies)
|
||||
report = {
|
||||
"format_version": 1,
|
||||
"model_path": str(model_path),
|
||||
"rounds_completed": len(latencies),
|
||||
"records_per_round": args.records_per_round,
|
||||
"writes": writes,
|
||||
"read_attempts": read_attempts,
|
||||
"read_hits": read_hits,
|
||||
"read_hit_rate": read_hits / max(1, read_attempts),
|
||||
"compactions": compactions,
|
||||
"archived_records": archived_records,
|
||||
"round_latency_seconds": {
|
||||
"mean": statistics.fmean(latencies) if latencies else 0.0,
|
||||
"p50": latencies_sorted[len(latencies_sorted) // 2] if latencies_sorted else 0.0,
|
||||
"p95": latencies_sorted[min(len(latencies_sorted) - 1, int(len(latencies_sorted) * 0.95))] if latencies_sorted else 0.0,
|
||||
"max": max(latencies, default=0.0),
|
||||
},
|
||||
"max_gpu": gpu_peak,
|
||||
"final_memory_stats": final_stats,
|
||||
"audit": audit,
|
||||
"errors": errors,
|
||||
"package_mutated": False,
|
||||
"warning": "This is an endurance/safety smoke test, not a quality benchmark.",
|
||||
}
|
||||
output = _project_path(args.output)
|
||||
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user