108 lines
4.5 KiB
Python
108 lines
4.5 KiB
Python
"""Item 31: synthetic memory tokens — per-layer KV prefix (Nils's design).
|
|
|
|
"The one read mechanism a frozen transformer natively possesses is
|
|
attention over the KV cache." A KVMemoryAdapter maps carried L30 states
|
|
to (k, v) entries appended to every band layer's attention — frozen
|
|
heads consume the state through their existing machinery. Keys live in
|
|
post-RoPE space (prefix-tuning convention: learned keys place themselves
|
|
where frozen queries already point).
|
|
|
|
Activation is global-flag-scoped like lora_band.LoopLoRA: a custom
|
|
attention interface "kvmem" wraps the stock eager path; when _MEM is
|
|
armed and the module is a band layer, memory kv columns are appended
|
|
(mask columns visible to every query).
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
|
|
from transformers.models.gemma4.modeling_gemma4 import (
|
|
eager_attention_forward)
|
|
|
|
_MEM = {"on": False, "kv": None, "band": (14, 30), "base": None}
|
|
|
|
|
|
def arm_memory(kv):
|
|
"""kv: dict layer_idx -> (k, v), each (B, kv_heads, M, head_dim)."""
|
|
_MEM["kv"] = kv
|
|
_MEM["on"] = kv is not None
|
|
|
|
|
|
def kvmem_attention(module, query, key, value, attention_mask, **kwargs):
|
|
if (_MEM["on"] and _MEM["kv"] is not None
|
|
and getattr(module, "layer_idx", -1) in _MEM["kv"]):
|
|
mk, mv, gate = _MEM["kv"][module.layer_idx]
|
|
m = mk.shape[2]
|
|
key = torch.cat([key, mk.to(dtype=key.dtype, device=key.device)], 2)
|
|
value = torch.cat([value, mv.to(dtype=value.dtype,
|
|
device=value.device)], 2)
|
|
if attention_mask is not None:
|
|
pad = torch.zeros(*attention_mask.shape[:-1], m,
|
|
dtype=attention_mask.dtype,
|
|
device=attention_mask.device) + gate.to(
|
|
attention_mask.dtype)
|
|
attention_mask = torch.cat([attention_mask, pad], -1)
|
|
base = _MEM["base"] or eager_attention_forward
|
|
return base(module, query, key, value, attention_mask, **kwargs)
|
|
|
|
|
|
def install(model):
|
|
"""Wrap the registered implementation IN PLACE — the config name is
|
|
untouched, so mask construction and every other branch stays
|
|
bit-identical; only the (q,k,v,mask) seam gains the memory append."""
|
|
if _MEM.get("installed"):
|
|
return
|
|
impl = model.config._attn_implementation
|
|
_MEM["base"] = ALL_ATTENTION_FUNCTIONS.get_interface(
|
|
impl, eager_attention_forward)
|
|
try:
|
|
ALL_ATTENTION_FUNCTIONS[impl] = kvmem_attention
|
|
except TypeError:
|
|
ALL_ATTENTION_FUNCTIONS.register(impl, kvmem_attention)
|
|
_MEM["installed"] = True
|
|
|
|
|
|
class KVMemoryAdapter(nn.Module):
|
|
"""Carried states (B, M, d) -> per-band-layer post-RoPE (k, v).
|
|
Per-layer geometry (head_dim, kv heads) read off the live modules —
|
|
gemma-4 mixes sliding (256) and global (512) head dims in the band."""
|
|
|
|
def __init__(self, model, band=(14, 30), code=512):
|
|
super().__init__()
|
|
self.band = band
|
|
tm = model.model.language_model
|
|
d = model.config.get_text_config().hidden_size
|
|
self.trunk = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, code),
|
|
nn.GELU())
|
|
self.geom = {}
|
|
self.k_heads = nn.ModuleDict()
|
|
self.v_heads = nn.ModuleDict()
|
|
tc = model.config.get_text_config()
|
|
for l in range(band[0], band[1] + 1):
|
|
att = tm.layers[l].self_attn
|
|
hd = att.head_dim
|
|
kvh = (tc.num_global_key_value_heads
|
|
if getattr(att, "use_alternative_attention", False)
|
|
else tc.num_key_value_heads)
|
|
self.geom[l] = (kvh, hd)
|
|
self.k_heads[str(l)] = nn.Linear(code, kvh * hd)
|
|
self.v_heads[str(l)] = nn.Linear(code, kvh * hd)
|
|
nn.init.zeros_(self.v_heads[str(l)].weight)
|
|
nn.init.zeros_(self.v_heads[str(l)].bias)
|
|
# learnable per-layer attention gate on the memory columns:
|
|
# init -10 -> memory is invisible until training opens it
|
|
self.gate = nn.ParameterDict(
|
|
{str(l): nn.Parameter(torch.tensor(-10.0))
|
|
for l in range(band[0], band[1] + 1)})
|
|
|
|
def forward(self, states):
|
|
B, M, _ = states.shape
|
|
c = self.trunk(states.float())
|
|
kv = {}
|
|
for l in range(self.band[0], self.band[1] + 1):
|
|
kvh, hd = self.geom[l]
|
|
k = self.k_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
|
|
v = self.v_heads[str(l)](c).view(B, M, kvh, hd).transpose(1, 2)
|
|
kv[l] = (k, v, self.gate[str(l)])
|
|
return kv
|