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,725 @@
|
||||
"""Full local comparison between Qwen3.5-4B and Natural Memory v1.
|
||||
|
||||
This is an engineering benchmark, not a claim of state-of-the-art performance.
|
||||
It evaluates the same frozen Qwen3.5-4B backbone with and without the internal
|
||||
memory path, using deterministic greedy decoding and locally generated cases.
|
||||
The report includes general ability, extra math and reasoning cases, long
|
||||
context retrieval, throughput, latency, VRAM, automatic write decisions,
|
||||
conflict replacement, unknown-fact refusal, reset, and shard-backed restart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import statistics
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import (
|
||||
load_memory_config,
|
||||
load_qwen_base,
|
||||
load_qwen_dynamic,
|
||||
load_tokenizer,
|
||||
)
|
||||
from .stream_chat_qwen_memory import _chat_tensor, _memory_system_prefix, _write_turn
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
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_memory_merged_v13",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data",
|
||||
default=r"H:\Memory\V2_dpskw\data\comprehensive_general.jsonl",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default=r"H:\Memory\V2_dpskw\natural_memory_v1_full_benchmark.json",
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=24)
|
||||
parser.add_argument("--perf-repeats", type=int, default=3)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def normalize(text: str) -> str:
|
||||
keep = str(text).lower()
|
||||
for char in " \t\r\n`*_#，。！？、；：,.!?;:'\"（）()[]{}<>|\\/：":
|
||||
keep = keep.replace(char, "")
|
||||
return keep
|
||||
|
||||
|
||||
def contains_answer(text: str, acceptable: list[str]) -> bool:
|
||||
normalized = normalize(text)
|
||||
for answer in acceptable:
|
||||
expected = normalize(str(answer))
|
||||
if not expected:
|
||||
continue
|
||||
if expected.isdigit() and len(expected) == 1:
|
||||
if any(
|
||||
normalized[index : index + 1] == expected
|
||||
and (index == 0 or not normalized[index - 1].isdigit())
|
||||
and (index + 1 == len(normalized) or not normalized[index + 1].isdigit())
|
||||
for index in range(len(normalized))
|
||||
):
|
||||
return True
|
||||
elif expected in normalized:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def load_jsonl(path: str | Path) -> list[dict[str, Any]]:
|
||||
rows = []
|
||||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||||
if line.strip():
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
|
||||
|
||||
def prompt_inputs(tokenizer: Any, prompt: str, device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt}],
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
return_dict=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return {
|
||||
key: value.to(device)
|
||||
for key, value in encoded.items()
|
||||
if isinstance(value, torch.Tensor)
|
||||
}
|
||||
|
||||
|
||||
def input_token_count(tokenizer: Any, prompt: str) -> int:
|
||||
encoded = tokenizer(prompt, add_special_tokens=False)
|
||||
return len(encoded["input_ids"])
|
||||
|
||||
|
||||
def generate_answer(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
prompt: str,
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
) -> tuple[str, int, float]:
|
||||
device = model._find_layer_device() if dynamic else model.get_input_embeddings().weight.device
|
||||
encoded = prompt_inputs(tokenizer, prompt, device)
|
||||
query = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
|
||||
query_ids = query["input_ids"].to(device)
|
||||
query_mask = query.get("attention_mask")
|
||||
if query_mask is None:
|
||||
query_mask = torch.ones_like(query_ids)
|
||||
query_mask = query_mask.to(device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
kwargs: dict[str, Any] = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"do_sample": False,
|
||||
"use_cache": True,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
if dynamic:
|
||||
kwargs.update(
|
||||
{
|
||||
"update_memory": False,
|
||||
"memory_query_input_ids": query_ids,
|
||||
"memory_query_attention_mask": query_mask,
|
||||
}
|
||||
)
|
||||
output = model.generate(**encoded, **kwargs)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
prompt_len = int(encoded["input_ids"].shape[1])
|
||||
response_ids = output[0, prompt_len:]
|
||||
response = tokenizer.decode(response_ids.detach().cpu().tolist(), skip_special_tokens=True).strip()
|
||||
return response, int(response_ids.numel()), elapsed
|
||||
|
||||
|
||||
def make_math_cases() -> list[dict[str, Any]]:
|
||||
raw = [
|
||||
("m01", "只输出最终整数：38 + 47 = ?", ["85"]),
|
||||
("m02", "只输出最终整数：900 - 376 = ?", ["524"]),
|
||||
("m03", "只输出最终整数：24 × 17 = ?", ["408"]),
|
||||
("m04", "只输出最终整数：936 ÷ 18 = ?", ["52"]),
|
||||
("m05", "只输出结果：2.75 + 3.6 = ?", ["6.35"]),
|
||||
("m06", "只输出结果：7/8 - 1/4 = ?", ["5/8", "0.625"]),
|
||||
("m07", "只输出整数：15 和 28 的最小公倍数是多少？", ["420"]),
|
||||
("m08", "只输出百分数：480 的 12.5% 是多少？", ["60"]),
|
||||
("m09", "只输出百分数：80 增长到 100，增长率是多少？", ["25%", "25"]),
|
||||
("m10", "只输出结果：3 的 5 次方是多少？", ["243"]),
|
||||
("m11", "只输出 x：4x + 7 = 31。", ["6"]),
|
||||
("m12", "只输出 x：9x - 18 = 0。", ["2"]),
|
||||
("m13", "只输出 x：2(x + 5) = 18。", ["4"]),
|
||||
("m14", "只输出 x：x/3 + 4 = 9。", ["15"]),
|
||||
("m15", "只输出 x：5x - 2 = 3x + 10。", ["6"]),
|
||||
("m16", "只输出下一个数：5，10，20，40，？", ["80"]),
|
||||
("m17", "只输出下一个数：3，6，11，18，27，？", ["38"]),
|
||||
("m18", "只输出下一个数：1，4，9，16，？", ["25"]),
|
||||
("m19", "只输出下一个数：2，3，5，8，12，？", ["17"]),
|
||||
("m20", "只输出整数：阶乘 6! 等于多少？", ["720"]),
|
||||
("m21", "只输出面积：长 12、宽 7 的矩形面积是多少？", ["84"]),
|
||||
("m22", "只输出周长：边长为 9 的正方形周长是多少？", ["36"]),
|
||||
("m23", "只输出面积：底 10、高 6 的三角形面积是多少？", ["30"]),
|
||||
("m24", "只输出角度：一个三角形两个角是 35 度和 65 度，第三个角是多少？", ["80"]),
|
||||
("m25", "只输出数量：3 件不同衬衫和 2 条不同裤子可以组成多少套穿搭？", ["6"]),
|
||||
("m26", "只输出数量：从 5 个人中选 2 个人，有多少种选法？", ["10"]),
|
||||
("m27", "只输出余数：17 除以 5 的余数是多少？", ["2"]),
|
||||
("m28", "只输出结果：平均数 8、12、16、20 是多少？", ["14"]),
|
||||
("m29", "只输出结果：一个商品原价 240 元，打八折后多少钱？", ["192"]),
|
||||
("m30", "只输出结果：2.4 × 0.5 = ?", ["1.2"]),
|
||||
]
|
||||
return [
|
||||
{"id": case_id, "category": "math", "prompt": prompt, "acceptable": answers}
|
||||
for case_id, prompt, answers in raw
|
||||
]
|
||||
|
||||
|
||||
def make_reasoning_cases() -> list[dict[str, Any]]:
|
||||
raw = [
|
||||
("r01", "只输出名字：甲比乙早到，乙比丙早到，谁最后到？", ["丙"]),
|
||||
("r02", "只输出名字：小李在小王左边，小王在小张左边，谁最右边？", ["小张"]),
|
||||
("r03", "只输出结论：所有鸟都有翅膀，企鹅是鸟，所以企鹅有翅膀吗？", ["是"]),
|
||||
("r04", "只输出结论：所有猫都是哺乳动物，鲸鱼是哺乳动物，所以鲸鱼是猫吗？", ["不是", "否"]),
|
||||
("r05", "只输出结论：如果下雨就带伞。现在下雨了，要不要带伞？", ["要"]),
|
||||
("r06", "只输出结论：只有持票者才能入场。小林没有票，他能入场吗？", ["不能", "不可以"]),
|
||||
("r07", "只输出星期：今天是星期三，五天后是星期几？", ["星期一", "周一"]),
|
||||
("r08", "只输出方向：你面向北，右转后面向哪个方向？", ["东"]),
|
||||
("r09", "只输出方向：你面向东，左转后面向哪个方向？", ["北"]),
|
||||
("r10", "只输出数量：盒子里有 4 个红球和 3 个蓝球，不看颜色拿出一个，至少有几个球？", ["1"]),
|
||||
("r11", "只输出名字：甲不是第一，乙在甲前面，丙在乙后面，谁可能是第一？", ["乙"]),
|
||||
("r12", "只输出结论：有些学生会游泳，小周是学生，能确定小周会游泳吗？", ["不能", "无法"]),
|
||||
("r13", "只输出下一个数：1，2，4，7，11，？", ["16"]),
|
||||
("r14", "只输出下一个数：81，27，9，3，？", ["1"]),
|
||||
("r15", "只输出名字：红色比蓝色重，绿色比红色轻但比蓝色重，哪个最轻？", ["蓝色"]),
|
||||
("r16", "只输出答案：苹果不是蔬菜，胡萝卜是蔬菜，香蕉是水果，哪个不是水果？", ["胡萝卜"]),
|
||||
("r17", "只输出结论：如果 A 大于 B 且 B 大于 C，那么 A 大于 C 吗？", ["是"]),
|
||||
("r18", "只输出结论：如果一个数能被 2 整除，它一定是偶数。14 能被 2 整除，它是偶数吗？", ["是"]),
|
||||
("r19", "只输出名字：小赵比小钱高，小孙比小赵矮但比小钱高，谁最高？", ["小赵"]),
|
||||
("r20", "只输出数量：一周中有几天的名字包含‘星’字？", ["7"]),
|
||||
("r21", "只输出结论：没有鱼是鸟，金鱼是鱼，所以金鱼是鸟吗？", ["不是", "否"]),
|
||||
("r22", "只输出顺序：春、夏、秋、冬之后又回到哪个季节？", ["春"]),
|
||||
("r23", "只输出结论：所有密码都需要保密，这个字符串是密码，所以它需要保密吗？", ["是"]),
|
||||
("r24", "只输出答案：小明有两个兄弟，每个兄弟都有一个姐姐，小明有几个姐姐？", ["1"]),
|
||||
]
|
||||
return [
|
||||
{"id": case_id, "category": "reasoning", "prompt": prompt, "acceptable": answers}
|
||||
for case_id, prompt, answers in raw
|
||||
]
|
||||
|
||||
|
||||
def make_context_cases(tokenizer: Any) -> list[dict[str, Any]]:
|
||||
rng = random.Random(20260904)
|
||||
filler = (
|
||||
"这是一段与问题无关的背景说明。系统记录了版本号、构建时间、测试批次、"
|
||||
"设备温度、日志摘要和普通项目备注。这些文字只是干扰项，不包含目标答案。"
|
||||
)
|
||||
cases: list[dict[str, Any]] = []
|
||||
for target_tokens in (512, 2048, 4096, 8192):
|
||||
for position in ("early", "middle", "late"):
|
||||
answer = f"CTX{target_tokens}-{position.upper()}-{rng.randrange(100, 999)}"
|
||||
needle = f"目标记录：本次检索需要返回的项目编码是 {answer}。"
|
||||
chunks: list[str] = []
|
||||
while input_token_count(tokenizer, " ".join(chunks + [filler, needle])) < target_tokens:
|
||||
chunks.append(filler)
|
||||
if position == "early":
|
||||
material = " ".join([needle] + chunks)
|
||||
elif position == "middle":
|
||||
half = len(chunks) // 2
|
||||
material = " ".join(chunks[:half] + [needle] + chunks[half:])
|
||||
else:
|
||||
material = " ".join(chunks + [needle])
|
||||
prompt = (
|
||||
"请阅读下面的材料，只输出目标记录中的项目编码，不要解释。\n"
|
||||
"---材料开始---\n"
|
||||
f"{material}\n"
|
||||
"---材料结束---\n"
|
||||
"问题：目标记录中的项目编码是什么？"
|
||||
)
|
||||
cases.append(
|
||||
{
|
||||
"id": f"ctx-{target_tokens}-{position}",
|
||||
"category": f"context_{target_tokens}",
|
||||
"prompt": prompt,
|
||||
"acceptable": [answer],
|
||||
"target_tokens": target_tokens,
|
||||
"position": position,
|
||||
}
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
def evaluate_cases(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
cases: list[dict[str, Any]],
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
categories: dict[str, list[float]] = {}
|
||||
started = time.perf_counter()
|
||||
for case in cases:
|
||||
if dynamic:
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
response, generated_tokens, elapsed = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
str(case["prompt"]),
|
||||
dynamic=dynamic,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
passed = contains_answer(response, list(case["acceptable"]))
|
||||
category = str(case["category"])
|
||||
categories.setdefault(category, []).append(float(passed))
|
||||
rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"category": category,
|
||||
"prompt_tokens": input_token_count(tokenizer, str(case["prompt"])),
|
||||
"acceptable": case["acceptable"],
|
||||
"generated": response,
|
||||
"generated_tokens": generated_tokens,
|
||||
"seconds": elapsed,
|
||||
"passed": passed,
|
||||
}
|
||||
)
|
||||
total = sum(sum(values) for values in categories.values())
|
||||
return {
|
||||
"cases": len(rows),
|
||||
"elapsed_seconds": time.perf_counter() - started,
|
||||
"overall_score": total / max(1, len(rows)),
|
||||
"categories": {
|
||||
category: {
|
||||
"count": len(values),
|
||||
"score": sum(values) / max(1, len(values)),
|
||||
}
|
||||
for category, values in sorted(categories.items())
|
||||
},
|
||||
"rows": rows,
|
||||
}
|
||||
|
||||
|
||||
def device_snapshot(model: Any) -> dict[str, Any]:
|
||||
device = model._find_layer_device() if hasattr(model, "_find_layer_device") else model.get_input_embeddings().weight.device
|
||||
params = sum(parameter.numel() for parameter in model.parameters())
|
||||
result: dict[str, Any] = {
|
||||
"device": str(device),
|
||||
"parameter_count": int(params),
|
||||
"parameter_count_billion": params / 1e9,
|
||||
}
|
||||
if device.type == "cuda":
|
||||
properties = torch.cuda.get_device_properties(device)
|
||||
result.update(
|
||||
{
|
||||
"gpu_name": properties.name,
|
||||
"gpu_total_memory_gb": properties.total_memory / 1024**3,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def measure_performance(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
cases: list[dict[str, Any]],
|
||||
*,
|
||||
dynamic: bool,
|
||||
max_new_tokens: int,
|
||||
repeats: int,
|
||||
) -> dict[str, Any]:
|
||||
selected = [cases[0]]
|
||||
for wanted in (512, 2048, 4096):
|
||||
matching = [case for case in cases if case.get("target_tokens") == wanted]
|
||||
if matching:
|
||||
selected.append(matching[1])
|
||||
rows = []
|
||||
device = model._find_layer_device() if dynamic else model.get_input_embeddings().weight.device
|
||||
for case in selected:
|
||||
latencies = []
|
||||
generated = 0
|
||||
for _ in range(max(1, repeats)):
|
||||
if dynamic:
|
||||
model.reset_memory(batch_size=1, device=device)
|
||||
_, token_count, elapsed = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
str(case["prompt"]),
|
||||
dynamic=dynamic,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
latencies.append(elapsed)
|
||||
generated += token_count
|
||||
rows.append(
|
||||
{
|
||||
"id": case["id"],
|
||||
"prompt_tokens": input_token_count(tokenizer, str(case["prompt"])),
|
||||
"median_seconds": statistics.median(latencies),
|
||||
"mean_seconds": statistics.mean(latencies),
|
||||
"tokens_per_second": generated / max(1e-9, sum(latencies)),
|
||||
"repeats": len(latencies),
|
||||
}
|
||||
)
|
||||
return {"rows": rows}
|
||||
|
||||
|
||||
def memory_payload(model: Any) -> dict[str, torch.Tensor]:
|
||||
payload: dict[str, torch.Tensor] = {
|
||||
"memory_state": model.runtime.state.detach().cpu().clone(),
|
||||
}
|
||||
if model.memory_config.natural_language_memory:
|
||||
for name in (
|
||||
"text_token_ids",
|
||||
"text_token_mask",
|
||||
"text_slot_valid",
|
||||
"text_slot_keys",
|
||||
"text_slot_age",
|
||||
"text_write_counter",
|
||||
"text_key_token_ids",
|
||||
"text_key_token_mask",
|
||||
):
|
||||
value = getattr(model.runtime, name)
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise RuntimeError(f"runtime memory field is unavailable: {name}")
|
||||
payload[name] = value.detach().cpu().clone()
|
||||
return payload
|
||||
|
||||
|
||||
def answer_memory_query(model: Any, tokenizer: Any, text: str, max_new_tokens: int) -> tuple[str, bool]:
|
||||
response, _, _ = generate_answer(
|
||||
model,
|
||||
tokenizer,
|
||||
text,
|
||||
dynamic=True,
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
return response, bool(model.runtime.text_prefix_used)
|
||||
|
||||
|
||||
def run_memory_benchmark(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
package_path: str | Path,
|
||||
*,
|
||||
max_new_tokens: int,
|
||||
no_4bit: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Run memory tests and restore the user's original embedded state."""
|
||||
|
||||
original = memory_payload(model)
|
||||
decision_rows = []
|
||||
positives = [
|
||||
"我叫林舟。",
|
||||
"我的常住城市是苏州。",
|
||||
"我最喜欢的水果是红富士苹果。",
|
||||
"我正在开发 Natural Memory v1 项目。",
|
||||
"以后请把代码默认写成 Python。",
|
||||
"我的常用时区是 Asia/Shanghai。",
|
||||
"这是我的长期偏好：使用简洁中文。",
|
||||
"这个项目的重要约束是不要修改原始 Qwen 权重。",
|
||||
]
|
||||
negatives = [
|
||||
"我叫什么？",
|
||||
"帮我解释向量数据库是什么。",
|
||||
"你觉得今天的天气怎么样？",
|
||||
"请把 memory 翻译成中文。",
|
||||
"计算一下 17 × 19。",
|
||||
"如果我选择 GPU，会发生什么？",
|
||||
"我之前有没有提到我的城市？",
|
||||
"给我一个自然语言记忆方案。",
|
||||
]
|
||||
try:
|
||||
for expected, items in ((True, positives), (False, negatives)):
|
||||
for text in items:
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
changed = _write_turn(model, tokenizer, text, model._find_layer_device())
|
||||
decision_rows.append(
|
||||
{
|
||||
"text": text,
|
||||
"expected_write": expected,
|
||||
"actual_write": changed,
|
||||
"write_probability": float(model.runtime.auto_memory_probability.mean())
|
||||
if isinstance(model.runtime.auto_memory_probability, torch.Tensor)
|
||||
else None,
|
||||
}
|
||||
)
|
||||
tp = sum(row["expected_write"] and row["actual_write"] for row in decision_rows)
|
||||
tn = sum((not row["expected_write"]) and (not row["actual_write"]) for row in decision_rows)
|
||||
fp = sum((not row["expected_write"]) and row["actual_write"] for row in decision_rows)
|
||||
fn = sum(row["expected_write"] and (not row["actual_write"]) for row in decision_rows)
|
||||
|
||||
model.reset_memory(batch_size=1, device=model._find_layer_device())
|
||||
first = "我的工作地点代号是NM-R7。"
|
||||
second = "我最喜欢的水果是青提。"
|
||||
replacement = "我的工作地点代号改为NM-K9。"
|
||||
writes = [
|
||||
{"text": first, "changed": _write_turn(model, tokenizer, first, model._find_layer_device())},
|
||||
{"text": second, "changed": _write_turn(model, tokenizer, second, model._find_layer_device())},
|
||||
{
|
||||
"text": replacement,
|
||||
"changed": _write_turn(model, tokenizer, replacement, model._find_layer_device()),
|
||||
},
|
||||
]
|
||||
before_restart = {}
|
||||
for name, query, expected in (
|
||||
("work_code", "我的工作地点代号是什么？", "NM-K9"),
|
||||
("fruit", "我最喜欢的水果是什么？", "青提"),
|
||||
("unknown", "我的血型是什么？如果没有记录，请明确说不知道。", "不知道"),
|
||||
):
|
||||
response, prefix_used = answer_memory_query(model, tokenizer, query, max_new_tokens)
|
||||
before_restart[name] = {
|
||||
"query": query,
|
||||
"expected": expected,
|
||||
"response": response,
|
||||
"expected_found": expected in response,
|
||||
"prefix_used": prefix_used,
|
||||
}
|
||||
|
||||
# Persist through ordinary natural-language turns. The caller releases
|
||||
# this model before loading a fresh process/model for the restart test;
|
||||
# keeping that lifecycle outside this function avoids two 4-bit Qwen
|
||||
# backbones occupying the GPU at the same time.
|
||||
model.save_embedded_memory_weights(package_path)
|
||||
return {
|
||||
"automatic_write_decision": {
|
||||
"rows": decision_rows,
|
||||
"true_positive": int(tp),
|
||||
"true_negative": int(tn),
|
||||
"false_positive": int(fp),
|
||||
"false_negative": int(fn),
|
||||
"precision": tp / max(1, tp + fp),
|
||||
"recall": tp / max(1, tp + fn),
|
||||
"specificity": tn / max(1, tn + fp),
|
||||
},
|
||||
"natural_language_writes": writes,
|
||||
"before_restart": before_restart,
|
||||
"embedded_write_persisted": True,
|
||||
}
|
||||
except Exception:
|
||||
# Best-effort restoration if a test fails halfway through.
|
||||
model._load_persistent_memory_payload(original)
|
||||
model.save_embedded_memory_weights(package_path)
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
raise
|
||||
|
||||
|
||||
def release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
tokenizer = load_tokenizer(args.base_model)
|
||||
memory_config_hint = load_memory_config(args.memory_model)
|
||||
memory_variant = (
|
||||
"natural_memory_v2"
|
||||
if memory_config_hint.memory_version >= 2 or memory_config_hint.hierarchical_memory
|
||||
else "natural_memory_v1"
|
||||
)
|
||||
general = load_jsonl(args.data)
|
||||
math_cases = make_math_cases()
|
||||
reasoning_cases = make_reasoning_cases()
|
||||
context_cases = make_context_cases(tokenizer)
|
||||
all_cases = general + math_cases + reasoning_cases + context_cases
|
||||
use_4bit = not args.no_4bit
|
||||
report: dict[str, Any] = {
|
||||
"benchmark": f"Natural Memory {memory_variant.rsplit('_', 1)[-1]} full local comparison",
|
||||
"date": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"seed": 20260904,
|
||||
"base_model": str(Path(args.base_model).resolve()),
|
||||
"memory_model": str(Path(args.memory_model).resolve()),
|
||||
"data": str(Path(args.data).resolve()),
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
"decoding": {"do_sample": False, "max_new_tokens": args.max_new_tokens},
|
||||
"case_counts": {
|
||||
"general_existing": len(general),
|
||||
"math": len(math_cases),
|
||||
"reasoning": len(reasoning_cases),
|
||||
"context": len(context_cases),
|
||||
"total": len(all_cases),
|
||||
},
|
||||
"context_targets": sorted({case["target_tokens"] for case in context_cases}),
|
||||
}
|
||||
|
||||
print(f"cases={len(all_cases)} quantization={report['quantization']}")
|
||||
print("loading Qwen3.5-4B baseline")
|
||||
started = time.perf_counter()
|
||||
baseline = load_qwen_base(args.base_model, load_in_4bit=use_4bit)
|
||||
baseline.eval()
|
||||
report["baseline"] = {
|
||||
"load_seconds": time.perf_counter() - started,
|
||||
"hardware": device_snapshot(baseline),
|
||||
}
|
||||
baseline_device = baseline.get_input_embeddings().weight.device
|
||||
if baseline_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(baseline_device)
|
||||
report["baseline"]["quality"] = evaluate_cases(
|
||||
baseline,
|
||||
tokenizer,
|
||||
all_cases,
|
||||
dynamic=False,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
)
|
||||
report["baseline"]["performance"] = measure_performance(
|
||||
baseline,
|
||||
tokenizer,
|
||||
context_cases,
|
||||
dynamic=False,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
repeats=args.perf_repeats,
|
||||
)
|
||||
if baseline_device.type == "cuda":
|
||||
report["baseline"]["peak_memory_allocated_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
||||
report["baseline"]["peak_memory_reserved_gb"] = torch.cuda.max_memory_reserved(baseline_device) / 1024**3
|
||||
release(baseline)
|
||||
|
||||
print(f"loading {memory_variant} embedded package")
|
||||
started = time.perf_counter()
|
||||
dynamic = load_qwen_dynamic(args.memory_model, load_in_4bit=use_4bit)
|
||||
dynamic.eval()
|
||||
report[memory_variant] = {
|
||||
"load_seconds": time.perf_counter() - started,
|
||||
"hardware": device_snapshot(dynamic),
|
||||
}
|
||||
dynamic_device = dynamic._find_layer_device()
|
||||
if dynamic_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(dynamic_device)
|
||||
report[memory_variant]["quality"] = evaluate_cases(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
all_cases,
|
||||
dynamic=True,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
)
|
||||
report[memory_variant]["performance"] = measure_performance(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
context_cases,
|
||||
dynamic=True,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
repeats=args.perf_repeats,
|
||||
)
|
||||
if dynamic_device.type == "cuda":
|
||||
report[memory_variant]["peak_memory_allocated_gb"] = torch.cuda.max_memory_allocated(dynamic_device) / 1024**3
|
||||
report[memory_variant]["peak_memory_reserved_gb"] = torch.cuda.max_memory_reserved(dynamic_device) / 1024**3
|
||||
|
||||
original_embedded_payload = memory_payload(dynamic)
|
||||
memory_report = run_memory_benchmark(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
args.memory_model,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
no_4bit=args.no_4bit,
|
||||
)
|
||||
# Release the first Qwen backbone before constructing the fresh model used
|
||||
# by the shard-backed restart check. This is important on a 12 GB GPU.
|
||||
release(dynamic)
|
||||
dynamic = None
|
||||
restarted = None
|
||||
try:
|
||||
restarted = load_qwen_dynamic(args.memory_model, load_in_4bit=use_4bit)
|
||||
restarted.eval()
|
||||
after_restart = {}
|
||||
for name, query, expected in (
|
||||
("work_code", "我的工作地点代号是什么？", "NM-K9"),
|
||||
("fruit", "我最喜欢的水果是什么？", "青提"),
|
||||
):
|
||||
response, prefix_used = answer_memory_query(restarted, tokenizer, query, args.max_new_tokens)
|
||||
after_restart[name] = {
|
||||
"query": query,
|
||||
"expected": expected,
|
||||
"response": response,
|
||||
"expected_found": expected in response,
|
||||
"prefix_used": prefix_used,
|
||||
}
|
||||
restarted.reset_memory(batch_size=1, device=restarted._find_layer_device())
|
||||
reset_response, reset_prefix = answer_memory_query(
|
||||
restarted,
|
||||
tokenizer,
|
||||
"我的工作地点代号是什么？",
|
||||
args.max_new_tokens,
|
||||
)
|
||||
reset_slots = int(restarted.runtime.text_slot_valid.sum().item())
|
||||
memory_report.update(
|
||||
{
|
||||
"after_restart_without_history_or_pt": after_restart,
|
||||
"restart_pass": all(item["expected_found"] for item in after_restart.values()),
|
||||
"reset": {
|
||||
"response": reset_response,
|
||||
"prefix_used": reset_prefix,
|
||||
"valid_slots": reset_slots,
|
||||
"cleared": reset_slots == 0 and not reset_prefix,
|
||||
},
|
||||
}
|
||||
)
|
||||
finally:
|
||||
if restarted is not None:
|
||||
# Restore the user's pre-benchmark state, so the benchmark itself
|
||||
# does not overwrite the active embedded memory snapshot.
|
||||
restarted._load_persistent_memory_payload(original_embedded_payload)
|
||||
restarted.save_embedded_memory_weights(args.memory_model)
|
||||
memory_report["state_restored"] = int(restarted.runtime.text_slot_valid.sum().item()) == int(
|
||||
original_embedded_payload["text_slot_valid"].sum().item()
|
||||
)
|
||||
release(restarted)
|
||||
report["memory"] = memory_report
|
||||
|
||||
baseline_quality = report["baseline"]["quality"]
|
||||
dynamic_quality = report[memory_variant]["quality"]
|
||||
categories = sorted(
|
||||
set(baseline_quality["categories"]) & set(dynamic_quality["categories"])
|
||||
)
|
||||
category_deltas = {
|
||||
category: dynamic_quality["categories"][category]["score"]
|
||||
- baseline_quality["categories"][category]["score"]
|
||||
for category in categories
|
||||
}
|
||||
report["comparison"] = {
|
||||
"overall_delta": dynamic_quality["overall_score"] - baseline_quality["overall_score"],
|
||||
"category_deltas": category_deltas,
|
||||
"peak_memory_allocated_delta_gb": report[memory_variant].get("peak_memory_allocated_gb", 0.0)
|
||||
- report["baseline"].get("peak_memory_allocated_gb", 0.0),
|
||||
"load_seconds_delta": report[memory_variant]["load_seconds"]
|
||||
- report["baseline"]["load_seconds"],
|
||||
}
|
||||
|
||||
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")
|
||||
compact = {
|
||||
"baseline_score": baseline_quality["overall_score"],
|
||||
f"{memory_variant}_score": dynamic_quality["overall_score"],
|
||||
"overall_delta": report["comparison"]["overall_delta"],
|
||||
"memory_restart_pass": report["memory"]["restart_pass"],
|
||||
"memory_reset_pass": report["memory"]["reset"]["cleared"],
|
||||
"cases": len(all_cases),
|
||||
"output": str(output),
|
||||
}
|
||||
print(json.dumps(compact, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user