configuration_livemem.py
2.6 KB · 62 lines · python Raw
1 """Config for the memory-augmented Qwen3 (主路 Qwen3Attention ‖ 边路 GDN2)."""
2 from __future__ import annotations
3
4 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
5
6
7 class LiveMemConfig(Qwen3Config):
8 """Qwen3 + per-attention-head GDN2 memory side-branch.
9
10 All base Qwen3 fields are inherited unchanged. The `mem_*` fields configure
11 the side branch and the two memory mechanisms (Design X / Design Y).
12 """
13
14 model_type = "livemem"
15
16 def __init__(
17 self,
18 # which mechanism: "X" = 连续扫描 (continuous scan, write_mask=None);
19 # "Y" = 门控读写解耦 (freeze gates on read tokens).
20 memory_design: str = "Y",
21 # layers that get a memory branch; None = all layers.
22 mem_layers: list[int] | None = None,
23 # GDN2 side-branch hyper-params. The legacy/default geometry is a
24 # full-MHA copy of Qwen3 (8 KV heads repeated 4x -> 32 heads), per plan
25 # 02 §1.3. Newer experiments can explicitly set mem_num_heads=8,
26 # mem_num_v_heads=8, mem_expand_v=4 to keep similar state capacity with
27 # fewer Q/K heads.
28 mem_head_dim: int | None = None,
29 mem_num_heads: int | None = None,
30 mem_num_v_heads: int | None = None,
31 mem_expand_v: float = 1.0,
32 mem_conv_size: int = 4,
33 mem_conv_bias: bool = False,
34 mem_norm_eps: float | None = None,
35 # zero-init the side o_proj so training starts ≈ original Qwen3 (plan T1).
36 mem_o_proj_zero_init: bool = True,
37 **kwargs,
38 ) -> None:
39 super().__init__(**kwargs)
40 if memory_design not in ("X", "Y"):
41 raise ValueError(f"memory_design must be 'X' or 'Y', got {memory_design!r}")
42 self.memory_design = memory_design
43 self.mem_layers = mem_layers
44 self.mem_head_dim = mem_head_dim if mem_head_dim is not None else self.head_dim
45 self.mem_num_heads = (
46 mem_num_heads if mem_num_heads is not None else self.num_attention_heads
47 )
48 self.mem_num_v_heads = (
49 mem_num_v_heads if mem_num_v_heads is not None else self.mem_num_heads
50 )
51 self.mem_expand_v = mem_expand_v
52 self.mem_conv_size = mem_conv_size
53 self.mem_conv_bias = mem_conv_bias
54 self.mem_norm_eps = mem_norm_eps if mem_norm_eps is not None else self.rms_norm_eps
55 self.mem_o_proj_zero_init = mem_o_proj_zero_init
56
57 @property
58 def memory_layer_indices(self) -> list[int]:
59 if self.mem_layers is None:
60 return list(range(self.num_hidden_layers))
61 return list(self.mem_layers)
62