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，读写关闭时与原生模型逐位相同
2026-09-19 11:11:31 +08:00
commit 643e22ecb9
484 changed files with 306821 additions and 0 deletions
@@ -0,0 +1,219 @@
+"""A small causal LM with an explicit writable memory state.
+
+This is intentionally independent from the Qwen checkpoint in the parent
+directory. It is a research reference implementation for architecture work.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+@dataclass
+class DynamicMemoryConfig:
+    vocab_size: int = 128
+    max_seq_len: int = 64
+    d_model: int = 128
+    n_layers: int = 4
+    n_heads: int = 4
+    mlp_ratio: int = 4
+    memory_slots: int = 8
+    dropout: float = 0.0
+
+
+@dataclass
+class DynamicMemoryOutput:
+    logits: Tensor
+    memory: Tensor
+    loss: Optional[Tensor] = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        return x * torch.rsqrt(variance + self.eps) * self.weight
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: DynamicMemoryConfig) -> None:
+        super().__init__()
+        if config.d_model % config.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
+        self.out = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.dropout = config.dropout
+
+    def forward(self, x: Tensor) -> Tensor:
+        batch, seq_len, dim = x.shape
+        qkv = self.qkv(x).view(batch, seq_len, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=True,
+        )
+        y = y.transpose(1, 2).contiguous().view(batch, seq_len, dim)
+        return self.out(y)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: DynamicMemoryConfig) -> None:
+        super().__init__()
+        hidden = config.d_model * config.mlp_ratio
+        self.up = nn.Linear(config.d_model, hidden, bias=False)
+        self.down = nn.Linear(hidden, config.d_model, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.down(F.silu(self.up(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: DynamicMemoryConfig) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(config.d_model)
+        self.attn = CausalSelfAttention(config)
+        self.norm2 = RMSNorm(config.d_model)
+        self.mlp = MLP(config)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class DynamicMemory(nn.Module):
+    """A differentiable key-value memory with explicit read/write behavior.
+
+    The memory tensor is returned to the caller and is not a model parameter.
+    It can therefore change during inference without changing the backbone.
+    """
+
+    def __init__(self, config: DynamicMemoryConfig) -> None:
+        super().__init__()
+        self.slots = config.memory_slots
+        self.dim = config.d_model
+        self.read_q = nn.Linear(self.dim, self.dim, bias=False)
+        self.read_k = nn.Linear(self.dim, self.dim, bias=False)
+        self.read_v = nn.Linear(self.dim, self.dim, bias=False)
+        self.read_out = nn.Linear(self.dim, self.dim, bias=False)
+        self.read_gate = nn.Linear(self.dim, 1)
+
+        self.slot_keys = nn.Parameter(torch.randn(self.slots, self.dim) / self.dim**0.5)
+        self.write_value = nn.Linear(self.dim, self.slots * self.dim, bias=False)
+        self.write_gate = nn.Linear(self.dim, self.slots)
+
+    def initial_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> Tensor:
+        return torch.zeros(batch_size, self.slots, self.dim, device=device, dtype=dtype)
+
+    def read(self, x: Tensor, memory: Tensor) -> Tensor:
+        q = self.read_q(x)
+        k = self.read_k(memory)
+        v = self.read_v(memory)
+        scores = torch.matmul(q, k.transpose(-1, -2)) / self.dim**0.5
+        retrieved = torch.matmul(scores.softmax(dim=-1), v)
+        retrieved = self.read_out(retrieved)
+        gate = torch.sigmoid(self.read_gate(x))
+        return gate * retrieved
+
+    def update(self, x: Tensor, memory: Tensor) -> Tensor:
+        # The final token is used as a compact summary of the newly observed
+        # chunk. This makes chunk boundaries explicit and keeps the experiment
+        # cheap enough to run repeatedly on a single consumer GPU.
+        summary = x[:, -1]
+        proposal = self.write_value(summary).view(-1, self.slots, self.dim)
+        address = (summary @ self.slot_keys.t()).softmax(dim=-1)
+        strength = torch.sigmoid(self.write_gate(summary)) * address
+        strength = strength.unsqueeze(-1)
+        return memory + strength * (proposal - memory)
+
+
+class DynamicMemoryLM(nn.Module):
+    """Decoder-only Transformer with a persistent, caller-owned memory state."""
+
+    def __init__(self, config: DynamicMemoryConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos_emb = nn.Embedding(config.max_seq_len, config.d_model)
+        self.memory = DynamicMemory(config)
+        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
+        self.norm = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _ensure_memory(self, memory: Optional[Tensor], batch_size: int, device: torch.device, dtype: torch.dtype) -> Tensor:
+        if memory is None:
+            return self.memory.initial_state(batch_size, device=device, dtype=dtype)
+        if memory.ndim != 3 or memory.shape[0] != batch_size:
+            raise ValueError("memory must have shape [batch, memory_slots, d_model]")
+        return memory
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        *,
+        memory: Optional[Tensor] = None,
+        update_memory: bool = True,
+        labels: Optional[Tensor] = None,
+    ) -> DynamicMemoryOutput:
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, seq]")
+        batch_size, seq_len = input_ids.shape
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f"sequence length {seq_len} exceeds max_seq_len={self.config.max_seq_len}")
+
+        x = self.token_emb(input_ids)
+        positions = torch.arange(seq_len, device=input_ids.device)
+        x = x + self.pos_emb(positions)[None, :, :]
+        memory = self._ensure_memory(memory, batch_size, input_ids.device, x.dtype)
+
+        # Read uses the state from before this chunk. The write happens only
+        # after logits are computed, preventing target-token leakage.
+        x = x + self.memory.read(x, memory)
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        new_memory = self.memory.update(x, memory) if update_memory else memory
+        loss = None
+        if labels is not None:
+            if labels.shape != input_ids.shape:
+                raise ValueError("labels must have the same shape as input_ids")
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.size(-1)),
+                labels[:, 1:].reshape(-1),
+                ignore_index=-100,
+            )
+        return DynamicMemoryOutput(logits=logits, memory=new_memory, loss=loss)
+
+
+def count_parameters(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())