modeling_livemem_gdn2.py
7.0 KB · 172 lines · python Raw
1 """GDN2 memory side-branch with an optional per-token write gate.
2
3 This subclasses fla's `GatedDeltaNet2` and overrides `forward` to apply a
4 `write_mask` that freezes the recurrence on read tokens. The GDN-2 update is
5
6 S_t = (I - k_t (b_t * k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t * v_t)^T
7
8 so setting g_t = b_t = w_t = 0 gives S_t = S_{t-1} (read-only); the output
9 o_t = q_t · S_t is still produced (the token reads the memory). Verified
10 against `fla.ops.gdn2.naive.naive_recurrent_gdn2`.
11
12 - Design Y: write_mask = is_evicted (open the gate only on the compress/evict
13 region; live + Q + A are frozen read-only).
14 - Design X: write_mask = None -> standard continuous scan over all tokens.
15 """
16 from __future__ import annotations
17
18 import torch
19 import torch.nn.functional as F
20 from einops import rearrange, repeat
21
22 from fla.layers.gdn2 import GatedDeltaNet2
23 from fla.layers.utils import (
24 get_layer_cache,
25 get_unpad_data,
26 index_first_axis,
27 pad_input,
28 update_layer_cache,
29 )
30 from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
31
32
33 class LiveMemGatedDeltaNet2(GatedDeltaNet2):
34 def forward(
35 self,
36 hidden_states: torch.Tensor,
37 write_mask: torch.Tensor | None = None,
38 attention_mask: torch.Tensor | None = None,
39 past_key_values=None,
40 use_cache: bool | None = False,
41 output_attentions: bool | None = False,
42 **kwargs,
43 ):
44 """Mirrors fla.GatedDeltaNet2.forward, adding `write_mask` support.
45
46 `write_mask`: float/bool [batch, seq_len]; 0 freezes the recurrence
47 (g=b=w=0 -> S_t = S_{t-1}). Padding (`attention_mask`) + `write_mask`
48 together are not supported (the unpad reorder would desync the mask);
49 the training path always passes `attention_mask=None`.
50 """
51 if attention_mask is not None:
52 assert len(attention_mask.shape) == 2, (
53 "Expected attention_mask as a [batch_size, seq_len] 0/1 padding mask."
54 )
55 assert write_mask is None, (
56 "write_mask is incompatible with a 2D padding mask (unpad reorder)."
57 )
58
59 cu_seqlens = kwargs.get("cu_seqlens")
60 batch_size, q_len, _ = hidden_states.shape
61 if cu_seqlens is not None and cu_seqlens.ndim == 2:
62 if use_cache:
63 raise ValueError("batched cu_seqlens training path does not support cache")
64 outs = []
65 for b in range(batch_size):
66 cu = cu_seqlens[b]
67 cu = cu[cu >= 0].contiguous()
68 wm = write_mask[b:b + 1] if write_mask is not None else None
69 o, _, _ = self.forward(
70 hidden_states[b:b + 1],
71 write_mask=wm,
72 attention_mask=None,
73 past_key_values=None,
74 use_cache=False,
75 output_attentions=output_attentions,
76 cu_seqlens=cu,
77 )
78 outs.append(o)
79 return torch.cat(outs, dim=0), None, past_key_values
80
81 mode = "fused_recurrent" if (q_len <= 64 and not self.training) else self.mode
82 if self.training:
83 assert mode == "chunk", "Only chunk mode is supported in training."
84
85 last_state = get_layer_cache(self, past_key_values)
86 indices = None
87 if cu_seqlens is None and attention_mask is not None:
88 indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
89 hidden_states = index_first_axis(
90 rearrange(hidden_states, "b s ... -> (b s) ..."), indices
91 ).unsqueeze(0)
92
93 if self.use_short_conv:
94 conv_state_q, conv_state_k, conv_state_v = None, None, None
95 if last_state is not None:
96 conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"]
97 q, conv_state_q = self.q_conv1d(
98 x=self.q_proj(hidden_states), cache=conv_state_q,
99 output_final_state=use_cache, cu_seqlens=cu_seqlens,
100 )
101 k, conv_state_k = self.k_conv1d(
102 x=self.k_proj(hidden_states), cache=conv_state_k,
103 output_final_state=use_cache, cu_seqlens=cu_seqlens,
104 )
105 v, conv_state_v = self.v_conv1d(
106 x=self.v_proj(hidden_states), cache=conv_state_v,
107 output_final_state=use_cache, cu_seqlens=cu_seqlens,
108 )
109 else:
110 q = F.silu(self.q_proj(hidden_states))
111 k = F.silu(self.k_proj(hidden_states))
112 v = F.silu(self.v_proj(hidden_states))
113
114 g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
115 b = self.b_proj(hidden_states).sigmoid()
116 w = self.w_proj(hidden_states).sigmoid()
117
118 q, k, g = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) for x in (q, k, g))
119 v = rearrange(v, "... (h d) -> ... h d", d=self.head_v_dim)
120 b = rearrange(b, "... (h d) -> ... h d", d=self.head_k_dim)
121 w = rearrange(w, "... (h d) -> ... h d", d=self.head_v_dim)
122 g = -self.A_log.float().exp().unsqueeze(-1) * g
123
124 # --- write gate (the only addition over fla) -----------------------
125 # Zeroing g/b/w on read tokens freezes the state at S_{t-1}.
126 if write_mask is not None:
127 wm = write_mask.to(g.dtype).view(write_mask.shape[0], write_mask.shape[1], 1, 1)
128 g = g * wm
129 b = b * wm.to(b.dtype)
130 w = w * wm.to(w.dtype)
131 # -------------------------------------------------------------------
132
133 if self.num_v_heads > self.num_heads:
134 q, k, g, b = (
135 repeat(x, "... h d -> ... (h g) d", g=self.num_v_heads // self.num_heads)
136 for x in (q, k, g, b)
137 )
138
139 if self.allow_neg_eigval:
140 b = b * 2.0
141
142 recurrent_state = last_state["recurrent_state"] if last_state is not None else None
143 if mode == "chunk":
144 o, recurrent_state = chunk_gdn2(
145 q=q, k=k, v=v, g=g, b=b, w=w,
146 initial_state=recurrent_state, output_final_state=use_cache,
147 use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens,
148 )
149 elif mode == "fused_recurrent":
150 o, recurrent_state = fused_recurrent_gdn2(
151 q=q, k=k, v=v, g=g, b=b, w=w,
152 initial_state=recurrent_state, output_final_state=use_cache,
153 use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens,
154 )
155 else:
156 raise NotImplementedError(f"Unsupported mode `{mode}`.")
157
158 update_layer_cache(
159 self, past_key_values,
160 recurrent_state=recurrent_state,
161 conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
162 offset=q_len,
163 )
164
165 o = self.o_norm(o, rearrange(self.g_proj(hidden_states), "... (h d) -> ... h d", d=self.head_v_dim))
166 o = rearrange(o, "b t h d -> b t (h d)")
167 o = self.o_proj(o)
168 if attention_mask is not None:
169 o = pad_input(o.squeeze(0), indices, batch_size, q_len)
170
171 return o, None, past_key_values
172