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