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:
@@ -0,0 +1,66 @@
|
||||
"""Check whether teacher-forcing and cached greedy generation agree on token 1."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import torch
|
||||
|
||||
from .benchmark_qwen import _generation_prompt
|
||||
from .qwen_integration import QwenMemoryConfig, load_qwen_dynamic, load_tokenizer
|
||||
from .train_qwen_memory import encode_messages, pad_batch
|
||||
|
||||
|
||||
def main() -> None:
|
||||
model = load_qwen_dynamic(
|
||||
".",
|
||||
memory_config=QwenMemoryConfig(mode="blend", blend_init=0.1),
|
||||
load_in_4bit=True,
|
||||
)
|
||||
model.load_memory_adapter("V2_dpskw/qwen_memory_adapter_full")
|
||||
model.eval()
|
||||
tokenizer = load_tokenizer(".")
|
||||
record = json.loads(
|
||||
next(open("V2_dpskw/data/benchmark_eval.jsonl", encoding="utf-8"))
|
||||
)
|
||||
device = model._find_layer_device()
|
||||
memory = encode_messages(tokenizer, record["memory"], 128)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], int(tokenizer.pad_token_id))
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
|
||||
with torch.inference_mode():
|
||||
memory_state = model(
|
||||
input_ids=memory_input.to(device),
|
||||
attention_mask=memory_mask.to(device),
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
).memory
|
||||
teacher_forcing = model(
|
||||
input_ids=prompt["input_ids"],
|
||||
attention_mask=prompt["attention_mask"],
|
||||
memory_state=memory_state,
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
generated = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=1,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
|
||||
teacher_id = int(teacher_forcing.logits[0, -1].argmax())
|
||||
generated_id = int(generated[0, -1])
|
||||
print(json.dumps({"expected": record["answer"]}, ensure_ascii=True))
|
||||
print(json.dumps({"teacher_id": teacher_id, "teacher_text": tokenizer.decode([teacher_id])}, ensure_ascii=True))
|
||||
print(json.dumps({"generated_id": generated_id, "generated_text": tokenizer.decode([generated_id])}, ensure_ascii=True))
|
||||
print(f"same={teacher_id == generated_id}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user