Add Natural Memory architecture and tooling

2026-09-05 08:53:41 +08:00
parent 0acf8b06ee
commit 516351f0b5
56 changed files with 18319 additions and 0 deletions
@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import unittest
+
+import torch
+from torch import nn
+
+from dynamic_memory_lab.qwen_integration import (
+    AutomaticMemoryPolicy,
+    MemoryLayerAdapter,
+    NaturalLanguageRetriever,
+    NativeQwenDynamicMemory,
+    QwenMemoryConfig,
+    QwenDynamicMemory,
+    _MemoryRuntime,
+    looks_like_question,
+    split_memory_candidates,
+)
+
+
+class _FakeAttention(nn.Module):
+    def __init__(self, *, fail_if_called: bool = False) -> None:
+        super().__init__()
+        self.fail_if_called = fail_if_called
+        self.called = False
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        self.called = True
+        if self.fail_if_called:
+            raise AssertionError("original token mixer was called in replace mode")
+        return hidden_states * 2.0, None
+
+
+class _FakeQwenLayer(nn.Module):
+    layer_type = "full_attention"
+
+    def __init__(self, *, fail_if_called: bool = False) -> None:
+        super().__init__()
+        self.input_layernorm = nn.Identity()
+        self.post_attention_layernorm = nn.Identity()
+        self.self_attn = _FakeAttention(fail_if_called=fail_if_called)
+        self.mlp = nn.Identity()
+
+    def forward(self, hidden_states, position_embeddings=None, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
+        return hidden_states + self.self_attn(hidden_states)[0]
+
+
+class QwenSurgeryTest(unittest.TestCase):
+    def _runtime(self):
+        memory = QwenDynamicMemory(
+            hidden_size=8,
+            config=QwenMemoryConfig(memory_slots=2, memory_dim=4),
+        )
+        runtime = _MemoryRuntime(memory)
+        runtime.state = memory.initial_state(1, device=torch.device("cpu"))
+        return runtime
+
+    def test_blend_exposes_trainable_mixer_weight(self) -> None:
+        runtime = self._runtime()
+        adapter = MemoryLayerAdapter(
+            _FakeQwenLayer(),
+            runtime,
+            read=True,
+            write=False,
+            mode="blend",
+            blend_init=0.5,
+        )
+        output = adapter(torch.ones(1, 3, 8))
+        output.sum().backward()
+        self.assertEqual(tuple(output.shape), (1, 3, 8))
+        self.assertIsNotNone(adapter.blend_logit.grad)
+
+    def test_replace_skips_original_token_mixer(self) -> None:
+        runtime = self._runtime()
+        layer = _FakeQwenLayer(fail_if_called=True)
+        adapter = MemoryLayerAdapter(layer, runtime, read=True, write=False, mode="replace")
+        output = adapter(torch.ones(1, 3, 8))
+        self.assertEqual(tuple(output.shape), (1, 3, 8))
+        self.assertFalse(layer.self_attn.called)
+
+    def test_raw_token_write_uses_output_projection_row(self) -> None:
+        memory = QwenDynamicMemory(
+            hidden_size=8,
+            config=QwenMemoryConfig(
+                memory_slots=2,
+                memory_dim=4,
+                write_token_offset=2,
+                raw_token_write=True,
+                broadcast_write=True,
+            ),
+        )
+        runtime = _MemoryRuntime(memory)
+        runtime.state = memory.initial_state(1, device=torch.device("cpu"))
+        runtime.read_enabled = False
+        runtime.update_enabled = True
+        runtime.input_ids = torch.tensor([[5, 6, 7, 8]])
+        runtime.attention_mask = torch.ones_like(runtime.input_ids)
+        runtime.output_embeddings = nn.Linear(8, 16, bias=False)
+        adapter = MemoryLayerAdapter(_FakeQwenLayer(), runtime, read=True, write=True, mode="residual")
+
+        adapter(torch.ones(1, 4, 8))
+
+        expected = runtime.output_embeddings.weight[7]
+        self.assertIsNotNone(runtime.raw_memory)
+        self.assertTrue(torch.allclose(runtime.raw_memory[0], expected))
+
+    def test_native_controller_exposes_write_forget_and_value_state(self) -> None:
+        memory = NativeQwenDynamicMemory(
+            hidden_size=8,
+            config=QwenMemoryConfig(memory_slots=2, memory_dim=4),
+        )
+        hidden = torch.randn(1, 3, 8)
+        state = memory.initial_state(1, device=torch.device("cpu"))
+        updated = memory.update(hidden, state, attention_mask=torch.ones(1, 5, dtype=torch.long))
+        self.assertEqual(tuple(updated.shape), (1, 2, 4))
+        self.assertEqual(tuple(memory.last_write_probability.shape), (1, 1))
+        self.assertEqual(tuple(memory.last_forget_probability.shape), (1, 2))
+        self.assertEqual(tuple(memory.last_write_summary.shape), (1, 8))
+        self.assertEqual(tuple(memory.last_write_representation.shape), (1, 8))
+
+    def test_natural_language_retriever_scores_single_and_multiple_slots(self) -> None:
+        retriever = NaturalLanguageRetriever(hidden_size=8, projection_size=4)
+        query = torch.randn(2, 8)
+        one_key = torch.randn(2, 8)
+        many_keys = torch.randn(2, 3, 8)
+        self.assertEqual(tuple(retriever(query, one_key).shape), (2,))
+        self.assertEqual(tuple(retriever(query, many_keys).shape), (2, 3))
+
+    def test_automatic_memory_policy_and_candidate_segmentation(self) -> None:
+        policy = AutomaticMemoryPolicy(hidden_size=8)
+        output = policy(torch.randn(3, 8))
+        self.assertEqual(tuple(output.shape), (3,))
+        self.assertEqual(
+            split_memory_candidates("我叫林浩，我正在开发星火项目。"),
+            ["我叫林浩，我正在开发星火项目。"],
+        )
+        self.assertTrue(looks_like_question("如果我选择 GPU，会发生什么？"))
+        self.assertFalse(looks_like_question("我住在上海，正在开发星火项目。"))
+
+
+if __name__ == "__main__":
+    unittest.main()