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
2026-09-19 11:11:31 +08:00
commit 643e22ecb9
484 changed files with 306821 additions and 0 deletions
+30
View File
@@ -0,0 +1,30 @@
"""Fast environment and gradient smoke test."""
from __future__ import annotations
import torch
from .model import DynamicMemoryConfig, DynamicMemoryLM, count_parameters
from .tasks import sample_associative_batch
def main() -> None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = DynamicMemoryConfig(vocab_size=64, max_seq_len=16, d_model=64, n_layers=2, n_heads=4, memory_slots=4)
model = DynamicMemoryLM(config).to(device)
batch = sample_associative_batch(batch_size=8, vocab_size=config.vocab_size, device=device)
memory = model(batch.learn_chunks[0]).memory
output = model(batch.query_input, memory=memory, update_memory=False, labels=batch.query_labels)
if output.loss is None or not torch.isfinite(output.loss):
raise RuntimeError("non-finite loss")
output.loss.backward()
gradients = [p.grad for p in model.parameters() if p.grad is not None]
if not gradients:
raise RuntimeError("no gradients produced")
print(f"smoke_ok device={device} parameters={count_parameters(model):,} loss={output.loss.detach().item():.4f}")
if device.type == "cuda":
print(f"gpu={torch.cuda.get_device_name(0)} memory_allocated_mb={torch.cuda.memory_allocated() / 1024**2:.1f}")
if __name__ == "__main__":
main()