WpyQwq
643e22ecb9
- 引入 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，读写关闭时与原生模型逐位相同
311 lines
13 KiB
Python
311 lines
13 KiB
Python
"""Stress-test Natural Memory retrieval on a large in-memory 4B library.
|
|
|
|
This is a routing/reader benchmark, not a language-generation benchmark. It
|
|
builds a multi-million-token address space with sparse project and personal
|
|
facts, then evaluates paraphrased answerable queries, explicit unknown queries,
|
|
and versioned conflict updates. Filler records stay in process RAM and only
|
|
the selected records are promoted to the model's bounded hot cache.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import json
|
|
import random
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from .qwen_integration import load_qwen_dynamic, load_tokenizer
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent
|
|
|
|
|
|
def _path(value: str | Path) -> Path:
|
|
path = Path(value)
|
|
return path if path.is_absolute() or path.exists() else PROJECT_ROOT / path
|
|
|
|
|
|
def _max_memory(gpu_memory_gb: float) -> dict[Any, str] | None:
|
|
if gpu_memory_gb <= 0.0 or not torch.cuda.is_available():
|
|
return None
|
|
return {0: f"{gpu_memory_gb:.1f}GiB", "cpu": "64GiB"}
|
|
|
|
|
|
def _set_cuda_process_cap(gpu_memory_gb: float) -> None:
|
|
if gpu_memory_gb <= 0.0 or not torch.cuda.is_available():
|
|
return
|
|
total = torch.cuda.get_device_properties(0).total_memory
|
|
fraction = min(0.95, max(0.05, gpu_memory_gb * 1024**3 / total))
|
|
torch.cuda.set_per_process_memory_fraction(fraction, device=0)
|
|
|
|
|
|
def _sync(device: torch.device) -> None:
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
def _make_facts(count: int) -> list[dict[str, str]]:
|
|
facts: list[dict[str, str]] = []
|
|
for index in range(count):
|
|
code = f"PROJ-{index:04d}"
|
|
if index % 4 == 0:
|
|
text = f"项目 {code} 的部署区域是 region-{index % 7}，责任服务是 service_{index:04d}。"
|
|
query = f"请查项目 {code} 的部署区域，直接给出 region。"
|
|
elif index % 4 == 1:
|
|
text = f"项目 {code} 的回滚命令是 rollback_{index:04d}，发布窗口为周三。"
|
|
query = f"项目 {code} 出问题时应该执行哪条回滚命令？"
|
|
elif index % 4 == 2:
|
|
text = f"个人偏好 {code}：通知时间设为 {8 + index % 5}:30，提醒渠道为 email。"
|
|
query = f"我在 {code} 里设置的通知时间是多少？"
|
|
else:
|
|
text = f"仓库 {code} 的关键文件是 src/module_{index:04d}.py，入口函数为 run_{index:04d}。"
|
|
query = f"仓库 {code} 的关键入口函数叫什么？"
|
|
facts.append({"code": code, "text": text, "query": query})
|
|
return facts
|
|
|
|
|
|
def _encode_texts(model: Any, tokenizer: Any, texts: list[str], device: torch.device) -> torch.Tensor:
|
|
encoded = tokenizer(
|
|
texts,
|
|
add_special_tokens=False,
|
|
padding=True,
|
|
return_tensors="pt",
|
|
)
|
|
input_ids = encoded["input_ids"].to(device)
|
|
mask = encoded["attention_mask"].to(device)
|
|
return model._encode_model_key(input_ids, mask).detach().cpu()
|
|
|
|
|
|
@torch.inference_mode()
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-model", default=r"H:\Memory")
|
|
parser.add_argument("--memory-model", default=r"H:\Memory\V2_dpskw\qwen3_5_4b_natural_memory_v2")
|
|
parser.add_argument("--output", default=r"H:\Memory\V2_dpskw\large_memory_4b.json")
|
|
parser.add_argument("--library-tokens", type=int, default=4_194_304)
|
|
parser.add_argument("--chunk-tokens", type=int, default=512)
|
|
parser.add_argument("--target-count", type=int, default=64)
|
|
parser.add_argument("--unknown-count", type=int, default=16)
|
|
parser.add_argument("--gpu-memory-gb", type=float, default=10.0)
|
|
parser.add_argument("--no-4bit", action="store_true")
|
|
args = parser.parse_args()
|
|
_set_cuda_process_cap(args.gpu_memory_gb)
|
|
|
|
tokenizer = load_tokenizer(_path(args.base_model))
|
|
use_4bit = not args.no_4bit
|
|
max_memory = _max_memory(args.gpu_memory_gb)
|
|
print("loading Natural Memory 4B")
|
|
model = load_qwen_dynamic(
|
|
_path(args.memory_model),
|
|
load_in_4bit=use_4bit,
|
|
max_memory=max_memory,
|
|
)
|
|
model.eval()
|
|
model.memory_config.memory_top_k_records = 2
|
|
model.memory_os_v2.bank.top_k_records = 2
|
|
device = model._find_layer_device()
|
|
model.reset_memory(batch_size=1, device=device)
|
|
|
|
target_count = max(1, min(int(args.target_count), 256))
|
|
unknown_count = max(1, min(int(args.unknown_count), 128))
|
|
facts = _make_facts(target_count)
|
|
chunk_count = max(1, int(args.library_tokens) // max(1, int(args.chunk_tokens)))
|
|
positions = random.Random(20260905).sample(range(chunk_count), min(target_count, chunk_count))
|
|
position_to_fact = dict(zip(positions, facts))
|
|
|
|
target_keys: dict[str, torch.Tensor] = {}
|
|
for start in range(0, len(facts), 16):
|
|
batch = facts[start : start + 16]
|
|
for fact, key in zip(batch, _encode_texts(model, tokenizer, [item["text"] for item in batch], device)):
|
|
target_keys[fact["code"]] = F.normalize(key.float(), dim=0)
|
|
|
|
generator = torch.Generator(device="cpu").manual_seed(20260905)
|
|
records: list[dict[str, Any]] = []
|
|
for chunk_index in range(chunk_count):
|
|
fact = position_to_fact.get(chunk_index)
|
|
if fact is None:
|
|
key = F.normalize(torch.randn(model.memory.hidden_size, generator=generator), dim=0)
|
|
records.append(
|
|
{
|
|
"text": f"library_chunk:{chunk_index}",
|
|
"key": key,
|
|
"summary": key,
|
|
"semantic_key": key,
|
|
"memory_type": "context_chunk",
|
|
"importance": 0.4,
|
|
"confidence": 0.8,
|
|
"source": "large_memory_benchmark",
|
|
"evidence": [f"library_chunk:{chunk_index}"],
|
|
"trusted": True,
|
|
"force": True,
|
|
}
|
|
)
|
|
continue
|
|
key = target_keys[fact["code"]]
|
|
token_ids = torch.tensor(tokenizer(fact["text"], add_special_tokens=False)["input_ids"], dtype=torch.long)
|
|
records.append(
|
|
{
|
|
"text": fact["text"],
|
|
"key": key,
|
|
"summary": key,
|
|
"semantic_key": key,
|
|
"memory_type": "fact",
|
|
"entity": fact["code"],
|
|
"attribute": "benchmark_fact",
|
|
"value": fact["text"],
|
|
"importance": 0.9,
|
|
"confidence": 0.95,
|
|
"source": "large_memory_benchmark",
|
|
"evidence": [f"library_chunk:{chunk_index}"],
|
|
"token_ids": token_ids,
|
|
"token_mask": torch.ones_like(token_ids, dtype=torch.bool),
|
|
"trusted": True,
|
|
"force": True,
|
|
}
|
|
)
|
|
|
|
_sync(device)
|
|
write_started = time.perf_counter()
|
|
model.memory_os_v2.write_batch(records)
|
|
_sync(device)
|
|
write_ms = (time.perf_counter() - write_started) * 1000.0
|
|
target_record_ids: dict[str, str] = {}
|
|
# Resolve target IDs by text after the batch write. This avoids depending
|
|
# on the bank's stable-ID timestamp while preserving an exact recall test.
|
|
for fact in facts:
|
|
matches = [record.record_id for record in model.memory_os_v2.bank.records.values() if record.text == fact["text"]]
|
|
if matches:
|
|
target_record_ids[fact["code"]] = matches[0]
|
|
|
|
answerable_queries = [fact["query"] for fact in facts]
|
|
unknown_queries = [f"我的不存在的字段 UNKNOWN-{index:04d} 是什么？如果没有登记就说不知道。" for index in range(unknown_count)]
|
|
all_queries = answerable_queries + unknown_queries
|
|
query_keys = _encode_texts(model, tokenizer, all_queries, device)
|
|
answerable_rows: list[dict[str, Any]] = []
|
|
unknown_rows: list[dict[str, Any]] = []
|
|
query_started = time.perf_counter()
|
|
for index, query in enumerate(all_queries):
|
|
encoded = tokenizer(query, add_special_tokens=False, return_tensors="pt")
|
|
query_ids = encoded["input_ids"][0]
|
|
if index < len(answerable_queries):
|
|
fact = facts[index]
|
|
expected_id = target_record_ids.get(fact["code"], "")
|
|
records_out, decision = model.memory_os_v2.read(
|
|
query_key=query_keys[index],
|
|
query_text=query,
|
|
query_token_ids=query_ids,
|
|
top_k_pages=4,
|
|
top_k_records=2,
|
|
max_hops=3,
|
|
)
|
|
returned_ids = [record.record_id for record in records_out]
|
|
answerable_rows.append(
|
|
{
|
|
"code": fact["code"],
|
|
"expected_record_id": expected_id,
|
|
"returned_record_ids": returned_ids,
|
|
"hit": expected_id in returned_ids,
|
|
"stop_reason": decision.stop_reason,
|
|
"coarse_candidates": model.memory_os_v2.bank._last_coarse_candidates,
|
|
}
|
|
)
|
|
else:
|
|
records_out, decision = model.memory_os_v2.read(
|
|
query_key=query_keys[index],
|
|
query_text=query,
|
|
query_token_ids=query_ids,
|
|
top_k_pages=4,
|
|
top_k_records=2,
|
|
max_hops=3,
|
|
)
|
|
unknown_rows.append(
|
|
{
|
|
"query": query,
|
|
"returned_records": len(records_out),
|
|
"abstained": not records_out,
|
|
"stop_reason": decision.stop_reason,
|
|
}
|
|
)
|
|
_sync(device)
|
|
query_ms = (time.perf_counter() - query_started) * 1000.0
|
|
|
|
conflict_entity = "PROJ-CONFLICT"
|
|
conflict_old = model.memory_os_v2.write(
|
|
text="项目 PROJ-CONFLICT 的负责人是 Alice。",
|
|
key=F.normalize(torch.randn(model.memory.hidden_size, generator=generator), dim=0),
|
|
entity=conflict_entity,
|
|
attribute="owner",
|
|
value="Alice",
|
|
source="large_memory_benchmark",
|
|
trusted=True,
|
|
force=True,
|
|
)[0]
|
|
conflict_new = model.memory_os_v2.write(
|
|
text="项目 PROJ-CONFLICT 的负责人是 Bob。",
|
|
key=F.normalize(torch.randn(model.memory.hidden_size, generator=generator), dim=0),
|
|
entity=conflict_entity,
|
|
attribute="owner",
|
|
value="Bob",
|
|
source="large_memory_benchmark",
|
|
trusted=True,
|
|
force=True,
|
|
)[0]
|
|
conflict_query = "项目 PROJ-CONFLICT 当前负责人是谁？"
|
|
conflict_key = _encode_texts(model, tokenizer, [conflict_query], device)[0]
|
|
conflict_out, conflict_decision = model.memory_os_v2.read(
|
|
query_key=conflict_key,
|
|
query_text=conflict_query,
|
|
query_token_ids=torch.tensor(tokenizer(conflict_query, add_special_tokens=False)["input_ids"]),
|
|
top_k_pages=4,
|
|
top_k_records=2,
|
|
max_hops=3,
|
|
)
|
|
conflict_pass = conflict_new.record_id in [record.record_id for record in conflict_out] and conflict_old.status != "active"
|
|
|
|
report = {
|
|
"benchmark": "large_memory_recall_4b",
|
|
"quantization": "4bit_nf4" if use_4bit else "none",
|
|
"library_tokens": int(args.library_tokens),
|
|
"chunk_tokens": int(args.chunk_tokens),
|
|
"chunk_count": chunk_count,
|
|
"stored_records": len(model.memory_os_v2.bank.records),
|
|
"answerable_count": len(answerable_rows),
|
|
"answerable_hits": sum(int(row["hit"]) for row in answerable_rows),
|
|
"answerable_recall": sum(int(row["hit"]) for row in answerable_rows) / max(1, len(answerable_rows)),
|
|
"unknown_count": len(unknown_rows),
|
|
"unknown_abstentions": sum(int(row["abstained"]) for row in unknown_rows),
|
|
"unknown_abstention_rate": sum(int(row["abstained"]) for row in unknown_rows) / max(1, len(unknown_rows)),
|
|
"conflict_update_pass": conflict_pass,
|
|
"conflict_stop_reason": conflict_decision.stop_reason,
|
|
"write_ms": write_ms,
|
|
"query_total_ms": query_ms,
|
|
"query_mean_ms": query_ms / max(1, len(all_queries)),
|
|
"max_coarse_candidates": max((row["coarse_candidates"] for row in answerable_rows), default=0),
|
|
"max_vram_gb": torch.cuda.max_memory_allocated(device) / 1024**3 if device.type == "cuda" else None,
|
|
"conflict_old_status": conflict_old.status,
|
|
"conflict_new_status": conflict_new.status,
|
|
"conflict_returned_ids": [record.record_id for record in conflict_out],
|
|
"answerable_rows": answerable_rows,
|
|
"unknown_rows": unknown_rows,
|
|
"warning": "Routing test uses a sparse lexical address side-index plus v9 neural reranking; generation quality still needs real chat/repository workloads.",
|
|
}
|
|
output = _path(args.output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
print(f"saved={output}")
|
|
del model
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|