wpy/natural-memory-nm21
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据

Browse Source 
- 引入 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
2026-09-19 11:11:31 +08:00
commit 643e22ecb9
484 changed files with 306821 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark_service_curves_4b.py
+497
View File
@@ -0,0 +1,497 @@
"""Compare long-context service curves on the local Qwen3.5 4B checkpoint.
This benchmark is intentionally framed around the production baselines that a
serving team would compare:
* dense/paged full prompt KV (until the checkpoint or GPU rejects it),
* a fixed sliding window,
* chunk retrieval followed by a short prompt-side rerank context, and
* Natural Memory with one bounded read at the request boundary.
The Natural Memory path never performs a memory read or write per generated
token. The reported ``reader_ms`` is measured inside the model wrapper and
``decode_tok_s`` is estimated from a one-token and a multi-token greedy run.
The synthetic task is deliberately simple and should be supplemented by
repository QA and agent-trace workloads before making a production claim.
"""
from __future__ import annotations
import argparse
import gc
import json
import math
import random
import time
from pathlib import Path
from typing import Any
import torch
from .qwen_integration import load_qwen_base, load_qwen_dynamic, load_tokenizer
from .stream_chat_qwen_memory import _chat_tensor
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 _sync(device: torch.device) -> None:
if device.type == "cuda":
torch.cuda.synchronize(device)
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:
"""Keep the benchmark inside the declared HBM operating point."""
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 _build_corpus(
tokenizer: Any,
target_tokens: int,
*,
chunk_tokens: int,
) -> tuple[torch.Tensor, torch.Tensor, str, str]:
"""Create a token-level corpus with a middle needle and a query."""
filler = (
"这是长上下文压力测试中的普通项目日志片段。它包含无关的版本号、"
"状态说明、时间戳和重复背景，不应被当作目标答案。"
)
needle = "用户档案字段 target_id 的值是 CURVE-7H2K-9P4M。"
query = "用户档案字段 target_id 的完整值是什么？请逐字符复制，不要省略。"
filler_ids = torch.tensor(
tokenizer(filler, add_special_tokens=False)["input_ids"], dtype=torch.long
)
needle_ids = torch.tensor(
tokenizer(needle, add_special_tokens=False)["input_ids"], dtype=torch.long
)
target_tokens = max(int(target_tokens), int(needle_ids.numel()) + 8)
filler_count = target_tokens - int(needle_ids.numel())
repeats = math.ceil(filler_count / max(1, filler_ids.numel()))
body = filler_ids.repeat(repeats)[:filler_count]
# Keep the needle outside a 32K hot window, while making it present in
# every measured prefix. The long-context axis then tests persistence,
# rather than accidentally testing a missing fact at the short end.
insert_at = min((body.numel() // 2) // max(1, chunk_tokens) * max(1, chunk_tokens), 65536)
corpus = torch.cat((body[:insert_at], needle_ids, body[insert_at:]), dim=0)
return corpus, needle_ids, needle, query
def _chat_from_content(tokenizer: Any, content: str, device: torch.device) -> dict[str, torch.Tensor]:
encoded = _chat_tensor(tokenizer, content)
return {key: value.to(device) for key, value in encoded.items()}
def _query_tensors(tokenizer: Any, query: str, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
encoded = tokenizer(query, add_special_tokens=False, return_tensors="pt")
ids = encoded["input_ids"].to(device)
mask = encoded.get("attention_mask")
if mask is None:
mask = torch.ones_like(ids)
return ids, mask.to(device)
@torch.inference_mode()
def _measure_generation(
model: Any,
encoded: dict[str, torch.Tensor],
tokenizer: Any,
device: torch.device,
*,
dynamic: bool,
query_ids: torch.Tensor | None = None,
query_mask: torch.Tensor | None = None,
query_text: str = "",
max_new_tokens: int = 16,
) -> dict[str, Any]:
def generate(count: int):
kwargs: dict[str, Any] = {
"max_new_tokens": count,
"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,
"memory_query_text": query_text,
}
)
_sync(device)
started = time.perf_counter()
output = model.generate(**encoded, **kwargs)
_sync(device)
return output, time.perf_counter() - started
try:
one, first_latency = generate(1)
many, total_latency = generate(max_new_tokens)
except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
if isinstance(exc, torch.cuda.OutOfMemoryError) or "out of memory" in str(exc).lower():
if device.type == "cuda":
torch.cuda.empty_cache()
return {"status": "cuda_oom", "error": str(exc)[:500]}
raise
generated = many[0, encoded["input_ids"].shape[1] :]
text = tokenizer.decode(generated.detach().cpu().tolist(), skip_special_tokens=True).strip()
extra = max(0.0, total_latency - first_latency)
generated_count = int(generated.numel())
runtime = getattr(model, "runtime", None)
prefix_tokens = int(getattr(runtime, "text_prefix_tokens", 0)) if dynamic else 0
public_input_tokens = int(encoded["input_ids"].shape[1])
row = {
"status": "ok",
"response": text,
"generated_tokens": generated_count,
"first_token_latency_s": first_latency,
"prefill_proxy_s": first_latency,
"total_latency_s": total_latency,
"batch_size": 1,
"public_input_tokens": public_input_tokens,
"memory_prefix_tokens": prefix_tokens,
"hot_kv_tokens": public_input_tokens + prefix_tokens,
"decode_tok_s": max(0, generated_count - 1) / max(extra, 1e-9),
"reader_ms": float(getattr(getattr(model, "runtime", None), "text_read_seconds", 0.0) * 1000.0)
if dynamic
else 0.0,
}
row["gpu_seconds_per_million_output_tokens"] = (
1_000_000.0 / max(1e-9, row["decode_tok_s"])
)
if device.type == "cuda":
row["peak_vram_gb"] = torch.cuda.max_memory_allocated(device) / 1024**3
row["peak_reserved_gb"] = torch.cuda.max_memory_reserved(device) / 1024**3
del one, many
torch.cuda.empty_cache()
return row
def _release(model: Any) -> None:
del model
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def _quality_fields(row: dict[str, Any], *, target_reachable: bool, expected: str) -> dict[str, Any]:
"""Separate answer correctness from whether the target was in the hot context."""
found = expected in row.get("response", "")
row["target_reachable"] = bool(target_reachable)
row["quality_pass"] = found
row["quality_expected"] = bool(target_reachable)
row["quality_correct"] = bool(found == target_reachable) if row.get("status") == "ok" else False
return row
def _prompt_content(tokenizer: Any, corpus: torch.Tensor, query: str) -> str:
material = tokenizer.decode(corpus.tolist(), skip_special_tokens=True)
return (
"请从下面的材料中回答问题，不要使用材料外的信息。\n"
"---材料开始---\n" + material + "\n---材料结束---\n问题：" + query
)
def _chunk_retrieve(
tokenizer: Any,
corpus: torch.Tensor,
query: str,
*,
chunk_tokens: int,
top_k: int,
) -> tuple[torch.Tensor, float, int]:
query_ids = set(tokenizer(query, add_special_tokens=False)["input_ids"])
started = time.perf_counter()
scored: list[tuple[int, int, torch.Tensor]] = []
for start in range(0, int(corpus.numel()), chunk_tokens):
chunk = corpus[start : start + chunk_tokens]
overlap = len(query_ids.intersection(set(chunk.tolist())))
scored.append((overlap, -start, chunk))
scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
selected = [item[2] for item in scored[:top_k]]
return torch.cat(selected, dim=0), (time.perf_counter() - started) * 1000.0, len(scored)
@torch.inference_mode()
def _add_memory_chunks(
model: Any,
tokenizer: Any,
corpus: torch.Tensor,
previous_tokens: int,
target_tokens: int,
*,
chunk_tokens: int,
device: torch.device,
needle_ids: torch.Tensor,
) -> dict[str, Any]:
"""Add context chunks without putting the long corpus in the decode KV."""
if model.memory_os_v2 is None:
raise RuntimeError("Natural Memory v2 is required for this benchmark")
# Filler records use deterministic CPU keys. The target record gets an
# actual frozen-Qwen key, which keeps the retrieval measurement honest for
# the one fact being scored without encoding millions of filler tokens.
records: list[dict[str, Any]] = []
body_tokens = max(0, int(corpus.numel()) - int(needle_ids.numel()))
target_start = min(
(body_tokens // 2) // max(1, chunk_tokens) * max(1, chunk_tokens),
65536,
)
for start in range(previous_tokens, target_tokens, chunk_tokens):
end = min(start + chunk_tokens, target_tokens)
ids = corpus[start:end]
if ids.numel() == 0:
continue
if start <= target_start < end:
key = model._encode_model_key(ids.unsqueeze(0).to(device), torch.ones((1, ids.numel()), dtype=torch.long, device=device))[0].cpu()
else:
generator = torch.Generator(device="cpu").manual_seed(1701 + start)
key = torch.randn(model.memory.hidden_size, generator=generator)
key = torch.nn.functional.normalize(key, dim=0)
record_text = (
tokenizer.decode(ids.tolist(), skip_special_tokens=True)
if start <= target_start < end
else f"curve_chunk:{start}:{end}"
)
records.append(
{
"text": record_text,
"key": key,
"summary": key,
"semantic_key": key,
"memory_type": "context_chunk",
"importance": 0.55,
"confidence": 0.85,
"source": "service_curve_benchmark",
"evidence": [f"token_range:{start}:{end}"],
"token_ids": ids,
"token_mask": torch.ones_like(ids, dtype=torch.bool),
"trusted": True,
"force": True,
}
)
started = time.perf_counter()
model.memory_os_v2.write_batch(records)
return {
"added_records": len(records),
"write_ms": (time.perf_counter() - started) * 1000.0,
}
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\service_curves_4b.json")
parser.add_argument("--lengths", default="131072,262144,524288,1048576,4194304")
parser.add_argument("--hot-window", type=int, default=32768)
parser.add_argument("--chunk-tokens", type=int, default=512)
parser.add_argument("--rag-top-k", type=int, default=2)
parser.add_argument(
"--memory-top-k-records",
type=int,
default=2,
help="Natural Memory records injected into hot KV; default matches RAG chunk count",
)
parser.add_argument("--max-new-tokens", type=int, default=32)
parser.add_argument("--gpu-hourly-cost", type=float, default=0.0)
parser.add_argument(
"--gpu-memory-gb",
type=float,
default=8.0,
help="hard CUDA placement cap; 0 disables the cap",
)
parser.add_argument("--no-4bit", action="store_true")
args = parser.parse_args()
_set_cuda_process_cap(args.gpu_memory_gb)
lengths = [int(item.strip()) for item in args.lengths.split(",") if item.strip()]
use_4bit = not args.no_4bit
tokenizer = load_tokenizer(_path(args.base_model))
max_length = max(lengths)
corpus, needle_ids, needle, query = _build_corpus(
tokenizer,
max_length,
chunk_tokens=args.chunk_tokens,
)
report: dict[str, Any] = {
"benchmark": "production_baseline_service_curves_4b",
"model": str(_path(args.base_model)),
"memory_model": str(_path(args.memory_model)),
"quantization": "4bit_nf4" if use_4bit else "none",
"lengths": lengths,
"hot_window": args.hot_window,
"chunk_tokens": args.chunk_tokens,
"rag_top_k": args.rag_top_k,
"memory_top_k_records": args.memory_top_k_records,
"query": query,
"expected": "CURVE-7H2K-9P4M",
"cost_note": "usd_per_million_output_tokens is null unless --gpu-hourly-cost is supplied",
"gpu_hourly_cost": args.gpu_hourly_cost if args.gpu_hourly_cost > 0 else None,
"gpu_memory_cap_gb": args.gpu_memory_gb if args.gpu_memory_gb > 0 else None,
"systems": {},
}
print("loading 4B baseline")
max_memory = _max_memory(args.gpu_memory_gb)
base = load_qwen_base(_path(args.base_model), load_in_4bit=use_4bit, max_memory=max_memory)
base.eval()
base_device = base.get_input_embeddings().weight.device
systems: dict[str, list[dict[str, Any]]] = {
"dense_full_kv": [],
"sliding_window": [],
"matched_hot_window": [],
"chunk_rag": [],
}
for target in lengths:
current = corpus[:target]
row_base = {"context_tokens": target}
if target <= 262144:
try:
prompt = _prompt_content(tokenizer, current, query)
encoded = _chat_from_content(tokenizer, prompt, base_device)
if base_device.type == "cuda":
torch.cuda.reset_peak_memory_stats(base_device)
row = _measure_generation(base, encoded, tokenizer, base_device, dynamic=False, max_new_tokens=args.max_new_tokens)
_quality_fields(row, target_reachable=True, expected="CURVE-7H2K-9P4M")
row_base.update(row)
except Exception as exc:
row_base.update({"status": "error", "error": str(exc)[:500]})
_quality_fields(row_base, target_reachable=True, expected="CURVE-7H2K-9P4M")
else:
row_base.update({"status": "unsupported_by_base_max_position", "quality_pass": False})
_quality_fields(row_base, target_reachable=True, expected="CURVE-7H2K-9P4M")
systems["dense_full_kv"].append(row_base)
window = current[-args.hot_window :]
prompt = _prompt_content(tokenizer, window, query)
encoded = _chat_from_content(tokenizer, prompt, base_device)
if base_device.type == "cuda":
torch.cuda.reset_peak_memory_stats(base_device)
row = _measure_generation(base, encoded, tokenizer, base_device, dynamic=False, max_new_tokens=args.max_new_tokens)
row.update({"context_tokens": target, "visible_tokens": int(window.numel())})
_quality_fields(row, target_reachable=False, expected="CURVE-7H2K-9P4M")
systems["sliding_window"].append(row)
matched_tokens = max(1, int(args.chunk_tokens) * int(args.rag_top_k))
matched_window = current[-matched_tokens:]
prompt = _prompt_content(tokenizer, matched_window, query)
encoded = _chat_from_content(tokenizer, prompt, base_device)
if base_device.type == "cuda":
torch.cuda.reset_peak_memory_stats(base_device)
row = _measure_generation(base, encoded, tokenizer, base_device, dynamic=False, max_new_tokens=args.max_new_tokens)
row.update({"context_tokens": target, "visible_tokens": int(matched_window.numel()), "matched_hot_window_tokens": matched_tokens})
_quality_fields(row, target_reachable=False, expected="CURVE-7H2K-9P4M")
systems["matched_hot_window"].append(row)
retrieved, retrieve_ms, chunk_count = _chunk_retrieve(tokenizer, current, query, chunk_tokens=args.chunk_tokens, top_k=args.rag_top_k)
prompt = _prompt_content(tokenizer, retrieved, query)
encoded = _chat_from_content(tokenizer, prompt, base_device)
if base_device.type == "cuda":
torch.cuda.reset_peak_memory_stats(base_device)
row = _measure_generation(base, encoded, tokenizer, base_device, dynamic=False, max_new_tokens=args.max_new_tokens)
row.update({"context_tokens": target, "retrieved_tokens": int(retrieved.numel()), "chunk_count": chunk_count, "retriever_ms": retrieve_ms})
_quality_fields(row, target_reachable=True, expected="CURVE-7H2K-9P4M")
systems["chunk_rag"].append(row)
print(f"baseline context={target}")
_release(base)
base = None
print("loading Natural Memory 4B")
memory = load_qwen_dynamic(
_path(args.memory_model), load_in_4bit=use_4bit, max_memory=max_memory
)
memory.eval()
memory.memory_config.memory_top_k_records = max(1, int(args.memory_top_k_records))
memory.memory_os_v2.bank.top_k_records = memory.memory_config.memory_top_k_records
memory_device = memory._find_layer_device()
memory.reset_memory(batch_size=1, device=memory_device)
previous = 0
systems["natural_memory"] = []
query_ids, query_mask = _query_tensors(tokenizer, query, memory_device)
for target in lengths:
write_info = _add_memory_chunks(
memory,
tokenizer,
corpus,
previous,
target,
chunk_tokens=args.chunk_tokens,
device=memory_device,
needle_ids=needle_ids,
)
previous = target
encoded = _chat_from_content(tokenizer, query, memory_device)
if memory_device.type == "cuda":
torch.cuda.reset_peak_memory_stats(memory_device)
row = _measure_generation(
memory,
encoded,
tokenizer,
memory_device,
dynamic=True,
query_ids=query_ids,
query_mask=query_mask,
query_text=query,
max_new_tokens=args.max_new_tokens,
)
decision = memory.runtime.v2_last_decisions[-1] if memory.runtime.v2_last_decisions else {}
row.update(
{
"context_tokens": target,
"stored_records": memory.memory_v2_stats().get("records", 0),
"coarse_candidates": decision.get("coarse_candidates", memory.memory_v2_stats().get("last_coarse_candidates", 0)),
"page_ids": decision.get("page_ids", []),
"record_ids": decision.get("record_ids", []),
"stop_reason": decision.get("stop_reason", ""),
"retrieved_records": len(decision.get("record_ids", [])),
"memory_write_ms": write_info["write_ms"],
}
)
_quality_fields(row, target_reachable=True, expected="CURVE-7H2K-9P4M")
systems["natural_memory"].append(row)
print(f"natural_memory context={target}")
_release(memory)
memory = None
if args.gpu_hourly_cost > 0.0:
for rows in systems.values():
for row in rows:
if row.get("status") != "ok":
continue
row["usd_per_million_output_tokens"] = (
1_000_000.0 / max(1e-9, row.get("decode_tok_s", 0.0))
* args.gpu_hourly_cost
/ 3600.0
)
report["systems"] = systems
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}")
if __name__ == "__main__":
main()