143 lines
5.4 KiB
Python
143 lines
5.4 KiB
Python
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()
|