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，读写关闭时与原生模型逐位相同
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""Synthetic streaming tasks for testing dynamic learning behavior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
@dataclass
|
|
class AssociativeBatch:
|
|
learn_chunks: list[Tensor]
|
|
query_input: Tensor
|
|
query_labels: Tensor
|
|
expected: Tensor
|
|
|
|
|
|
def _rand_tokens(batch_size: int, low: int, high: int, device: torch.device) -> Tensor:
|
|
return torch.randint(low, high, (batch_size,), device=device)
|
|
|
|
|
|
def sample_associative_batch(
|
|
*,
|
|
batch_size: int,
|
|
vocab_size: int,
|
|
device: torch.device,
|
|
overwrite: bool = False,
|
|
) -> AssociativeBatch:
|
|
"""Generate key-value observations followed by a separated query.
|
|
|
|
The query cannot see the learning chunks through attention. It can only
|
|
answer by using the returned dynamic memory state.
|
|
"""
|
|
|
|
key_low, key_high = 4, vocab_size // 2
|
|
value_low, value_high = vocab_size // 2, vocab_size
|
|
keys = _rand_tokens(batch_size, key_low, key_high, device)
|
|
value = _rand_tokens(batch_size, value_low, value_high, device)
|
|
learn_chunks = [torch.stack((keys, value), dim=1)]
|
|
|
|
expected = value
|
|
if overwrite:
|
|
replacement = _rand_tokens(batch_size, value_low, value_high, device)
|
|
learn_chunks.append(torch.stack((keys, replacement), dim=1))
|
|
expected = replacement
|
|
|
|
query_input = torch.stack((keys, expected), dim=1)
|
|
query_labels = torch.full_like(query_input, -100)
|
|
query_labels[:, 1] = expected
|
|
return AssociativeBatch(
|
|
learn_chunks=learn_chunks,
|
|
query_input=query_input,
|
|
query_labels=query_labels,
|
|
expected=expected,
|
|
)
|