modeling_livemem.py
21.3 KB · 497 lines · python Raw
1 """Memory-augmented Qwen3: 主路 Qwen3Attention ‖ 边路 GDN2, o = o_main + o_side.
2
3 Design X (continuous scan) and Design Y (gated read/write) share one skeleton
4 and the same attention eviction mask; they differ only in the RNN `write_mask`:
5 - X: write_mask = None (RNN scans every token; state = compression of all)
6 - Y: write_mask = is_evicted (open gate only on the evict/compress region)
7 """
8 from __future__ import annotations
9
10 import torch
11 import torch.nn as nn
12
13 from transformers.cache_utils import Cache, DynamicCache
14 from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
15 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
16 from transformers.models.qwen3.modeling_qwen3 import (
17 Qwen3Attention,
18 Qwen3ForCausalLM,
19 Qwen3Model,
20 Qwen3PreTrainedModel,
21 )
22
23 from .configuration_livemem import LiveMemConfig
24 from .modeling_livemem_gdn2 import LiveMemGatedDeltaNet2
25
26
27 def make_memory_and_mask(
28 is_evicted: torch.Tensor | None = None,
29 segment_ids: torch.Tensor | None = None,
30 seq_ids: torch.Tensor | None = None,
31 chunk_id: torch.Tensor | None = None,
32 evict_step: torch.Tensor | None = None,
33 ):
34 """Build an `and_mask_function` for `create_causal_mask` from up to four
35 constraints, AND-combined with causal by `create_causal_mask`:
36
37 - dynamic eviction (`chunk_id [B,T]`, `evict_step [B,T]` int): the real
38 training path. Each token belongs to a chunk; `evict_step[t]` is the chunk
39 step at which that token's chunk is evicted to the RNN state. keep(q,kv) =
40 `evict_step[kv] > chunk_id[q]` — key kv's chunk must still be live when
41 query q's chunk is processed (Design X: RNN scans all, attention evicts).
42 - static eviction (`is_evicted [B,T]` bool): the simple synthetic variant.
43 keep(q,kv) = (kv live) OR (q evicted).
44 - segments (`segment_ids [B,T]` int): keep(q,kv) = (kv shared, id==0) OR
45 (same segment). PACK block-diagonal QA.
46 - documents (`seq_ids [B,T]` int): keep(q,kv) = same document. cu_seqlens
47 packing isolation.
48
49 Returns None if no constraint is given. Works for flex_attention (BlockMask)
50 and sdpa/eager (vmapped) backends.
51 """
52 preds = []
53 if chunk_id is not None and evict_step is not None:
54 preds.append(lambda b, q, kv: evict_step[b, kv] > chunk_id[b, q])
55 if is_evicted is not None:
56 preds.append(lambda b, q, kv: (~is_evicted[b, kv]) | is_evicted[b, q])
57 if segment_ids is not None:
58 preds.append(lambda b, q, kv: (segment_ids[b, kv] == 0) | (segment_ids[b, kv] == segment_ids[b, q]))
59 if seq_ids is not None:
60 preds.append(lambda b, q, kv: seq_ids[b, kv] == seq_ids[b, q])
61 if not preds:
62 return None
63
64 def and_mask(b, h, q, kv):
65 out = preds[0](b, q, kv)
66 for p in preds[1:]:
67 out = out & p(b, q, kv)
68 return out
69
70 return and_mask
71
72
73 # Backwards-compatible alias.
74 def make_evict_and_mask(is_evicted: torch.Tensor):
75 return make_memory_and_mask(is_evicted)
76
77
78 class LiveMemAttention(nn.Module):
79 """Wraps the original Qwen3Attention (main path) and adds a GDN2 side branch.
80
81 Per-forward memory control (`write_mask`, side-branch cache) is set as
82 attributes by the model loop rather than threaded through kwargs, so the
83 base attention path and HF decorators never see custom kwargs.
84 """
85
86 def __init__(self, base_attn: Qwen3Attention, config: LiveMemConfig) -> None:
87 super().__init__()
88 self.layer_idx = base_attn.layer_idx
89 self.attn = base_attn # main path: untouched Qwen3Attention
90 self.mem = LiveMemGatedDeltaNet2(
91 hidden_size=config.hidden_size,
92 expand_v=config.mem_expand_v,
93 head_dim=config.mem_head_dim,
94 num_heads=config.mem_num_heads,
95 num_v_heads=config.mem_num_v_heads,
96 mode="chunk",
97 use_short_conv=True,
98 conv_size=config.mem_conv_size,
99 conv_bias=config.mem_conv_bias,
100 layer_idx=base_attn.layer_idx,
101 norm_eps=config.mem_norm_eps,
102 )
103 # per-forward control, set by LiveMemModel.forward
104 self._mem_write_mask: torch.Tensor | None = None
105 self._mem_cache = None
106 self._mem_use_cache: bool = False
107 self._mem_cu_seqlens: torch.Tensor | None = None
108 # Training diagnostics. Disabled by default; train/sft/loop.py enables
109 # this on one layer so normal forward/inference pays no reduction cost.
110 self._record_o_stats: bool = False
111 self._last_o_stats: dict[str, torch.Tensor] = {}
112
113 def forward(
114 self,
115 hidden_states: torch.Tensor,
116 position_embeddings,
117 attention_mask: torch.Tensor | None,
118 past_key_values: Cache | None = None,
119 **kwargs,
120 ):
121 o_main, attn_weights = self.attn(
122 hidden_states,
123 position_embeddings=position_embeddings,
124 attention_mask=attention_mask,
125 past_key_values=past_key_values,
126 **kwargs,
127 )
128 o_side, _, _ = self.mem(
129 hidden_states,
130 write_mask=self._mem_write_mask,
131 past_key_values=self._mem_cache,
132 use_cache=self._mem_use_cache,
133 cu_seqlens=self._mem_cu_seqlens,
134 )
135 o_total = o_main + o_side
136 if self._record_o_stats:
137 with torch.no_grad():
138 main_abs = o_main.detach().float().abs().mean()
139 side_abs = o_side.detach().float().abs().mean()
140 total_abs = o_total.detach().float().abs().mean()
141 eps = torch.tensor(1e-12, device=total_abs.device, dtype=total_abs.dtype)
142 self._last_o_stats = {
143 "main_out_abs": main_abs,
144 "side_out_abs": side_abs,
145 "total_out_abs": total_abs,
146 "side_out_ratio": side_abs / torch.maximum(total_abs, eps),
147 "side_main_ratio": side_abs / torch.maximum(main_abs, eps),
148 }
149 return o_total, attn_weights
150
151
152 class LiveMemPreTrainedModel(Qwen3PreTrainedModel):
153 config: LiveMemConfig
154 _no_split_modules = ["Qwen3DecoderLayer"]
155
156
157 class LiveMemModel(LiveMemPreTrainedModel, Qwen3Model):
158 config_class = LiveMemConfig
159
160 def __init__(self, config: LiveMemConfig) -> None:
161 Qwen3Model.__init__(self, config)
162 # Replace self_attn with LiveMemAttention on the selected layers.
163 mem_layers = set(config.memory_layer_indices)
164 for idx in mem_layers:
165 layer = self.layers[idx]
166 layer.self_attn = LiveMemAttention(layer.self_attn, config)
167 self._mem_layer_indices = sorted(mem_layers)
168 self.post_init()
169
170 def _set_mem_control(self, write_mask, mem_cache, mem_use_cache, cu_seqlens=None) -> None:
171 for idx in self._mem_layer_indices:
172 m = self.layers[idx].self_attn
173 m._mem_write_mask = write_mask
174 m._mem_cache = mem_cache
175 m._mem_use_cache = mem_use_cache
176 m._mem_cu_seqlens = cu_seqlens
177
178 def _clear_mem_control(self) -> None:
179 self._set_mem_control(None, None, False, None)
180
181 def forward(
182 self,
183 input_ids: torch.LongTensor | None = None,
184 attention_mask: torch.Tensor | None = None,
185 position_ids: torch.LongTensor | None = None,
186 past_key_values: Cache | None = None,
187 inputs_embeds: torch.FloatTensor | None = None,
188 use_cache: bool | None = None,
189 is_evicted: torch.Tensor | None = None,
190 write_mask: torch.Tensor | None = None,
191 segment_ids: torch.Tensor | None = None,
192 seq_ids: torch.Tensor | None = None,
193 chunk_id: torch.Tensor | None = None,
194 evict_step: torch.Tensor | None = None,
195 cu_seqlens: torch.Tensor | None = None,
196 mem_cache=None,
197 **kwargs,
198 ) -> BaseModelOutputWithPast:
199 if (input_ids is None) ^ (inputs_embeds is not None):
200 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
201
202 if inputs_embeds is None:
203 inputs_embeds = self.embed_tokens(input_ids)
204
205 if use_cache and past_key_values is None:
206 past_key_values = DynamicCache(config=self.config)
207
208 if position_ids is None:
209 past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
210 position_ids = torch.arange(
211 inputs_embeds.shape[1], device=inputs_embeds.device
212 ).unsqueeze(0) + past_seen
213
214 # RNN write gate (per token). An explicit `write_mask` always wins (PACK
215 # freezes QA segments, segment-write uses 双位置); otherwise Design Y
216 # derives it from the eviction layout, and Design X scans continuously.
217 if write_mask is not None:
218 write_mask = write_mask.to(inputs_embeds.dtype)
219 elif is_evicted is not None and self.config.memory_design == "Y":
220 write_mask = is_evicted.to(inputs_embeds.dtype)
221
222 # Build the (eviction / segment / document-aware) causal mask, reused for
223 # all layers. `seq_ids` isolates packed sequences (cu_seqlens path);
224 # `chunk_id`/`evict_step` drive dynamic chunk eviction (real training).
225 if not isinstance(attention_mask, dict):
226 and_mask = make_memory_and_mask(is_evicted, segment_ids, seq_ids, chunk_id, evict_step)
227 mask_kwargs = {
228 "config": self.config,
229 "inputs_embeds": inputs_embeds,
230 "attention_mask": attention_mask,
231 "past_key_values": past_key_values,
232 "position_ids": position_ids,
233 "and_mask_function": and_mask,
234 }
235 causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
236 if self.has_sliding_layers:
237 causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
238 else:
239 causal_mask_mapping = attention_mask
240
241 hidden_states = inputs_embeds
242 position_embeddings = self.rotary_emb(hidden_states, position_ids)
243
244 # Set control on every forward (incl. None when no eviction), so there is
245 # no stale state. We deliberately do NOT clear afterwards: gradient
246 # checkpointing recomputes this forward during backward and must see the
247 # same write_mask. Training is sequential (forward→backward→next forward),
248 # so the values stay valid until the next forward overwrites them.
249 self._set_mem_control(write_mask, mem_cache, bool(use_cache), cu_seqlens)
250 for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
251 hidden_states = decoder_layer(
252 hidden_states,
253 attention_mask=causal_mask_mapping[self.config.layer_types[i]],
254 position_embeddings=position_embeddings,
255 position_ids=position_ids,
256 past_key_values=past_key_values,
257 use_cache=use_cache,
258 **kwargs,
259 )
260
261 hidden_states = self.norm(hidden_states)
262 return BaseModelOutputWithPast(
263 last_hidden_state=hidden_states,
264 past_key_values=past_key_values if use_cache else None,
265 )
266
267
268 class LiveMemForCausalLM(LiveMemPreTrainedModel, Qwen3ForCausalLM):
269 config_class = LiveMemConfig
270 _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
271
272 def __init__(self, config: LiveMemConfig) -> None:
273 # Build directly (don't call Qwen3ForCausalLM.__init__, which would
274 # construct a throwaway base Qwen3Model first).
275 Qwen3PreTrainedModel.__init__(self, config)
276 self.model = LiveMemModel(config)
277 self.vocab_size = config.vocab_size
278 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
279 self.post_init()
280 # Honor zero-init of the side o_proj for *both* construction paths
281 # (post_init randomizes it, so this must run last). For from_qwen3 the
282 # copy-init also zeros it; here it covers from-scratch construction.
283 if config.mem_o_proj_zero_init:
284 for idx in config.memory_layer_indices:
285 nn.init.zeros_(self.model.layers[idx].self_attn.mem.o_proj.weight)
286 def forward(
287 self,
288 input_ids: torch.LongTensor | None = None,
289 attention_mask: torch.Tensor | None = None,
290 position_ids: torch.LongTensor | None = None,
291 past_key_values: Cache | None = None,
292 inputs_embeds: torch.FloatTensor | None = None,
293 labels: torch.LongTensor | None = None,
294 use_cache: bool | None = None,
295 is_evicted: torch.Tensor | None = None,
296 write_mask: torch.Tensor | None = None,
297 segment_ids: torch.Tensor | None = None,
298 seq_ids: torch.Tensor | None = None,
299 chunk_id: torch.Tensor | None = None,
300 evict_step: torch.Tensor | None = None,
301 cu_seqlens: torch.Tensor | None = None,
302 mem_cache=None,
303 logits_to_keep: int | torch.Tensor = 0,
304 **kwargs,
305 ) -> CausalLMOutputWithPast:
306 outputs = self.model(
307 input_ids=input_ids,
308 attention_mask=attention_mask,
309 position_ids=position_ids,
310 past_key_values=past_key_values,
311 inputs_embeds=inputs_embeds,
312 use_cache=use_cache,
313 is_evicted=is_evicted,
314 write_mask=write_mask,
315 segment_ids=segment_ids,
316 seq_ids=seq_ids,
317 chunk_id=chunk_id,
318 evict_step=evict_step,
319 cu_seqlens=cu_seqlens,
320 mem_cache=mem_cache,
321 **kwargs,
322 )
323 hidden_states = outputs.last_hidden_state
324
325 if labels is not None:
326 # Answer-only logits: gather just the supervised positions and run
327 # lm_head on those, so we never materialize [B, L, vocab] (≈40GB at
328 # L=128k). Mathematically identical to full-seq CE with ignore_index
329 # = -100 (mean over answer tokens -> 按 answer token 归一).
330 shift_hidden = hidden_states[:, :-1, :]
331 shift_labels = labels[:, 1:].to(hidden_states.device)
332 sel = shift_labels != -100
333 sel_hidden = shift_hidden[sel] # [n_answer, H] bf16
334 sel_lab = shift_labels[sel] # [n_answer]
335 n = sel_hidden.shape[0]
336 # Chunked lm_head + CE over the answer tokens: never materialize the
337 # full [n_answer, vocab] fp32 logits (≈30GB when a 64k pack is mostly
338 # answer, e.g. long open-ended replies -> OOM). sum/n == mean CE.
339 if n == 0:
340 loss = hidden_states.sum() * 0.0 # keep graph; no supervised token
341 else:
342 CH = 8192
343 tot = hidden_states.new_zeros((), dtype=torch.float32)
344 for s in range(0, n, CH):
345 lg = self.lm_head(sel_hidden[s:s + CH]).float()
346 tot = tot + nn.functional.cross_entropy(
347 lg, sel_lab[s:s + CH], reduction="sum")
348 loss = tot / n
349 return CausalLMOutputWithPast(loss=loss, logits=None,
350 past_key_values=outputs.past_key_values)
351
352 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
353 logits = self.lm_head(hidden_states[:, slice_indices, :])
354 return CausalLMOutputWithPast(
355 loss=None,
356 logits=logits,
357 past_key_values=outputs.past_key_values,
358 hidden_states=outputs.last_hidden_state,
359 )
360
361 # ------------------------------------------------------------------ init
362 @classmethod
363 def from_qwen3(
364 cls,
365 qwen3_path: str,
366 memory_design: str = "Y",
367 mem_layers: list[int] | None = None,
368 mem_o_proj_zero_init: bool = True,
369 dtype: torch.dtype | None = torch.bfloat16,
370 device_map: str | None = None,
371 attn_implementation: str | None = None,
372 **config_overrides,
373 ) -> "LiveMemForCausalLM":
374 """Build a LiveMem from a pretrained Qwen3: load base weights, then
375 copy-init each side branch from the backbone attention geometry."""
376 base = Qwen3ForCausalLM.from_pretrained(
377 qwen3_path, dtype=dtype, attn_implementation=attn_implementation
378 )
379 config = LiveMemConfig(
380 memory_design=memory_design,
381 mem_layers=mem_layers,
382 mem_o_proj_zero_init=mem_o_proj_zero_init,
383 **{**base.config.to_dict(), **config_overrides},
384 )
385 if attn_implementation is not None:
386 config._attn_implementation = attn_implementation
387
388 model = cls(config)
389 if dtype is not None:
390 model = model.to(dtype=dtype)
391
392 # 1) load all base weights that map directly (embed/mlp/norms/lm_head and,
393 # for wrapped layers, the main attention under `.self_attn.attn.*`).
394 sd = _remap_base_state_dict(base.state_dict(), config.memory_layer_indices)
395 missing, unexpected = model.load_state_dict(sd, strict=False)
396 # side-branch params (model.layers.*.self_attn.mem.*) are expected-missing
397 leftover = [k for k in missing if ".self_attn.mem." not in k]
398 if leftover:
399 raise RuntimeError(f"Unexpected missing keys after base load: {leftover[:8]} ...")
400 if unexpected:
401 raise RuntimeError(f"Unexpected keys when loading base: {unexpected[:8]} ...")
402
403 # 2) copy-init each side branch from its (now-loaded) main attention.
404 for idx in config.memory_layer_indices:
405 mattn = model.model.layers[idx].self_attn
406 copy_init_side_branch(mattn.mem, mattn.attn, config.mem_o_proj_zero_init)
407
408 del base
409 if device_map is not None:
410 model = model.to(device_map)
411 return model
412
413
414 def _remap_base_state_dict(state_dict: dict, mem_layers: list[int]) -> dict:
415 """Insert `.attn` into self_attn keys for wrapped layers so base attention
416 weights land on LiveMemAttention.attn.*; all other keys pass through."""
417 mem_set = set(mem_layers)
418 sub = ("q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm")
419 out = {}
420 for k, v in state_dict.items():
421 nk = k
422 if ".self_attn." in k:
423 parts = k.split(".")
424 try:
425 li = parts.index("layers")
426 layer_idx = int(parts[li + 1])
427 except (ValueError, IndexError):
428 layer_idx = None
429 if layer_idx in mem_set and any(f".self_attn.{s}." in k for s in sub):
430 nk = k.replace(".self_attn.", ".self_attn.attn.", 1)
431 out[nk] = v
432 return out
433
434
435 @torch.no_grad()
436 def copy_init_side_branch(
437 side: LiveMemGatedDeltaNet2, attn: Qwen3Attention, zero_o: bool
438 ) -> None:
439 """Copy Qwen3 QKVO into the GDN2 side branch.
440
441 Supports both the legacy full-MHA side branch (32 Q/K/V heads for Qwen3-4B)
442 and the compact KV-head branch (8 Q/K/V heads + expanded V):
443 - Q: direct copy if head counts match; if target heads match backbone KV
444 heads, average the corresponding GQA Q group.
445 - K: copy/adapt from backbone KV heads.
446 - V: copy/adapt from backbone KV heads, then block-repeat each V head along
447 its channel dimension when `expand_v > 1`.
448 - O: copied only when shapes match; normally zero-initialized for training.
449 """
450 hd = attn.head_dim
451 dt = side.q_proj.weight.dtype
452 qh = attn.q_proj.weight.shape[0] // hd
453 kvh = attn.k_proj.weight.shape[0] // hd
454
455 def adapt_heads(heads: torch.Tensor, target_heads: int, name: str) -> torch.Tensor:
456 src_heads = heads.shape[0]
457 if target_heads == src_heads:
458 return heads
459 if target_heads > src_heads and target_heads % src_heads == 0:
460 return heads.repeat_interleave(target_heads // src_heads, dim=0)
461 if src_heads > target_heads and src_heads % target_heads == 0:
462 return heads.view(target_heads, src_heads // target_heads, hd, -1).mean(dim=1)
463 raise ValueError(f"cannot adapt {name} heads from {src_heads} to {target_heads}")
464
465 q_heads = attn.q_proj.weight.view(qh, hd, -1)
466 if side.num_heads == qh:
467 q_init = q_heads
468 elif qh % kvh == 0 and side.num_heads == kvh:
469 q_init = q_heads.view(kvh, qh // kvh, hd, -1).mean(dim=1)
470 else:
471 q_init = adapt_heads(q_heads, side.num_heads, "q")
472 side.q_proj.weight.copy_(q_init.reshape(side.q_proj.weight.shape).to(dt))
473
474 k_heads = attn.k_proj.weight.view(kvh, hd, -1)
475 k_init = adapt_heads(k_heads, side.num_heads, "k")
476 side.k_proj.weight.copy_(k_init.reshape(side.k_proj.weight.shape).to(dt))
477
478 v_heads = adapt_heads(attn.v_proj.weight.view(kvh, hd, -1), side.num_v_heads, "v")
479 if side.head_v_dim % hd != 0:
480 raise ValueError(
481 f"side.head_v_dim={side.head_v_dim} must be a multiple of backbone head_dim={hd} "
482 "for copy initialization"
483 )
484 v_expand = side.head_v_dim // hd
485 v_init = v_heads.repeat(1, v_expand, 1)
486 side.v_proj.weight.copy_(v_init.reshape(side.v_proj.weight.shape).to(dt))
487
488 if zero_o:
489 side.o_proj.weight.zero_()
490 else:
491 if side.o_proj.weight.shape != attn.o_proj.weight.shape:
492 raise ValueError(
493 f"cannot copy-init o_proj with shape {tuple(side.o_proj.weight.shape)} "
494 f"from backbone shape {tuple(attn.o_proj.weight.shape)}; use zero_o=True"
495 )
496 side.o_proj.weight.copy_(attn.o_proj.weight.to(side.o_proj.weight.dtype))
497