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，读写关闭时与原生模型逐位相同
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
"""Debug the raw token pointer memory path on one benchmark record."""
|
|
|
|
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:
|
|
config = QwenMemoryConfig(
|
|
mode="blend",
|
|
blend_init=0.1,
|
|
write_token_offset=4,
|
|
broadcast_write=True,
|
|
raw_token_write=True,
|
|
raw_logit_scale=30.0,
|
|
)
|
|
model = load_qwen_dynamic(".", memory_config=config, load_in_4bit=True)
|
|
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))
|
|
print("memory_ids", memory_input.tolist())
|
|
print("memory_mask", memory_mask.tolist())
|
|
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
|
print("prompt_ids", prompt["input_ids"].tolist())
|
|
with torch.inference_mode():
|
|
output = 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,
|
|
)
|
|
print("raw_memory_shape", tuple(model.runtime.raw_memory.shape))
|
|
print("raw_memory_norm", float(model.runtime.raw_memory.float().norm()))
|
|
generated = model.generate(
|
|
**prompt,
|
|
max_new_tokens=1,
|
|
do_sample=False,
|
|
update_memory=False,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
print("generated", generated.tolist())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|