3568 lines
171 KiB
Python
3568 lines
171 KiB
Python
"""Integrate dynamic memory into a local Qwen3.5 checkpoint.
|
||
|
||
The original Hugging Face checkpoint is loaded unchanged. Decoder layers are
|
||
wrapped after loading, and only the separate memory module is trainable by
|
||
default.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import asdict, dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
import weakref
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch import Tensor, nn
|
||
|
||
from .memory_os_v2 import (
|
||
KVBudgetManagerV2,
|
||
MemoryOSV2,
|
||
MemoryRouterV2,
|
||
PagedMemoryBankV2,
|
||
memory_record_to_dict,
|
||
)
|
||
from .tiered_memory_store_v2 import TieredMemoryStoreV2
|
||
|
||
|
||
# Qwen's local tokenizer already contains this unused FIM control token. It
|
||
# gives us a reset signal without changing the base vocabulary or resizing a
|
||
# quantized embedding table. Callers may still pass any other token id.
|
||
DEFAULT_MEMORY_RESET_TOKEN = "<|fim_prefix|>"
|
||
|
||
|
||
def split_memory_candidates(text: str, *, max_candidates: int = 8) -> list[str]:
|
||
"""Split one user turn into bounded fact-sized write candidates.
|
||
|
||
This is only a write-side segmentation aid. The model's internal policy
|
||
still decides whether each candidate is durable; the read path remains
|
||
entirely inside the model-owned memory bank.
|
||
"""
|
||
|
||
text = " ".join(str(text).strip().split())
|
||
if not text:
|
||
return []
|
||
# A comma is not a safe write boundary. In natural conversation it is
|
||
# common for the second half of a sentence to negate or qualify the
|
||
# first half (for example, "mentioned X, but it is not important").
|
||
# Splitting there would make the policy see a misleading fragment and can
|
||
# permanently store noise. Keep clauses intact and only split at hard
|
||
# sentence boundaries.
|
||
parts = [
|
||
part.strip()
|
||
for part in re.split(r"(?<=[。！？!?；;])\s*", text)
|
||
if part.strip()
|
||
]
|
||
if len(parts) <= 1:
|
||
return [text]
|
||
if len(parts) <= max_candidates:
|
||
return parts
|
||
# Preserve all text when a long paragraph exceeds the safety bound. The
|
||
# exact sequence can still be stored in one slot without silent loss.
|
||
return [text]
|
||
|
||
|
||
def looks_like_question(text: str) -> bool:
|
||
"""Reject question-shaped turns from the automatic durable-write path."""
|
||
|
||
normalized = " ".join(str(text).strip().split())
|
||
if not normalized:
|
||
return False
|
||
if any(mark in normalized for mark in ("？", "?")):
|
||
return True
|
||
return normalized.startswith(
|
||
(
|
||
"请问",
|
||
"什么",
|
||
"为什么",
|
||
"怎么",
|
||
"如何",
|
||
"能不能",
|
||
"是否",
|
||
"有没有",
|
||
"如果",
|
||
"帮我",
|
||
"解释",
|
||
)
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class QwenMemoryConfig:
|
||
memory_slots: int = 16
|
||
memory_dim: int = 512
|
||
layer_indices: Optional[tuple[int, ...]] = None
|
||
read_scale: float = 1.0
|
||
write_scale: float = 1.0
|
||
mode: str = "residual"
|
||
blend_init: float = 0.0
|
||
direct_logit_scale: float = 0.0
|
||
write_token_offset: Optional[int] = None
|
||
broadcast_write: bool = False
|
||
raw_token_write: bool = False
|
||
raw_logit_scale: float = 0.0
|
||
native_mode: bool = False
|
||
persistent_memory: bool = False
|
||
reset_token_id: Optional[int] = None
|
||
summary_pooling: bool = True
|
||
# Natural-language episodic memory keeps the exact token sequence of a
|
||
# selected write inside the model-owned memory state. The model retrieves
|
||
# and injects these tokens internally; callers never need to replay chat
|
||
# history or assemble a prompt by hand.
|
||
natural_language_memory: bool = False
|
||
text_memory_tokens: int = 256
|
||
text_memory_top_k: int = 2
|
||
text_memory_threshold: float = 0.40
|
||
text_memory_write_threshold: float = 0.5
|
||
text_memory_replace_threshold: float = 0.35
|
||
text_memory_key_tokens: int = 128
|
||
text_memory_overlap_threshold: float = 0.22
|
||
automatic_memory: bool = False
|
||
auto_memory_threshold: float = 0.35
|
||
# v2 adds a separately trained natural-language forget decision. Version
|
||
# 1 keeps the old one-logit policy byte-compatible with existing adapters.
|
||
automatic_memory_policy_version: int = 1
|
||
auto_forget_threshold: float = 0.50
|
||
# Natural Memory v2: compact addressed pages sit beside the v1 hot text
|
||
# bank. V2 is opt-in so existing checkpoints remain byte-compatible.
|
||
memory_version: int = 1
|
||
hierarchical_memory: bool = False
|
||
memory_router_dim: int = 128
|
||
memory_router_heads: int = 8
|
||
memory_page_capacity: int = 32
|
||
# 32K pages x 32 records/page gives a one-million-record address space.
|
||
# Pages are allocated lazily; this is a capacity, not a startup tensor.
|
||
memory_max_pages: int = 32768
|
||
memory_hot_pages: int = 8
|
||
memory_top_k_pages: int = 4
|
||
memory_top_k_records: int = 8
|
||
memory_max_hops: int = 3
|
||
memory_coarse_index_bits: int = 20
|
||
memory_v2_read_threshold: float = 0.65
|
||
memory_v2_write_threshold: float = 0.50
|
||
# ``embedded`` keeps the portable user snapshot inside the safetensors
|
||
# package. ``tiered`` additionally opens a durable page store so cold
|
||
# records can leave RAM while the same model-owned reader remains active.
|
||
memory_storage_mode: str = "embedded"
|
||
memory_storage_path: Optional[str] = None
|
||
memory_resident_pages: int = 256
|
||
# Embedded mode keeps every record in process RAM. Only this bounded hot
|
||
# cache is copied to the execution device; it is not a second store.
|
||
memory_gpu_cache_records: int = 256
|
||
memory_gpu_cache_tokens: int = 131072
|
||
# Keep a large VRAM safety margin for the model weights, activations and
|
||
# KV cache. Memory records that do not fit remain in system RAM.
|
||
memory_gpu_cache_reserve_mb: int = 2048
|
||
memory_gpu_cache_adaptive: bool = True
|
||
# HF generation can keep the active cache on CPU while the current layer
|
||
# is executing. This is opt-in because it trades GPU memory for PCIe
|
||
# traffic and is therefore not always faster on short prompts.
|
||
kv_offload: bool = False
|
||
kv_cache_implementation: Optional[str] = None
|
||
kv_offload_only_non_sliding: bool = True
|
||
# When a prompt exceeds the hot-window budget, archive the old prefix as
|
||
# V2 context records and run Qwen only on the recent window. The reader
|
||
# remains model-owned; callers do not have to assemble a retrieval prompt.
|
||
auto_compact_context: bool = True
|
||
context_chunk_tokens: int = 512
|
||
context_archive_max_records: int = 1_048_576
|
||
kv_budget_tokens: int = 32768
|
||
kv_hard_max_tokens: int = 131072
|
||
kv_compaction_trigger: float = 0.90
|
||
kv_keep_recent_tokens: int = 8192
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.mode not in {"residual", "blend", "replace"}:
|
||
raise ValueError("mode must be one of: residual, blend, replace")
|
||
if not 0.0 <= self.blend_init <= 1.0:
|
||
raise ValueError("blend_init must be between 0 and 1")
|
||
if self.direct_logit_scale < 0.0:
|
||
raise ValueError("direct_logit_scale must be non-negative")
|
||
if self.write_token_offset is not None and self.write_token_offset < 1:
|
||
raise ValueError("write_token_offset must be positive")
|
||
if self.raw_logit_scale < 0.0:
|
||
raise ValueError("raw_logit_scale must be non-negative")
|
||
if self.text_memory_tokens < 1:
|
||
raise ValueError("text_memory_tokens must be positive")
|
||
if self.text_memory_top_k < 1:
|
||
raise ValueError("text_memory_top_k must be positive")
|
||
if not 0.0 <= self.text_memory_threshold <= 1.0:
|
||
raise ValueError("text_memory_threshold must be between 0 and 1")
|
||
if not 0.0 <= self.text_memory_write_threshold <= 1.0:
|
||
raise ValueError("text_memory_write_threshold must be between 0 and 1")
|
||
if not -1.0 <= self.text_memory_replace_threshold <= 1.0:
|
||
raise ValueError("text_memory_replace_threshold must be between -1 and 1")
|
||
if self.text_memory_key_tokens < 1:
|
||
raise ValueError("text_memory_key_tokens must be positive")
|
||
if not 0.0 <= self.text_memory_overlap_threshold <= 1.0:
|
||
raise ValueError("text_memory_overlap_threshold must be between 0 and 1")
|
||
if not 0.0 <= self.auto_memory_threshold <= 1.0:
|
||
raise ValueError("auto_memory_threshold must be between 0 and 1")
|
||
if self.automatic_memory_policy_version < 1:
|
||
raise ValueError("automatic_memory_policy_version must be positive")
|
||
if not 0.0 <= self.auto_forget_threshold <= 1.0:
|
||
raise ValueError("auto_forget_threshold must be between 0 and 1")
|
||
if self.memory_version < 1:
|
||
raise ValueError("memory_version must be positive")
|
||
if self.memory_router_heads < 1:
|
||
raise ValueError("memory_router_heads must be positive")
|
||
if self.memory_router_dim < 8 or self.memory_router_dim % self.memory_router_heads != 0:
|
||
raise ValueError("memory_router_dim must be divisible by memory_router_heads")
|
||
if self.memory_page_capacity < 1 or self.memory_max_pages < 1:
|
||
raise ValueError("memory page limits must be positive")
|
||
if self.memory_top_k_pages < 1 or self.memory_top_k_records < 1:
|
||
raise ValueError("memory top-k limits must be positive")
|
||
if self.memory_max_hops < 1:
|
||
raise ValueError("memory_max_hops must be positive")
|
||
if not 0.0 <= self.memory_v2_read_threshold <= 1.0:
|
||
raise ValueError("memory_v2_read_threshold must be between 0 and 1")
|
||
if not 0.0 <= self.memory_v2_write_threshold <= 1.0:
|
||
raise ValueError("memory_v2_write_threshold must be between 0 and 1")
|
||
if self.memory_storage_mode not in {"embedded", "tiered"}:
|
||
raise ValueError("memory_storage_mode must be embedded or tiered")
|
||
if self.memory_resident_pages < self.memory_hot_pages:
|
||
raise ValueError("memory_resident_pages must be >= memory_hot_pages")
|
||
if self.memory_gpu_cache_records < 0 or self.memory_gpu_cache_tokens < 0:
|
||
raise ValueError("memory GPU cache limits must be non-negative")
|
||
if self.memory_gpu_cache_reserve_mb < 0:
|
||
raise ValueError("memory_gpu_cache_reserve_mb must be non-negative")
|
||
if self.kv_cache_implementation is not None and not str(self.kv_cache_implementation).strip():
|
||
raise ValueError("kv_cache_implementation must be non-empty when provided")
|
||
if self.context_chunk_tokens < 1:
|
||
raise ValueError("context_chunk_tokens must be positive")
|
||
if self.context_archive_max_records < 1:
|
||
raise ValueError("context_archive_max_records must be positive")
|
||
if not 1 <= self.kv_budget_tokens <= self.kv_hard_max_tokens:
|
||
raise ValueError("kv_budget_tokens must be inside [1, kv_hard_max_tokens]")
|
||
|
||
def resolved_layers(self, num_hidden_layers: int) -> tuple[int, ...]:
|
||
if self.layer_indices is not None:
|
||
layers = tuple(sorted(set(self.layer_indices)))
|
||
else:
|
||
stride = max(1, num_hidden_layers // 4)
|
||
layers = tuple(min(num_hidden_layers - 1, stride * i - 1) for i in range(1, 5))
|
||
layers = tuple(sorted(set(layers)))
|
||
if not layers or min(layers) < 0 or max(layers) >= num_hidden_layers:
|
||
raise ValueError(f"layer_indices must be inside [0, {num_hidden_layers})")
|
||
return layers
|
||
|
||
|
||
def load_memory_config(adapter_dir: str | Path) -> QwenMemoryConfig:
|
||
"""Load the architecture/configuration metadata saved with an adapter."""
|
||
|
||
metadata_path = Path(adapter_dir) / "memory_config.json"
|
||
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||
saved = metadata.get("memory_config", {})
|
||
saved_layers = metadata.get("layer_indices", saved.get("layer_indices"))
|
||
return QwenMemoryConfig(
|
||
memory_slots=int(saved.get("memory_slots", 16)),
|
||
memory_dim=int(saved.get("memory_dim", 512)),
|
||
layer_indices=tuple(saved_layers) if saved_layers is not None else None,
|
||
read_scale=float(saved.get("read_scale", 1.0)),
|
||
write_scale=float(saved.get("write_scale", 1.0)),
|
||
mode=str(saved.get("mode", "residual")),
|
||
blend_init=float(saved.get("blend_init", 0.0)),
|
||
direct_logit_scale=float(saved.get("direct_logit_scale", 0.0)),
|
||
write_token_offset=(
|
||
int(saved["write_token_offset"]) if saved.get("write_token_offset") is not None else None
|
||
),
|
||
broadcast_write=bool(saved.get("broadcast_write", False)),
|
||
raw_token_write=bool(saved.get("raw_token_write", False)),
|
||
raw_logit_scale=float(saved.get("raw_logit_scale", 0.0)),
|
||
native_mode=bool(saved.get("native_mode", False)),
|
||
persistent_memory=bool(saved.get("persistent_memory", False)),
|
||
reset_token_id=(int(saved["reset_token_id"]) if saved.get("reset_token_id") is not None else None),
|
||
summary_pooling=bool(saved.get("summary_pooling", True)),
|
||
natural_language_memory=bool(saved.get("natural_language_memory", False)),
|
||
text_memory_tokens=int(saved.get("text_memory_tokens", 256)),
|
||
text_memory_top_k=int(saved.get("text_memory_top_k", 2)),
|
||
text_memory_threshold=float(saved.get("text_memory_threshold", 0.40)),
|
||
text_memory_write_threshold=float(saved.get("text_memory_write_threshold", 0.5)),
|
||
text_memory_replace_threshold=float(saved.get("text_memory_replace_threshold", 0.35)),
|
||
text_memory_key_tokens=int(saved.get("text_memory_key_tokens", 128)),
|
||
text_memory_overlap_threshold=float(saved.get("text_memory_overlap_threshold", 0.22)),
|
||
automatic_memory=bool(saved.get("automatic_memory", False)),
|
||
auto_memory_threshold=float(saved.get("auto_memory_threshold", 0.35)),
|
||
automatic_memory_policy_version=int(saved.get("automatic_memory_policy_version", 1)),
|
||
auto_forget_threshold=float(saved.get("auto_forget_threshold", 0.50)),
|
||
memory_version=int(saved.get("memory_version", 1)),
|
||
hierarchical_memory=bool(saved.get("hierarchical_memory", False)),
|
||
memory_router_dim=int(saved.get("memory_router_dim", 128)),
|
||
memory_router_heads=int(saved.get("memory_router_heads", 8)),
|
||
memory_page_capacity=int(saved.get("memory_page_capacity", 32)),
|
||
memory_max_pages=int(saved.get("memory_max_pages", 32768)),
|
||
memory_hot_pages=int(saved.get("memory_hot_pages", 8)),
|
||
memory_top_k_pages=int(saved.get("memory_top_k_pages", 4)),
|
||
memory_top_k_records=int(saved.get("memory_top_k_records", 8)),
|
||
memory_max_hops=int(saved.get("memory_max_hops", 3)),
|
||
memory_coarse_index_bits=int(saved.get("memory_coarse_index_bits", 20)),
|
||
memory_v2_read_threshold=float(saved.get("memory_v2_read_threshold", 0.50)),
|
||
memory_v2_write_threshold=float(saved.get("memory_v2_write_threshold", 0.50)),
|
||
memory_storage_mode=str(saved.get("memory_storage_mode", "embedded")),
|
||
memory_storage_path=(str(saved["memory_storage_path"]) if saved.get("memory_storage_path") else None),
|
||
memory_resident_pages=int(saved.get("memory_resident_pages", 256)),
|
||
memory_gpu_cache_records=int(saved.get("memory_gpu_cache_records", 256)),
|
||
memory_gpu_cache_tokens=int(saved.get("memory_gpu_cache_tokens", 131072)),
|
||
memory_gpu_cache_reserve_mb=int(saved.get("memory_gpu_cache_reserve_mb", 2048)),
|
||
memory_gpu_cache_adaptive=bool(saved.get("memory_gpu_cache_adaptive", True)),
|
||
kv_offload=bool(saved.get("kv_offload", False)),
|
||
kv_cache_implementation=(
|
||
str(saved["kv_cache_implementation"])
|
||
if saved.get("kv_cache_implementation") is not None
|
||
else None
|
||
),
|
||
kv_offload_only_non_sliding=bool(saved.get("kv_offload_only_non_sliding", True)),
|
||
auto_compact_context=bool(saved.get("auto_compact_context", True)),
|
||
context_chunk_tokens=int(saved.get("context_chunk_tokens", 512)),
|
||
context_archive_max_records=int(saved.get("context_archive_max_records", 1_048_576)),
|
||
kv_budget_tokens=int(saved.get("kv_budget_tokens", 32768)),
|
||
kv_hard_max_tokens=int(saved.get("kv_hard_max_tokens", 131072)),
|
||
kv_compaction_trigger=float(saved.get("kv_compaction_trigger", 0.90)),
|
||
kv_keep_recent_tokens=int(saved.get("kv_keep_recent_tokens", 8192)),
|
||
)
|
||
|
||
|
||
def resolve_memory_reset_token(tokenizer: Any, token: str = DEFAULT_MEMORY_RESET_TOKEN) -> int:
|
||
"""Resolve a tokenizer control token suitable for clearing memory."""
|
||
|
||
token_id = tokenizer.convert_tokens_to_ids(token)
|
||
unknown_id = getattr(tokenizer, "unk_token_id", None)
|
||
if token_id is None or (unknown_id is not None and int(token_id) == int(unknown_id)):
|
||
raise ValueError(
|
||
f"reset token {token!r} is not present in this tokenizer; pass --reset-token-id instead"
|
||
)
|
||
return int(token_id)
|
||
|
||
|
||
class QwenDynamicMemory(nn.Module):
|
||
"""A compact float32 key-value memory connected to Qwen hidden states."""
|
||
|
||
def __init__(self, hidden_size: int, config: QwenMemoryConfig) -> None:
|
||
super().__init__()
|
||
self.hidden_size = hidden_size
|
||
self.config = config
|
||
dim = config.memory_dim
|
||
self.query = nn.Linear(hidden_size, dim, bias=False)
|
||
self.key = nn.Linear(dim, dim, bias=False)
|
||
self.value = nn.Linear(dim, dim, bias=False)
|
||
self.read_out = nn.Linear(dim, hidden_size, bias=False)
|
||
self.read_gate = nn.Linear(hidden_size, 1)
|
||
self.slot_keys = nn.Parameter(torch.randn(config.memory_slots, dim) / dim**0.5)
|
||
self.write_key = nn.Linear(hidden_size, dim, bias=False)
|
||
self.write_value = nn.Linear(hidden_size, dim, bias=False)
|
||
self.write_gate = nn.Linear(hidden_size, config.memory_slots)
|
||
self.apply(self._init_weights)
|
||
self.last_read_address: Optional[Tensor] = None
|
||
self.last_read_relevance: Optional[Tensor] = None
|
||
self.last_write_address: Optional[Tensor] = None
|
||
self.last_write_probability: Optional[Tensor] = None
|
||
|
||
@staticmethod
|
||
def _init_weights(module: nn.Module) -> None:
|
||
if isinstance(module, nn.Linear):
|
||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||
|
||
def initial_state(self, batch_size: int, *, device: torch.device) -> Tensor:
|
||
return torch.zeros(
|
||
batch_size,
|
||
self.config.memory_slots,
|
||
self.config.memory_dim,
|
||
device=device,
|
||
dtype=self.slot_keys.dtype,
|
||
)
|
||
|
||
def read(self, hidden_states: Tensor, memory: Tensor) -> Tensor:
|
||
work_dtype = self.slot_keys.dtype
|
||
hidden = hidden_states.to(work_dtype)
|
||
memory = memory.to(device=hidden.device, dtype=work_dtype)
|
||
q = self.query(hidden)
|
||
k = self.key(memory)
|
||
v = self.value(memory)
|
||
scores = torch.matmul(q, k.transpose(-1, -2)) / self.config.memory_dim**0.5
|
||
address = scores.softmax(dim=-1)
|
||
retrieved = torch.matmul(address, v)
|
||
retrieved = self.read_out(retrieved)
|
||
gate = torch.sigmoid(self.read_gate(hidden))
|
||
self.last_read_address = address
|
||
self.last_read_relevance = gate * address.max(dim=-1, keepdim=True).values
|
||
return (self.config.read_scale * gate * retrieved).to(hidden_states.dtype)
|
||
|
||
def update(
|
||
self,
|
||
hidden_states: Tensor,
|
||
memory: Tensor,
|
||
attention_mask: Optional[Tensor] = None,
|
||
) -> Tensor:
|
||
work_dtype = self.slot_keys.dtype
|
||
hidden = hidden_states.to(work_dtype)
|
||
memory = memory.to(device=hidden.device, dtype=work_dtype)
|
||
if self.config.write_token_offset is None:
|
||
summary = hidden[:, -1]
|
||
else:
|
||
if hidden.shape[1] < self.config.write_token_offset:
|
||
raise ValueError("write_token_offset exceeds the memory sequence length")
|
||
summary = hidden[:, -self.config.write_token_offset]
|
||
proposal = self.write_value(summary)
|
||
if self.config.broadcast_write:
|
||
strength = torch.ones(
|
||
hidden.shape[0],
|
||
self.config.memory_slots,
|
||
device=hidden.device,
|
||
dtype=work_dtype,
|
||
)
|
||
address = strength / float(self.config.memory_slots)
|
||
else:
|
||
address = (self.write_key(summary) @ self.slot_keys.t()).softmax(dim=-1)
|
||
strength = torch.sigmoid(self.write_gate(summary)) * address
|
||
self.last_write_address = address
|
||
self.last_write_probability = strength.sum(dim=-1, keepdim=True)
|
||
strength = (self.config.write_scale * strength).unsqueeze(-1)
|
||
proposal = proposal[:, None, :].expand(-1, self.config.memory_slots, -1)
|
||
return memory + strength * (proposal - memory)
|
||
|
||
|
||
class NaturalLanguageRetriever(nn.Module):
|
||
"""Trainable pair scorer for query-to-episodic-memory retrieval."""
|
||
|
||
def __init__(self, hidden_size: int, projection_size: int = 256) -> None:
|
||
super().__init__()
|
||
self.query_projection = nn.Linear(hidden_size, projection_size, bias=False)
|
||
self.key_projection = nn.Linear(hidden_size, projection_size, bias=False)
|
||
self.pair_scorer = nn.Sequential(
|
||
nn.Linear(projection_size * 4, projection_size),
|
||
nn.SiLU(),
|
||
nn.Linear(projection_size, 1),
|
||
)
|
||
self.apply(self._init_weights)
|
||
|
||
@staticmethod
|
||
def _init_weights(module: nn.Module) -> None:
|
||
if isinstance(module, nn.Linear):
|
||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||
|
||
def forward(self, query: Tensor, keys: Tensor) -> Tensor:
|
||
query = F.normalize(query.float(), dim=-1)
|
||
keys = F.normalize(keys.float(), dim=-1)
|
||
query_projection = self.query_projection(query)
|
||
key_projection = self.key_projection(keys)
|
||
if key_projection.ndim == 3:
|
||
query_projection = query_projection.unsqueeze(1)
|
||
features = torch.cat(
|
||
(
|
||
query_projection.expand_as(key_projection),
|
||
key_projection,
|
||
query_projection.expand_as(key_projection) * key_projection,
|
||
(query_projection.expand_as(key_projection) - key_projection).abs(),
|
||
),
|
||
dim=-1,
|
||
)
|
||
return self.pair_scorer(features).squeeze(-1)
|
||
|
||
|
||
class AutomaticMemoryPolicy(nn.Module):
|
||
"""High-recall controller for deciding whether a user turn is durable."""
|
||
|
||
def __init__(self, hidden_size: int, projection_size: int = 256) -> None:
|
||
super().__init__()
|
||
self.encoder = nn.Sequential(
|
||
nn.Linear(hidden_size, projection_size),
|
||
nn.LayerNorm(projection_size),
|
||
nn.SiLU(),
|
||
nn.Linear(projection_size, projection_size // 2),
|
||
nn.SiLU(),
|
||
)
|
||
self.importance = nn.Linear(projection_size // 2, 1)
|
||
self.apply(self._init_weights)
|
||
|
||
@staticmethod
|
||
def _init_weights(module: nn.Module) -> None:
|
||
if isinstance(module, nn.Linear):
|
||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||
if module.bias is not None:
|
||
nn.init.zeros_(module.bias)
|
||
|
||
def forward(self, hidden: Tensor) -> Tensor:
|
||
return self.importance(self.encoder(hidden.float())).squeeze(-1)
|
||
|
||
|
||
class AutomaticMemoryPolicyV2(AutomaticMemoryPolicy):
|
||
"""Write/forget policy trained on the same natural-language turn state."""
|
||
|
||
def __init__(self, hidden_size: int, projection_size: int = 256) -> None:
|
||
super().__init__(hidden_size, projection_size)
|
||
self.forget = nn.Linear(projection_size // 2, 1)
|
||
nn.init.normal_(self.forget.weight, mean=0.0, std=0.02)
|
||
nn.init.zeros_(self.forget.bias)
|
||
|
||
def _features(self, hidden: Tensor) -> Tensor:
|
||
return self.encoder(hidden.float())
|
||
|
||
def forward(self, hidden: Tensor) -> Tensor:
|
||
return self.importance(self._features(hidden)).squeeze(-1)
|
||
|
||
def forget_logits(self, hidden: Tensor) -> Tensor:
|
||
return self.forget(self._features(hidden)).squeeze(-1)
|
||
|
||
|
||
class NativeQwenDynamicMemory(QwenDynamicMemory):
|
||
"""Learned write/forget controller for model-owned persistent memory."""
|
||
|
||
def __init__(self, hidden_size: int, config: QwenMemoryConfig) -> None:
|
||
super().__init__(hidden_size, config)
|
||
dim = config.memory_dim
|
||
self.summary_score = nn.Linear(hidden_size, 1, bias=False)
|
||
self.write_decision = nn.Linear(hidden_size, 1)
|
||
# Forgetting must depend on both the incoming candidate and what is
|
||
# already stored. A hidden-only gate can learn "this sentence looks
|
||
# like a fact", but it cannot learn "this fact conflicts with the
|
||
# value in the addressed slot".
|
||
self.forget_gate = nn.Sequential(
|
||
nn.Linear(hidden_size + (2 * dim), dim),
|
||
nn.SiLU(),
|
||
nn.Linear(dim, config.memory_slots),
|
||
)
|
||
nn.init.normal_(self.summary_score.weight, mean=0.0, std=0.02)
|
||
nn.init.normal_(self.write_decision.weight, mean=0.0, std=0.02)
|
||
for module in self.forget_gate:
|
||
if isinstance(module, nn.Linear):
|
||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||
nn.init.constant_(self.write_decision.bias, -1.5)
|
||
nn.init.constant_(self.forget_gate[-1].bias, -2.0)
|
||
self.last_write_probability: Optional[Tensor] = None
|
||
self.last_forget_probability: Optional[Tensor] = None
|
||
# Keep the semantic summary available to the automatic policy. The
|
||
# older ``last_write_representation`` is optimized for the memory
|
||
# value path and may discard distinctions such as negation or a
|
||
# question-shaped turn.
|
||
self.last_write_summary: Optional[Tensor] = None
|
||
self.last_write_representation: Optional[Tensor] = None
|
||
|
||
def _summary(self, hidden_states: Tensor, attention_mask: Optional[Tensor]) -> Tensor:
|
||
work_dtype = self.slot_keys.dtype
|
||
hidden = hidden_states.to(work_dtype)
|
||
if self.config.write_token_offset is not None:
|
||
if hidden.shape[1] < self.config.write_token_offset:
|
||
raise ValueError("write_token_offset exceeds the memory sequence length")
|
||
return hidden[:, -self.config.write_token_offset]
|
||
if not self.config.summary_pooling:
|
||
return hidden[:, -1]
|
||
scores = self.summary_score(hidden).squeeze(-1)
|
||
if attention_mask is not None and attention_mask.ndim == 2:
|
||
# During cached generation the decoder may expose only the newest
|
||
# token while the attention mask still covers the full sequence.
|
||
# Align the mask to the score sequence instead of assuming both
|
||
# lengths are identical.
|
||
score_length = scores.shape[1]
|
||
mask = attention_mask.to(device=hidden.device, dtype=torch.bool)
|
||
if mask.shape[1] > score_length:
|
||
mask = mask[:, -score_length:]
|
||
elif mask.shape[1] < score_length:
|
||
pad = torch.ones(
|
||
mask.shape[0],
|
||
score_length - mask.shape[1],
|
||
device=mask.device,
|
||
dtype=mask.dtype,
|
||
)
|
||
mask = torch.cat((mask, pad), dim=1)
|
||
scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
|
||
weights = scores.softmax(dim=-1)
|
||
return (weights.unsqueeze(-1) * hidden).sum(dim=1)
|
||
|
||
def update(
|
||
self,
|
||
hidden_states: Tensor,
|
||
memory: Tensor,
|
||
attention_mask: Optional[Tensor] = None,
|
||
) -> Tensor:
|
||
work_dtype = self.slot_keys.dtype
|
||
memory = memory.to(device=hidden_states.device, dtype=work_dtype)
|
||
summary = self._summary(hidden_states, attention_mask)
|
||
self.last_write_summary = summary
|
||
proposal = self.write_value(summary)
|
||
self.last_write_representation = self.read_out(self.value(proposal))
|
||
address = (self.write_key(summary) @ self.slot_keys.t()).softmax(dim=-1)
|
||
slot_probability = torch.sigmoid(self.write_gate(summary))
|
||
write_probability = torch.sigmoid(self.write_decision(summary))
|
||
self.last_write_address = address
|
||
matched_memory = torch.sum(address.unsqueeze(-1) * memory, dim=1)
|
||
forget_features = torch.cat((summary, proposal, matched_memory), dim=-1)
|
||
forget_probability = torch.sigmoid(self.forget_gate(forget_features))
|
||
strength = self.config.write_scale * write_probability * slot_probability * address
|
||
erase = strength * forget_probability
|
||
proposal = proposal[:, None, :]
|
||
updated = memory * (1.0 - erase.unsqueeze(-1))
|
||
updated = updated + strength.unsqueeze(-1) * (proposal - updated)
|
||
self.last_write_probability = write_probability
|
||
self.last_forget_probability = forget_probability
|
||
return updated
|
||
|
||
|
||
class _MemoryRuntime:
|
||
def __init__(self, memory: QwenDynamicMemory) -> None:
|
||
self.memory = memory
|
||
self.state: Optional[Tensor] = None
|
||
self.read_enabled = True
|
||
self.update_enabled = True
|
||
self.last_read: Optional[Tensor] = None
|
||
self.raw_memory: Optional[Tensor] = None
|
||
self.input_ids: Optional[Tensor] = None
|
||
self.attention_mask: Optional[Tensor] = None
|
||
self.output_embeddings: Optional[nn.Module] = None
|
||
self.reset_mask: Optional[Tensor] = None
|
||
self.text_token_ids: Optional[Tensor] = None
|
||
self.text_token_mask: Optional[Tensor] = None
|
||
self.text_slot_valid: Optional[Tensor] = None
|
||
self.text_slot_keys: Optional[Tensor] = None
|
||
self.text_slot_age: Optional[Tensor] = None
|
||
self.text_write_counter: Optional[Tensor] = None
|
||
self.text_key_token_ids: Optional[Tensor] = None
|
||
self.text_key_token_mask: Optional[Tensor] = None
|
||
self.text_last_written_slot: Optional[Tensor] = None
|
||
self.text_read_slots: Optional[Tensor] = None
|
||
self.text_read_relevance: Optional[Tensor] = None
|
||
self.text_read_overlap: Optional[Tensor] = None
|
||
self.v2_query_key: Optional[Tensor] = None
|
||
self.v2_last_decisions: list[Any] = []
|
||
self.text_prefix_used: bool = False
|
||
self.auto_memory_probability: Optional[Tensor] = None
|
||
self.auto_memory_forget_probability: Optional[Tensor] = None
|
||
self.context_compaction: Optional[dict[str, Any]] = None
|
||
# Persistent checkpoints are the durable user state. Evaluation,
|
||
# branch execution and a new temporary conversation must be able to
|
||
# reset only their working state without mutating that checkpoint.
|
||
self.use_persistent_state: bool = True
|
||
|
||
|
||
class MemoryLayerAdapter(nn.Module):
|
||
"""Wrap one Qwen decoder layer and optionally replace its token mixer."""
|
||
|
||
def __init__(
|
||
self,
|
||
inner: nn.Module,
|
||
runtime: _MemoryRuntime,
|
||
*,
|
||
read: bool,
|
||
write: bool,
|
||
mode: str = "residual",
|
||
blend_init: float = 0.0,
|
||
) -> None:
|
||
super().__init__()
|
||
if mode not in {"residual", "blend", "replace"}:
|
||
raise ValueError("mode must be one of: residual, blend, replace")
|
||
self.inner = inner
|
||
self._runtime_ref = weakref.ref(runtime)
|
||
self.read_enabled = read
|
||
self.write_enabled = write
|
||
self.mode = mode
|
||
if mode == "blend":
|
||
blend_init = min(max(float(blend_init), 1e-4), 1.0 - 1e-4)
|
||
logit = torch.logit(torch.tensor(blend_init, dtype=torch.float32))
|
||
inner_device = next(
|
||
(parameter.device for parameter in inner.parameters() if parameter.device.type != "meta"),
|
||
None,
|
||
)
|
||
if inner_device is not None:
|
||
logit = logit.to(inner_device)
|
||
self.blend_logit = nn.Parameter(logit)
|
||
|
||
def _original_token_mixer(
|
||
self,
|
||
normalized_hidden: Tensor,
|
||
position_embeddings: Any,
|
||
attention_mask: Optional[Tensor],
|
||
position_ids: Optional[Tensor],
|
||
past_key_values: Any,
|
||
kwargs: dict[str, Any],
|
||
) -> Tensor:
|
||
layer_type = getattr(self.inner, "layer_type", None)
|
||
if layer_type == "linear_attention":
|
||
output = self.inner.linear_attn(
|
||
hidden_states=normalized_hidden,
|
||
cache_params=past_key_values,
|
||
attention_mask=attention_mask,
|
||
**kwargs,
|
||
)
|
||
return output[0] if isinstance(output, tuple) else output
|
||
if layer_type == "full_attention":
|
||
output = self.inner.self_attn(
|
||
hidden_states=normalized_hidden,
|
||
attention_mask=attention_mask,
|
||
position_ids=position_ids,
|
||
past_key_values=past_key_values,
|
||
position_embeddings=position_embeddings,
|
||
**kwargs,
|
||
)
|
||
return output[0] if isinstance(output, tuple) else output
|
||
raise TypeError(f"unsupported Qwen layer_type for surgery: {layer_type!r}")
|
||
|
||
def _prefetch_cache_layer(self, past_key_values: Any) -> None:
|
||
"""Make an offloaded Qwen hybrid-cache layer ready before it is read."""
|
||
|
||
if past_key_values is None or not getattr(past_key_values, "offloading", False):
|
||
return
|
||
component = getattr(self.inner, "self_attn", None)
|
||
if component is None:
|
||
component = getattr(self.inner, "linear_attn", None)
|
||
layer_index = getattr(component, "layer_idx", None)
|
||
if layer_index is None or not hasattr(past_key_values, "prefetch"):
|
||
return
|
||
only_non_sliding = bool(getattr(past_key_values, "only_non_sliding", True))
|
||
parameter = next(component.parameters(), None)
|
||
device = parameter.device if parameter is not None else torch.device("cpu")
|
||
past_key_values.prefetch(int(layer_index), only_non_sliding)
|
||
# The stock offload helper uses a separate CUDA stream. Qwen3.5
|
||
# reads linear-attention state before it calls ``update_*``; make the
|
||
# current layer's ownership synchronous here so a CPU state can never
|
||
# reach a CUDA convolution or attention kernel.
|
||
if 0 <= int(layer_index) < len(past_key_values.layers):
|
||
cache_layer = past_key_values.layers[int(layer_index)]
|
||
for name in ("keys", "values", "conv_states", "recurrent_states"):
|
||
value = getattr(cache_layer, name, None)
|
||
if isinstance(value, Tensor) and value.device != device:
|
||
setattr(cache_layer, name, value.to(device=device, non_blocking=False))
|
||
|
||
def _memory_read(self, hidden_states: Tensor, runtime: _MemoryRuntime) -> Tensor:
|
||
if runtime.read_enabled and self.read_enabled and runtime.state is not None:
|
||
output = runtime.memory.read(hidden_states, runtime.state)
|
||
else:
|
||
output = torch.zeros_like(hidden_states)
|
||
if self.write_enabled:
|
||
runtime.last_read = output
|
||
return output
|
||
|
||
def _surgical_forward(
|
||
self,
|
||
hidden_states: Tensor,
|
||
position_embeddings: Any,
|
||
attention_mask: Optional[Tensor],
|
||
position_ids: Optional[Tensor],
|
||
past_key_values: Any,
|
||
kwargs: dict[str, Any],
|
||
runtime: _MemoryRuntime,
|
||
) -> Tensor:
|
||
residual = hidden_states
|
||
normalized_hidden = self.inner.input_layernorm(hidden_states)
|
||
memory_output = self._memory_read(normalized_hidden, runtime)
|
||
|
||
if self.mode == "replace":
|
||
token_mixer = memory_output
|
||
else:
|
||
original_output = self._original_token_mixer(
|
||
normalized_hidden,
|
||
position_embeddings,
|
||
attention_mask,
|
||
position_ids,
|
||
past_key_values,
|
||
kwargs,
|
||
)
|
||
mix = torch.sigmoid(self.blend_logit).to(dtype=original_output.dtype)
|
||
# A zero/empty memory must be an exact no-op. This prevents a
|
||
# trained blend coefficient from attenuating the original Qwen
|
||
# token mixer on ordinary prompts that have no user memory.
|
||
if runtime.state is None or not runtime.read_enabled:
|
||
# Exact zero in the forward pass, but retain a surrogate
|
||
# gradient so the blend parameter remains trainable in unit
|
||
# tests and future calibration runs.
|
||
mix = mix - mix.detach()
|
||
else:
|
||
memory_activity = runtime.state.detach().abs().mean()
|
||
if float(memory_activity) <= 1e-6:
|
||
mix = mix - mix.detach()
|
||
token_mixer = (1.0 - mix) * original_output + mix * memory_output
|
||
|
||
hidden_states = residual + token_mixer
|
||
residual = hidden_states
|
||
hidden_states = self.inner.post_attention_layernorm(hidden_states)
|
||
hidden_states = self.inner.mlp(hidden_states)
|
||
return residual + hidden_states
|
||
|
||
def forward(
|
||
self,
|
||
hidden_states: Tensor,
|
||
position_embeddings: Any = None,
|
||
attention_mask: Optional[Tensor] = None,
|
||
position_ids: Optional[Tensor] = None,
|
||
past_key_values: Any = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
runtime = self._runtime_ref()
|
||
if runtime is None:
|
||
return self.inner(
|
||
hidden_states,
|
||
position_embeddings,
|
||
attention_mask=attention_mask,
|
||
position_ids=position_ids,
|
||
past_key_values=past_key_values,
|
||
**kwargs,
|
||
)
|
||
|
||
self._prefetch_cache_layer(past_key_values)
|
||
if self.mode == "residual":
|
||
hidden_states = hidden_states + self._memory_read(hidden_states, runtime)
|
||
output = self.inner(
|
||
hidden_states,
|
||
position_embeddings,
|
||
attention_mask=attention_mask,
|
||
position_ids=position_ids,
|
||
past_key_values=past_key_values,
|
||
**kwargs,
|
||
)
|
||
else:
|
||
output = self._surgical_forward(
|
||
hidden_states,
|
||
position_embeddings,
|
||
attention_mask,
|
||
position_ids,
|
||
past_key_values,
|
||
kwargs,
|
||
runtime,
|
||
)
|
||
hidden_output = output[0] if isinstance(output, tuple) else output
|
||
if runtime.update_enabled and self.write_enabled and runtime.state is not None:
|
||
updated_state = runtime.memory.update(
|
||
hidden_output,
|
||
runtime.state,
|
||
attention_mask=runtime.attention_mask,
|
||
)
|
||
if runtime.reset_mask is not None:
|
||
reset_mask = runtime.reset_mask.to(device=updated_state.device, dtype=torch.bool)
|
||
updated_state = torch.where(
|
||
reset_mask[:, None, None],
|
||
torch.zeros_like(updated_state),
|
||
updated_state,
|
||
)
|
||
runtime.reset_mask = None
|
||
runtime.state = updated_state
|
||
if (
|
||
runtime.memory.config.raw_token_write
|
||
and not runtime.read_enabled
|
||
and runtime.input_ids is not None
|
||
):
|
||
if runtime.output_embeddings is None:
|
||
raise RuntimeError("raw_token_write requires output embeddings")
|
||
offset = runtime.memory.config.write_token_offset or 1
|
||
if runtime.attention_mask is not None and runtime.attention_mask.ndim == 2:
|
||
last_valid = runtime.attention_mask.sum(dim=-1).to(torch.long) - 1
|
||
positions = (last_valid - (offset - 1)).clamp_min(0)
|
||
token_ids = runtime.input_ids.gather(1, positions[:, None]).squeeze(1)
|
||
else:
|
||
token_ids = runtime.input_ids[:, -offset]
|
||
# Store the output-projection row, not the input embedding row.
|
||
# Qwen checkpoints are allowed to untie these matrices, and the
|
||
# generation-time pointer is consumed through output_embeddings.
|
||
runtime.raw_memory = runtime.output_embeddings.weight[token_ids].detach()
|
||
return output
|
||
|
||
|
||
class QwenDynamicMemoryOutput:
|
||
"""Proxy for a normal HF output with one added field: ``memory``."""
|
||
|
||
def __init__(self, base_output: Any, memory: Optional[Tensor]) -> None:
|
||
self.base_output = base_output
|
||
self.memory = memory
|
||
|
||
def __getattr__(self, name: str) -> Any:
|
||
return getattr(self.base_output, name)
|
||
|
||
def __getitem__(self, key: Any) -> Any:
|
||
return self.base_output[key]
|
||
|
||
|
||
class QwenDynamicMemoryModel(nn.Module):
|
||
"""Qwen3.5 with caller-owned persistent dynamic memory state."""
|
||
|
||
def __init__(
|
||
self,
|
||
base_model: nn.Module,
|
||
memory_config: Optional[QwenMemoryConfig] = None,
|
||
*,
|
||
freeze_backbone: bool = True,
|
||
) -> None:
|
||
super().__init__()
|
||
self.base_model = base_model
|
||
text_config = base_model.config.text_config
|
||
self.memory_config = memory_config or QwenMemoryConfig()
|
||
self.layer_indices = self.memory_config.resolved_layers(text_config.num_hidden_layers)
|
||
memory_type = NativeQwenDynamicMemory if self.memory_config.native_mode else QwenDynamicMemory
|
||
self.memory = memory_type(text_config.hidden_size, self.memory_config)
|
||
self.text_retriever = (
|
||
NaturalLanguageRetriever(text_config.hidden_size)
|
||
if self.memory_config.natural_language_memory
|
||
else None
|
||
)
|
||
policy_type = (
|
||
AutomaticMemoryPolicyV2
|
||
if self.memory_config.automatic_memory_policy_version >= 2
|
||
else AutomaticMemoryPolicy
|
||
)
|
||
self.memory_policy = (
|
||
policy_type(text_config.hidden_size)
|
||
if self.memory_config.natural_language_memory and self.memory_config.automatic_memory
|
||
else None
|
||
)
|
||
self.memory_router_v2: Optional[MemoryRouterV2] = None
|
||
self.memory_os_v2: Optional[MemoryOSV2] = None
|
||
self._memory_router_v2_ready = False
|
||
if self.memory_config.hierarchical_memory or self.memory_config.memory_version >= 2:
|
||
self.memory_router_v2 = MemoryRouterV2(
|
||
text_config.hidden_size,
|
||
router_dim=self.memory_config.memory_router_dim,
|
||
num_heads=self.memory_config.memory_router_heads,
|
||
max_hops=self.memory_config.memory_max_hops,
|
||
)
|
||
self.memory_os_v2 = self._new_memory_os_v2(text_config.hidden_size)
|
||
self._text_retriever_ready = False
|
||
self._memory_policy_ready = False
|
||
self.runtime = _MemoryRuntime(self.memory)
|
||
self.runtime.output_embeddings = self.base_model.get_output_embeddings()
|
||
self._persistent_memory: Optional[Tensor] = None
|
||
self.register_buffer("persistent_memory", torch.empty(0), persistent=True)
|
||
self.register_buffer(
|
||
"persistent_text_token_ids",
|
||
torch.empty(0, dtype=torch.long),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_token_mask",
|
||
torch.empty(0, dtype=torch.bool),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_slot_valid",
|
||
torch.empty(0, dtype=torch.bool),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_slot_keys",
|
||
torch.empty(0),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_slot_age",
|
||
torch.empty(0, dtype=torch.long),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_write_counter",
|
||
torch.empty(0, dtype=torch.long),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_key_token_ids",
|
||
torch.empty(0, dtype=torch.long),
|
||
persistent=True,
|
||
)
|
||
self.register_buffer(
|
||
"persistent_text_key_token_mask",
|
||
torch.empty(0, dtype=torch.bool),
|
||
persistent=True,
|
||
)
|
||
|
||
if freeze_backbone:
|
||
for parameter in self.base_model.parameters():
|
||
parameter.requires_grad_(False)
|
||
|
||
language_model = self.base_model.model.language_model
|
||
self._memory_adapters: list[MemoryLayerAdapter] = []
|
||
for layer_index in self.layer_indices:
|
||
layer = language_model.layers[layer_index]
|
||
adapter = MemoryLayerAdapter(
|
||
layer,
|
||
self.runtime,
|
||
read=True,
|
||
write=layer_index == self.layer_indices[-1],
|
||
mode=self.memory_config.mode,
|
||
blend_init=self.memory_config.blend_init,
|
||
)
|
||
language_model.layers[layer_index] = adapter
|
||
self._memory_adapters.append(adapter)
|
||
|
||
self.memory.to(self._find_layer_device())
|
||
if self.memory_router_v2 is not None:
|
||
self.memory_router_v2.to(self._find_layer_device())
|
||
if self.text_retriever is not None:
|
||
self.text_retriever.to(self._find_layer_device())
|
||
if self.memory_policy is not None:
|
||
self.memory_policy.to(self._find_layer_device())
|
||
if self.memory_config.persistent_memory:
|
||
self.persistent_memory = self.memory.initial_state(
|
||
1,
|
||
device=self._find_layer_device(),
|
||
)
|
||
self._persistent_memory = self.persistent_memory.detach()
|
||
if hasattr(self.base_model.config, "use_cache"):
|
||
self.base_model.config.use_cache = False
|
||
if hasattr(self.base_model.config, "text_config"):
|
||
self.base_model.config.text_config.use_cache = False
|
||
|
||
def _new_memory_os_v2(self, hidden_size: int) -> MemoryOSV2:
|
||
"""Create a V2 store from the configured page and candidate limits."""
|
||
|
||
if self.memory_router_v2 is None:
|
||
raise RuntimeError("cannot create V2 memory without a router")
|
||
tier_store = None
|
||
if self.memory_config.memory_storage_mode == "tiered":
|
||
if not self.memory_config.memory_storage_path:
|
||
raise ValueError("tiered memory requires memory_storage_path")
|
||
tier_store = TieredMemoryStoreV2(
|
||
self.memory_config.memory_storage_path,
|
||
key_dim=self.memory_config.memory_router_dim,
|
||
page_capacity=self.memory_config.memory_page_capacity,
|
||
)
|
||
bank = PagedMemoryBankV2(
|
||
hidden_size,
|
||
page_capacity=self.memory_config.memory_page_capacity,
|
||
max_pages=self.memory_config.memory_max_pages,
|
||
hot_pages=self.memory_config.memory_hot_pages,
|
||
top_k_pages=self.memory_config.memory_top_k_pages,
|
||
top_k_records=self.memory_config.memory_top_k_records,
|
||
max_hops=self.memory_config.memory_max_hops,
|
||
router=self.memory_router_v2,
|
||
key_dim=self.memory_config.memory_router_dim,
|
||
coarse_index_bits=self.memory_config.memory_coarse_index_bits,
|
||
tier_store=tier_store,
|
||
max_resident_pages=self.memory_config.memory_resident_pages,
|
||
runtime_device=self._find_layer_device(),
|
||
gpu_cache_records=self.memory_config.memory_gpu_cache_records,
|
||
gpu_cache_tokens=self.memory_config.memory_gpu_cache_tokens,
|
||
gpu_cache_reserve_mb=self.memory_config.memory_gpu_cache_reserve_mb,
|
||
gpu_cache_adaptive=self.memory_config.memory_gpu_cache_adaptive,
|
||
)
|
||
return MemoryOSV2(
|
||
hidden_size,
|
||
router=self.memory_router_v2,
|
||
bank=bank,
|
||
kv_budget=KVBudgetManagerV2(
|
||
max_tokens=self.memory_config.kv_budget_tokens,
|
||
hard_max_tokens=self.memory_config.kv_hard_max_tokens,
|
||
compaction_trigger=self.memory_config.kv_compaction_trigger,
|
||
keep_recent_tokens=self.memory_config.kv_keep_recent_tokens,
|
||
),
|
||
read_threshold=self.memory_config.memory_v2_read_threshold,
|
||
write_threshold=self.memory_config.memory_v2_write_threshold,
|
||
)
|
||
|
||
def _ensure_text_memory(self, batch_size: int, *, device: torch.device) -> None:
|
||
"""Create the fixed-size model-owned natural-language memory bank."""
|
||
|
||
if not self.memory_config.natural_language_memory:
|
||
return
|
||
shape = (
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_tokens,
|
||
)
|
||
if tuple(self.persistent_text_token_ids.shape) != shape:
|
||
self.persistent_text_token_ids = torch.zeros(
|
||
shape,
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
self.persistent_text_token_mask = torch.zeros(
|
||
shape,
|
||
dtype=torch.bool,
|
||
device=device,
|
||
)
|
||
self.persistent_text_slot_valid = torch.zeros(
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
dtype=torch.bool,
|
||
device=device,
|
||
)
|
||
hidden_size = self.memory.hidden_size
|
||
self.persistent_text_slot_keys = torch.zeros(
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
hidden_size,
|
||
dtype=torch.float32,
|
||
device=device,
|
||
)
|
||
self.persistent_text_slot_age = torch.full(
|
||
(batch_size, self.memory_config.memory_slots),
|
||
-1,
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
self.persistent_text_write_counter = torch.zeros(
|
||
batch_size,
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
key_shape = (
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_key_tokens,
|
||
)
|
||
self.persistent_text_key_token_ids = torch.zeros(
|
||
key_shape,
|
||
dtype=torch.long,
|
||
device=device,
|
||
)
|
||
self.persistent_text_key_token_mask = torch.zeros(
|
||
key_shape,
|
||
dtype=torch.bool,
|
||
device=device,
|
||
)
|
||
|
||
def _bind_text_memory(self, batch_size: int, *, device: torch.device) -> None:
|
||
"""Bind the current text bank to the per-call runtime."""
|
||
|
||
if not self.memory_config.natural_language_memory:
|
||
self.runtime.text_token_ids = None
|
||
self.runtime.text_token_mask = None
|
||
self.runtime.text_slot_valid = None
|
||
return
|
||
if self.memory_config.persistent_memory and self.runtime.use_persistent_state:
|
||
self._ensure_text_memory(1, device=device)
|
||
else:
|
||
self._ensure_text_memory(batch_size, device=device)
|
||
if not (self.memory_config.persistent_memory and self.runtime.use_persistent_state):
|
||
shape = (
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_tokens,
|
||
)
|
||
key_shape = (
|
||
batch_size,
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_key_tokens,
|
||
)
|
||
if (
|
||
self.runtime.text_token_ids is None
|
||
or tuple(self.runtime.text_token_ids.shape) != shape
|
||
or self.runtime.text_key_token_ids is None
|
||
or tuple(self.runtime.text_key_token_ids.shape) != key_shape
|
||
):
|
||
self.runtime.text_token_ids = torch.zeros(shape, dtype=torch.long, device=device)
|
||
self.runtime.text_token_mask = torch.zeros(shape, dtype=torch.bool, device=device)
|
||
self.runtime.text_slot_valid = torch.zeros(
|
||
batch_size, self.memory_config.memory_slots, dtype=torch.bool, device=device
|
||
)
|
||
self.runtime.text_slot_keys = torch.zeros(
|
||
batch_size, self.memory_config.memory_slots, self.memory.hidden_size,
|
||
dtype=torch.float32, device=device
|
||
)
|
||
self.runtime.text_slot_age = torch.full(
|
||
(batch_size, self.memory_config.memory_slots), -1, dtype=torch.long, device=device
|
||
)
|
||
self.runtime.text_write_counter = torch.zeros(batch_size, dtype=torch.long, device=device)
|
||
self.runtime.text_key_token_ids = torch.zeros(key_shape, dtype=torch.long, device=device)
|
||
self.runtime.text_key_token_mask = torch.zeros(key_shape, dtype=torch.bool, device=device)
|
||
return
|
||
if self.persistent_text_token_ids.shape[0] == batch_size:
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask
|
||
return
|
||
if self.persistent_text_token_ids.shape[0] == 1:
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids.expand(
|
||
batch_size, -1, -1
|
||
).clone()
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask.expand(
|
||
batch_size, -1, -1
|
||
).clone()
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid.expand(
|
||
batch_size, -1
|
||
).clone()
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys.expand(
|
||
batch_size, -1, -1
|
||
).clone()
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age.expand(
|
||
batch_size, -1
|
||
).clone()
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter.expand(
|
||
batch_size
|
||
).clone()
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids.expand(
|
||
batch_size, -1, -1
|
||
).clone()
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask.expand(
|
||
batch_size, -1, -1
|
||
).clone()
|
||
return
|
||
raise ValueError(
|
||
"natural-language memory batch does not match the loaded persistent bank"
|
||
)
|
||
|
||
def _clear_text_memory(self, mask: Optional[Tensor] = None) -> None:
|
||
if self.runtime.text_token_ids is None:
|
||
return
|
||
if mask is None:
|
||
self.runtime.text_token_ids.zero_()
|
||
self.runtime.text_token_mask.zero_()
|
||
self.runtime.text_slot_valid.zero_()
|
||
self.runtime.text_slot_keys.zero_()
|
||
self.runtime.text_slot_age.fill_(-1)
|
||
self.runtime.text_write_counter.zero_()
|
||
self.runtime.text_key_token_ids.zero_()
|
||
self.runtime.text_key_token_mask.zero_()
|
||
self.runtime.text_last_written_slot = None
|
||
self.runtime.text_prefix_used = False
|
||
return
|
||
mask = mask.to(device=self.runtime.text_token_ids.device, dtype=torch.bool)
|
||
self.runtime.text_token_ids[mask] = 0
|
||
self.runtime.text_token_mask[mask] = False
|
||
self.runtime.text_slot_valid[mask] = False
|
||
self.runtime.text_slot_keys[mask] = 0
|
||
self.runtime.text_slot_age[mask] = -1
|
||
self.runtime.text_write_counter[mask] = 0
|
||
self.runtime.text_key_token_ids[mask] = 0
|
||
self.runtime.text_key_token_mask[mask] = False
|
||
self.runtime.text_last_written_slot = None
|
||
|
||
@torch.no_grad()
|
||
def _encode_model_key(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
|
||
"""Encode a memory key with the frozen original Qwen representation."""
|
||
|
||
previous_read = self.runtime.read_enabled
|
||
previous_update = self.runtime.update_enabled
|
||
self.runtime.read_enabled = False
|
||
self.runtime.update_enabled = False
|
||
try:
|
||
embedding_layer = self.base_model.get_input_embeddings()
|
||
input_ids = input_ids.to(embedding_layer.weight.device)
|
||
attention_mask = attention_mask.to(input_ids.device)
|
||
output = self.base_model(
|
||
input_ids=input_ids,
|
||
attention_mask=attention_mask,
|
||
use_cache=False,
|
||
output_hidden_states=True,
|
||
)
|
||
hidden = output.hidden_states[-1].float()
|
||
weights = attention_mask.to(device=hidden.device, dtype=hidden.dtype).unsqueeze(-1)
|
||
key = (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1.0)
|
||
return F.normalize(key, dim=-1)
|
||
finally:
|
||
self.runtime.read_enabled = previous_read
|
||
self.runtime.update_enabled = previous_update
|
||
|
||
@torch.no_grad()
|
||
def compact_context_for_kv(
|
||
self,
|
||
input_ids: Tensor,
|
||
attention_mask: Optional[Tensor] = None,
|
||
*,
|
||
archive: bool = True,
|
||
chunk_tokens: Optional[int] = None,
|
||
) -> tuple[Tensor, Tensor, dict[str, Any]]:
|
||
"""Archive the old prefix and return a bounded hot-context window.
|
||
|
||
This is the model-side boundary between the exact working KV and the
|
||
lossy long-term memory path. It never sends the current token to all
|
||
pages: each archived chunk gets one address, and future reads still
|
||
use the V2 coarse-page -> exact-record -> Top-K route.
|
||
|
||
The method accepts already-tokenized input because it is also used by
|
||
generation. The archived token sequence is kept losslessly inside a
|
||
V2 record; its semantic address is produced by the frozen Qwen hidden
|
||
representation. A caller may disable ``archive`` for measurement,
|
||
in which case the method only reports the retention plan.
|
||
"""
|
||
|
||
if input_ids.ndim != 2:
|
||
raise ValueError("input_ids must have shape [batch, sequence]")
|
||
if attention_mask is None:
|
||
attention_mask = torch.ones_like(input_ids)
|
||
if attention_mask.shape != input_ids.shape:
|
||
raise ValueError("attention_mask must have the same shape as input_ids")
|
||
|
||
budget = self.memory_os_v2.kv_budget if self.memory_os_v2 is not None else None
|
||
if (
|
||
self.memory_os_v2 is None
|
||
or not self.memory_config.auto_compact_context
|
||
or budget is None
|
||
):
|
||
return input_ids, attention_mask, {
|
||
"compacted": False,
|
||
"reason": "disabled",
|
||
"archived_records": 0,
|
||
"original_tokens": [int(item) for item in attention_mask.sum(dim=-1).tolist()],
|
||
"retained_tokens": [int(item) for item in attention_mask.sum(dim=-1).tolist()],
|
||
}
|
||
|
||
chunk_size = int(chunk_tokens or self.memory_config.context_chunk_tokens)
|
||
if chunk_size < 1:
|
||
raise ValueError("chunk_tokens must be positive")
|
||
hot_limit = int(budget.max_tokens)
|
||
pending: list[dict[str, Any]] = []
|
||
pending_previous: list[Optional[int]] = []
|
||
retained: list[Tensor] = []
|
||
original_lengths: list[int] = []
|
||
retained_lengths: list[int] = []
|
||
|
||
for batch_index in range(input_ids.shape[0]):
|
||
valid = input_ids[batch_index][attention_mask[batch_index].to(dtype=torch.bool)]
|
||
original_lengths.append(int(valid.numel()))
|
||
if valid.numel() <= hot_limit:
|
||
retained.append(valid)
|
||
retained_lengths.append(int(valid.numel()))
|
||
continue
|
||
|
||
archived_prefix = valid[:-hot_limit]
|
||
retained_window = valid[-hot_limit:]
|
||
retained.append(retained_window)
|
||
retained_lengths.append(int(retained_window.numel()))
|
||
previous_index: Optional[int] = None
|
||
for start in range(0, int(archived_prefix.numel()), chunk_size):
|
||
end = min(start + chunk_size, int(archived_prefix.numel()))
|
||
if len(pending) >= self.memory_config.context_archive_max_records:
|
||
raise RuntimeError(
|
||
"context archive record limit reached; increase "
|
||
"context_archive_max_records before compacting this input"
|
||
)
|
||
chunk = archived_prefix[start:end].detach()
|
||
chunk_mask = torch.ones_like(chunk, dtype=torch.bool)
|
||
key = self._encode_model_key(chunk.unsqueeze(0), chunk_mask.unsqueeze(0))[0]
|
||
pending.append(
|
||
{
|
||
"text": f"context_chunk:{batch_index}:{start}:{end}",
|
||
"key": key,
|
||
"summary": key,
|
||
"memory_type": "context_chunk",
|
||
# Context chunks are ordered evidence, not competing
|
||
# values of one fact. Leaving the fact-conflict
|
||
# fields empty prevents ordinary version resolution
|
||
# from superseding every earlier chunk in a document.
|
||
"entity": "",
|
||
"attribute": "",
|
||
"value": "",
|
||
"importance": 0.55,
|
||
"confidence": 0.80,
|
||
"source": "kv_compaction",
|
||
"evidence": [
|
||
f"batch:{batch_index}",
|
||
f"token_range:{start}:{end}",
|
||
f"token_count:{end - start}",
|
||
],
|
||
"token_ids": chunk,
|
||
"token_mask": chunk_mask,
|
||
"trusted": True,
|
||
"force": True,
|
||
}
|
||
)
|
||
pending_previous.append(previous_index)
|
||
previous_index = len(pending) - 1
|
||
|
||
results: list[tuple[Any, str]] = []
|
||
if pending and archive:
|
||
results = self.memory_os_v2.write_batch(pending)
|
||
# Add a forward chain after the atomic insert. This gives the
|
||
# multi-hop reader a deterministic path through adjacent chunks
|
||
# without making page allocation depend on future record IDs.
|
||
for result_index, previous_index in enumerate(pending_previous):
|
||
if previous_index is None:
|
||
continue
|
||
record = results[result_index][0]
|
||
previous_record = results[previous_index][0]
|
||
if previous_record.record_id not in record.related_ids:
|
||
record.related_ids.append(previous_record.record_id)
|
||
self.memory_os_v2.bank._store_record(record)
|
||
|
||
max_retained = max(retained_lengths, default=0)
|
||
compacted_ids = torch.zeros(
|
||
(input_ids.shape[0], max_retained), dtype=input_ids.dtype, device=input_ids.device
|
||
)
|
||
compacted_mask = torch.zeros(
|
||
(input_ids.shape[0], max_retained), dtype=attention_mask.dtype, device=input_ids.device
|
||
)
|
||
for batch_index, values in enumerate(retained):
|
||
if values.numel() == 0:
|
||
continue
|
||
compacted_ids[batch_index, : values.numel()] = values.to(input_ids.device)
|
||
compacted_mask[batch_index, : values.numel()] = 1
|
||
return compacted_ids, compacted_mask, {
|
||
"compacted": any(original > retained for original, retained in zip(original_lengths, retained_lengths)),
|
||
"archived_records": len(results),
|
||
"archived_tokens": sum(original - retained for original, retained in zip(original_lengths, retained_lengths)),
|
||
"original_tokens": original_lengths,
|
||
"retained_tokens": retained_lengths,
|
||
"hot_limit": hot_limit,
|
||
"archive_requested": bool(archive),
|
||
}
|
||
|
||
@torch.no_grad()
|
||
def _rank_v2_text_matches(
|
||
self,
|
||
query_key: Tensor,
|
||
query_token_ids: Tensor,
|
||
) -> list[tuple[float, int, Any]]:
|
||
"""Rank active V2 text records without materializing their payloads."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
return []
|
||
active = [
|
||
record
|
||
for record in self.memory_os_v2.bank.records.values()
|
||
if record.status == "active"
|
||
]
|
||
if not active:
|
||
return []
|
||
keys = torch.stack([record.key for record in active], dim=0)
|
||
scores = self.memory_os_v2.bank._score_candidates(query_key, keys).reshape(-1).cpu()
|
||
query_unique = torch.unique(query_token_ids.detach().reshape(-1).cpu())
|
||
ranked: list[tuple[float, int, Any]] = []
|
||
for index, record in enumerate(active):
|
||
shared = 0
|
||
if record.token_ids is not None and query_unique.numel() > 0:
|
||
own = torch.unique(record.token_ids.detach().reshape(-1).cpu())
|
||
shared = int(torch.isin(query_unique, own).sum().item())
|
||
ranked.append((float(scores[index].item()), shared, record))
|
||
ranked.sort(key=lambda item: (item[0], item[1], item[2].last_access), reverse=True)
|
||
return ranked
|
||
|
||
@torch.no_grad()
|
||
def _forget_text_memory_by_key(
|
||
self,
|
||
key_input_ids: Tensor,
|
||
key_attention_mask: Tensor,
|
||
) -> Tensor:
|
||
"""Erase the addressed natural-language record for a learned forget turn."""
|
||
|
||
erased = torch.zeros(
|
||
key_input_ids.shape[0], dtype=torch.bool, device=key_input_ids.device
|
||
)
|
||
if (
|
||
self.runtime.text_slot_valid is None
|
||
or self.runtime.text_key_token_ids is None
|
||
or self.runtime.text_key_token_mask is None
|
||
or not bool(self.runtime.text_slot_valid.any())
|
||
):
|
||
return erased
|
||
query_keys = self._encode_model_key(key_input_ids, key_attention_mask)
|
||
for batch_index in range(key_input_ids.shape[0]):
|
||
valid_slots = self.runtime.text_slot_valid[batch_index]
|
||
if not bool(valid_slots.any()):
|
||
continue
|
||
query_ids = key_input_ids[batch_index][key_attention_mask[batch_index].bool()]
|
||
query_ids = query_ids[: self.memory_config.text_memory_key_tokens]
|
||
old_ids = self.runtime.text_key_token_ids[batch_index]
|
||
old_mask = self.runtime.text_key_token_mask[batch_index]
|
||
shared = (
|
||
(old_ids[:, :, None] == query_ids[None, None, :])
|
||
& old_mask[:, :, None]
|
||
& valid_slots[:, None, None]
|
||
).any(dim=-1).sum(dim=-1).float()
|
||
old_lengths = old_mask.sum(dim=-1).float()
|
||
overlap = shared / torch.sqrt(
|
||
old_lengths * float(max(1, query_ids.numel()))
|
||
).clamp_min(1.0)
|
||
overlap = overlap.masked_fill(~valid_slots, -1.0)
|
||
best_overlap, best_slot = overlap.max(dim=-1)
|
||
best_score = best_overlap
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
slot_keys = F.normalize(
|
||
self.runtime.text_slot_keys[batch_index].to(query_keys.device, dtype=query_keys.dtype),
|
||
dim=-1,
|
||
)
|
||
learned = torch.sigmoid(
|
||
self.text_retriever(query_keys[batch_index].unsqueeze(0), slot_keys.unsqueeze(0))[0]
|
||
).masked_fill(~valid_slots.to(query_keys.device), -1.0)
|
||
learned_score, learned_slot = learned.max(dim=-1)
|
||
if float(learned_score) >= float(best_score):
|
||
best_score = learned_score
|
||
best_slot = learned_slot
|
||
if float(best_score) < self.memory_config.text_memory_overlap_threshold and not (
|
||
self.text_retriever is not None
|
||
and self._text_retriever_ready
|
||
and float(best_score) >= 0.65
|
||
):
|
||
continue
|
||
slot = int(best_slot.item())
|
||
self._clear_legacy_memory_slot(slot)
|
||
if self.memory_os_v2 is not None and self.runtime.use_persistent_state:
|
||
ranked = self._rank_v2_text_matches(
|
||
query_keys[batch_index],
|
||
query_ids,
|
||
)
|
||
for score, shared, record in ranked:
|
||
# The hot-bank matcher has already located the target.
|
||
# V2 records created as independent fragments have no
|
||
# slot index, so use the frozen-Qwen address plus token
|
||
# evidence to retract only the matching old record.
|
||
same_hot_slot = record.slot_index == slot
|
||
if same_hot_slot or (score >= 0.60 and shared >= 2):
|
||
self.memory_os_v2.retract_record(record.record_id)
|
||
if not same_hot_slot and score < 0.78:
|
||
break
|
||
erased[batch_index] = True
|
||
return erased
|
||
|
||
def _write_text_memory(
|
||
self,
|
||
input_ids: Optional[Tensor],
|
||
attention_mask: Optional[Tensor],
|
||
text_input_ids: Optional[Tensor] = None,
|
||
text_attention_mask: Optional[Tensor] = None,
|
||
key_input_ids: Optional[Tensor] = None,
|
||
key_attention_mask: Optional[Tensor] = None,
|
||
storage_input_ids: Optional[Tensor] = None,
|
||
storage_attention_mask: Optional[Tensor] = None,
|
||
force_write: bool = False,
|
||
memory_text: Optional[str] = None,
|
||
) -> None:
|
||
"""Commit one conversational write into the selected internal slot.
|
||
|
||
The learned controller decides whether a turn is a durable fact and
|
||
which slot it belongs to. The exact token sequence is then copied into
|
||
the model-owned bank so multi-token names, values and punctuation are
|
||
not lossy-compressed into a single output token.
|
||
"""
|
||
|
||
# The bank is consumed as an internal prefix before the next Qwen
|
||
# chat turn. Store the valid system-message form for generation;
|
||
# keep the exact user fact separately in ``text_key_token_ids`` for
|
||
# duplicate detection and retrieval overlap. A raw fact placed
|
||
# before ``<|im_start|>user`` is not a reliable conditioning channel.
|
||
storage_ids = (
|
||
text_input_ids
|
||
if text_input_ids is not None
|
||
else (storage_input_ids if storage_input_ids is not None else input_ids)
|
||
)
|
||
storage_mask = (
|
||
text_attention_mask
|
||
if text_input_ids is not None
|
||
else (storage_attention_mask if storage_input_ids is not None else attention_mask)
|
||
)
|
||
key_ids = key_input_ids if key_input_ids is not None else storage_ids
|
||
key_mask = key_attention_mask if key_input_ids is not None else storage_mask
|
||
if (
|
||
not self.memory_config.natural_language_memory
|
||
or storage_ids is None
|
||
or self.runtime.text_token_ids is None
|
||
or self.runtime.text_slot_keys is None
|
||
or self.runtime.text_key_token_ids is None
|
||
):
|
||
return
|
||
self.runtime.text_last_written_slot = torch.full(
|
||
(storage_ids.shape[0],),
|
||
-1,
|
||
dtype=torch.long,
|
||
device=self.runtime.text_slot_valid.device,
|
||
)
|
||
address = self.memory.last_write_address
|
||
write_probability = self.memory.last_write_probability
|
||
if address is None or key_ids is None:
|
||
return
|
||
if write_probability is None:
|
||
write_probability = address.max(dim=-1, keepdim=True).values
|
||
if force_write:
|
||
should_write = torch.ones(
|
||
storage_ids.shape[0], dtype=torch.bool, device=storage_ids.device
|
||
)
|
||
elif (
|
||
self.memory_config.automatic_memory
|
||
and self._memory_policy_ready
|
||
and self.runtime.auto_memory_probability is not None
|
||
):
|
||
# High recall is intentional: the policy is trained to recognize
|
||
# durable user information. It is the final automatic decision;
|
||
# the older native write gate is not OR-ed here because it was
|
||
# trained for continuous-state updates and can be over-eager on
|
||
# question-shaped inputs.
|
||
should_write = self.runtime.auto_memory_probability >= self.memory_config.auto_memory_threshold
|
||
else:
|
||
should_write = write_probability.squeeze(-1) >= self.memory_config.text_memory_write_threshold
|
||
forget_probability = self.runtime.auto_memory_forget_probability
|
||
if (
|
||
not force_write
|
||
and forget_probability is not None
|
||
and bool((forget_probability >= self.memory_config.auto_forget_threshold).any())
|
||
):
|
||
forget_mask = forget_probability >= self.memory_config.auto_forget_threshold
|
||
forget_key_mask = (
|
||
key_mask
|
||
if key_mask is not None
|
||
else torch.ones_like(key_ids, dtype=torch.bool)
|
||
)
|
||
erased = self._forget_text_memory_by_key(key_ids, forget_key_mask)
|
||
should_write = should_write & ~(forget_mask & erased)
|
||
if not force_write and memory_text:
|
||
if looks_like_question(memory_text):
|
||
should_write = torch.zeros_like(should_write, dtype=torch.bool)
|
||
if storage_mask is None:
|
||
storage_mask = torch.ones_like(storage_ids, dtype=torch.bool)
|
||
else:
|
||
storage_mask = storage_mask.to(device=storage_ids.device, dtype=torch.bool)
|
||
if key_mask is None:
|
||
key_mask = torch.ones_like(key_ids, dtype=torch.bool)
|
||
else:
|
||
key_mask = key_mask.to(device=key_ids.device, dtype=torch.bool)
|
||
text_keys = self._encode_model_key(key_ids, key_mask)
|
||
max_tokens = self.memory_config.text_memory_tokens
|
||
for batch_index in range(storage_ids.shape[0]):
|
||
if not bool(should_write[batch_index]):
|
||
continue
|
||
valid_ids = storage_ids[batch_index][storage_mask[batch_index]]
|
||
if valid_ids.numel() == 0:
|
||
continue
|
||
valid_key_ids = key_ids[batch_index][key_mask[batch_index]]
|
||
valid_key_ids = valid_key_ids[: self.memory_config.text_memory_key_tokens].detach().to(
|
||
device=self.runtime.text_key_token_ids.device,
|
||
dtype=torch.long,
|
||
)
|
||
valid_ids = valid_ids[:max_tokens].detach().to(
|
||
device=self.runtime.text_token_ids.device,
|
||
dtype=torch.long,
|
||
)
|
||
key = text_keys[batch_index]
|
||
valid_slots = self.runtime.text_slot_valid[batch_index]
|
||
if bool(valid_slots.any()):
|
||
old_key_ids = self.runtime.text_key_token_ids[batch_index]
|
||
old_key_mask = self.runtime.text_key_token_mask[batch_index]
|
||
equal = old_key_ids[:, :, None] == valid_key_ids[None, None, :]
|
||
overlap_count = (
|
||
equal
|
||
& old_key_mask[:, :, None]
|
||
& torch.ones(
|
||
1,
|
||
1,
|
||
valid_key_ids.numel(),
|
||
dtype=torch.bool,
|
||
device=old_key_ids.device,
|
||
)
|
||
).any(dim=-1).sum(dim=-1).float()
|
||
old_lengths = old_key_mask.sum(dim=-1).float()
|
||
similarities = overlap_count / torch.sqrt(
|
||
old_lengths * float(max(1, valid_key_ids.numel()))
|
||
).clamp_min(1.0)
|
||
masked_similarities = similarities.masked_fill(
|
||
~valid_slots,
|
||
torch.finfo(similarities.dtype).min,
|
||
)
|
||
best_similarity, best_slot = masked_similarities.max(dim=-1)
|
||
# Exact token equality is a stronger duplicate signal than
|
||
# any learned semantic score. It makes repeated natural
|
||
# language facts idempotent, even when the retriever has not
|
||
# seen that exact name/value during training.
|
||
key_length = valid_key_ids.numel()
|
||
exact_match = (
|
||
(old_key_ids[:, :key_length] == valid_key_ids[None, :])
|
||
| (~old_key_mask[:, :key_length])
|
||
).all(dim=-1) & valid_slots & (
|
||
old_key_mask.sum(dim=-1) == valid_key_ids.numel()
|
||
)
|
||
exact_slots = exact_match.nonzero(as_tuple=False).flatten()
|
||
if exact_slots.numel() > 0:
|
||
best_similarity = torch.tensor(1.0, device=key.device)
|
||
best_slot = exact_slots[0]
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
if exact_slots.numel() == 0:
|
||
learned_similarity = torch.sigmoid(
|
||
self.text_retriever(
|
||
key.unsqueeze(0),
|
||
self.runtime.text_slot_keys[batch_index].unsqueeze(0),
|
||
)[0]
|
||
)
|
||
learned_similarity = learned_similarity.masked_fill(
|
||
~valid_slots,
|
||
torch.finfo(learned_similarity.dtype).min,
|
||
)
|
||
learned_best, learned_slot = learned_similarity.max(dim=-1)
|
||
# Shared words such as "我" and "请记住" are common in
|
||
# many facts. Once the retriever is trained, use its
|
||
# semantic same-attribute score for replacement so a new
|
||
# fruit fact cannot evict a work-location fact merely
|
||
# because both contain the same pronouns.
|
||
if float(learned_best) >= 0.65:
|
||
best_similarity = learned_best
|
||
best_slot = learned_slot
|
||
# Keep the exact token-overlap score when the learned
|
||
# retriever is uncertain. Overwriting it with -1
|
||
# made obvious same-attribute updates look unrelated
|
||
# merely because the paraphrase was outside the
|
||
# retriever's training distribution.
|
||
else:
|
||
best_similarity = torch.tensor(-1.0, device=key.device)
|
||
best_slot = torch.tensor(0, dtype=torch.long, device=key.device)
|
||
if float(best_similarity) >= self.memory_config.text_memory_replace_threshold:
|
||
slot = int(best_slot.item())
|
||
else:
|
||
free_slots = (~valid_slots).nonzero(as_tuple=False).flatten()
|
||
if free_slots.numel() > 0:
|
||
slot = int(free_slots[0].item())
|
||
else:
|
||
slot = int(self.runtime.text_slot_age[batch_index].argmin().item())
|
||
self.runtime.text_token_ids[batch_index, slot].zero_()
|
||
self.runtime.text_token_mask[batch_index, slot].zero_()
|
||
self.runtime.text_token_ids[batch_index, slot, : valid_ids.numel()].copy_(valid_ids)
|
||
self.runtime.text_token_mask[batch_index, slot, : valid_ids.numel()] = True
|
||
self.runtime.text_slot_valid[batch_index, slot] = True
|
||
self.runtime.text_slot_keys[batch_index, slot].copy_(key.to(self.runtime.text_slot_keys.device))
|
||
self.runtime.text_key_token_ids[batch_index, slot].zero_()
|
||
self.runtime.text_key_token_mask[batch_index, slot].zero_()
|
||
self.runtime.text_key_token_ids[batch_index, slot, : valid_key_ids.numel()].copy_(valid_key_ids)
|
||
self.runtime.text_key_token_mask[batch_index, slot, : valid_key_ids.numel()] = True
|
||
self.runtime.text_write_counter[batch_index] += 1
|
||
self.runtime.text_slot_age[batch_index, slot] = self.runtime.text_write_counter[batch_index]
|
||
self.runtime.text_last_written_slot[batch_index] = slot
|
||
if self.memory_os_v2 is not None and self.runtime.use_persistent_state:
|
||
confidence = float(write_probability[batch_index].reshape(-1)[0].item())
|
||
if force_write:
|
||
confidence = 1.0
|
||
# The 16-slot bank is only the hot cache. New, unrelated
|
||
# conversational fragments must remain independent V2
|
||
# records instead of superseding whatever happened to be in
|
||
# the same hot slot. Reuse a V2 slot index only when the
|
||
# hot-bank matcher already established that this is a
|
||
# duplicate/update of an existing fact.
|
||
v2_slot_index = (
|
||
slot
|
||
if float(best_similarity) >= self.memory_config.text_memory_replace_threshold
|
||
else -1
|
||
)
|
||
if v2_slot_index >= 0 and self.memory_os_v2 is not None:
|
||
# A semantic replacement should retire the matching
|
||
# active record, while unrelated fragments remain active
|
||
# even if they once occupied the same hot-cache slot.
|
||
for score, shared, record in self._rank_v2_text_matches(
|
||
key,
|
||
valid_key_ids,
|
||
):
|
||
if record.slot_index == slot or (score >= 0.60 and shared >= 2):
|
||
self.memory_os_v2.retract_record(record.record_id)
|
||
if record.slot_index != slot and score < 0.78:
|
||
break
|
||
self.memory_os_v2.write(
|
||
text=memory_text or f"memory_slot:{slot}",
|
||
key=key,
|
||
summary=key,
|
||
memory_type="episodic_text",
|
||
importance=max(confidence, 0.5),
|
||
confidence=confidence,
|
||
source="explicit" if force_write else "automatic",
|
||
evidence=[f"token_count:{int(valid_ids.numel())}"],
|
||
slot_index=v2_slot_index,
|
||
token_ids=valid_ids,
|
||
token_mask=torch.ones_like(valid_ids, dtype=torch.bool),
|
||
trusted=force_write or confidence >= self.memory_config.memory_v2_write_threshold,
|
||
force=force_write,
|
||
)
|
||
|
||
def _probe_text_retrieval(
|
||
self,
|
||
input_ids: Tensor,
|
||
attention_mask: Tensor,
|
||
) -> tuple[Optional[Tensor], Optional[Tensor]]:
|
||
"""Compute query-dependent similarity against internal text keys."""
|
||
|
||
if (
|
||
not self.memory_config.natural_language_memory
|
||
or self.runtime.text_slot_valid is None
|
||
or self.runtime.text_slot_keys is None
|
||
or self.runtime.text_key_token_ids is None
|
||
or self.runtime.text_key_token_mask is None
|
||
or not bool(self.runtime.text_slot_valid.any())
|
||
):
|
||
return None, None
|
||
query_key = self._encode_model_key(input_ids, attention_mask)
|
||
self.runtime.v2_query_key = query_key.detach()
|
||
slot_keys = F.normalize(
|
||
self.runtime.text_slot_keys.to(device=query_key.device, dtype=query_key.dtype),
|
||
dim=-1,
|
||
)
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
learned_logits = self.text_retriever(query_key, slot_keys)
|
||
dense_similarity = torch.sigmoid(learned_logits)
|
||
else:
|
||
dense_similarity = torch.einsum("bh,bsh->bs", query_key, slot_keys)
|
||
# Dense similarity alone is too permissive for short Chinese queries:
|
||
# common tokens such as "我" and "的" can make unrelated memories
|
||
# look relevant. Combine it with exact token overlap from the
|
||
# internally stored key, yielding a conservative hybrid score.
|
||
key_token_ids = self.runtime.text_key_token_ids.to(device=input_ids.device)
|
||
key_token_mask = self.runtime.text_key_token_mask.to(device=input_ids.device, dtype=torch.bool)
|
||
query_token_ids = input_ids.to(device=key_token_ids.device)
|
||
query_token_mask = attention_mask.to(device=key_token_ids.device, dtype=torch.bool)
|
||
equal = query_token_ids[:, None, :, None] == key_token_ids[:, :, None, :]
|
||
equal &= query_token_mask[:, None, :, None] & key_token_mask[:, :, None, :]
|
||
shared_tokens = equal.any(dim=-1).sum(dim=-1).float()
|
||
query_lengths = query_token_mask.sum(dim=-1, keepdim=True).float()
|
||
key_lengths = key_token_mask.sum(dim=-1).float()
|
||
overlap = shared_tokens / torch.sqrt(query_lengths * key_lengths).clamp_min(1.0)
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
scores = 0.85 * dense_similarity + 0.15 * overlap.clamp(0.0, 1.0)
|
||
else:
|
||
dense_score = (dense_similarity + 1.0).clamp(0.0, 2.0) * 0.5
|
||
scores = 0.4 * dense_score + 0.6 * overlap.clamp(0.0, 1.0)
|
||
self.runtime.text_read_overlap = overlap
|
||
scores = scores.masked_fill(
|
||
~self.runtime.text_slot_valid.to(device=scores.device),
|
||
torch.finfo(scores.dtype).min,
|
||
)
|
||
return scores, scores.max(dim=-1).values
|
||
|
||
def _build_text_prefix(
|
||
self,
|
||
input_ids: Tensor,
|
||
attention_mask: Tensor,
|
||
query_text: str = "",
|
||
) -> tuple[Optional[Tensor], Optional[Tensor], int]:
|
||
"""Retrieve text memory and build an internal prefix for Qwen."""
|
||
|
||
address, relevance = self._probe_text_retrieval(input_ids, attention_mask)
|
||
self.runtime.text_read_slots = None
|
||
self.runtime.v2_last_decisions = []
|
||
self.runtime.text_read_relevance = relevance
|
||
self.runtime.text_prefix_used = False
|
||
if (
|
||
self.memory_os_v2 is not None
|
||
and self._memory_router_v2_ready
|
||
and self.runtime.v2_query_key is None
|
||
):
|
||
self.runtime.v2_query_key = self._encode_model_key(input_ids, attention_mask).detach()
|
||
|
||
# V2 is the scalable address path. It first selects a bounded set of
|
||
# pages through the coarse index, then reranks only records in those
|
||
# pages. The legacy fixed bank remains a compatibility fallback for
|
||
# old checkpoints that do not contain a trained V2 router.
|
||
if (
|
||
self.memory_os_v2 is not None
|
||
and self._memory_router_v2_ready
|
||
and self.runtime.v2_query_key is not None
|
||
):
|
||
v2_prefix_parts: list[Tensor] = []
|
||
v2_prefix_masks: list[Tensor] = []
|
||
v2_slots: list[list[int]] = []
|
||
for batch_index in range(input_ids.shape[0]):
|
||
records, decision = self.memory_os_v2.read(
|
||
query_key=self.runtime.v2_query_key[batch_index],
|
||
query_text=query_text,
|
||
query_token_ids=input_ids[batch_index],
|
||
top_k_pages=self.memory_config.memory_top_k_pages,
|
||
top_k_records=self.memory_config.memory_top_k_records,
|
||
max_hops=self.memory_config.memory_max_hops,
|
||
)
|
||
self.runtime.v2_last_decisions.append(asdict(decision))
|
||
ids_parts: list[Tensor] = []
|
||
slot_list: list[int] = []
|
||
for record in records:
|
||
ids: Optional[Tensor] = None
|
||
if record.token_ids is not None:
|
||
cached_ids, cached_mask = self.memory_os_v2.bank.gpu_record_payload(record)
|
||
if cached_ids is not None:
|
||
if cached_mask is not None and cached_mask.numel() == cached_ids.numel():
|
||
ids = cached_ids[cached_mask]
|
||
else:
|
||
ids = cached_ids
|
||
elif (
|
||
self.runtime.text_token_ids is not None
|
||
and self.runtime.text_token_mask is not None
|
||
and 0 <= record.slot_index < self.runtime.text_token_ids.shape[1]
|
||
):
|
||
slot = record.slot_index
|
||
slot_mask = self.runtime.text_token_mask[batch_index, slot]
|
||
ids = self.runtime.text_token_ids[batch_index, slot][slot_mask].detach().cpu()
|
||
if ids is not None and ids.numel() > 0:
|
||
ids_parts.append(ids)
|
||
if record.slot_index >= 0:
|
||
slot_list.append(record.slot_index)
|
||
if ids_parts:
|
||
v2_prefix_parts.append(torch.cat(ids_parts, dim=0))
|
||
v2_prefix_masks.append(torch.ones_like(v2_prefix_parts[-1], dtype=torch.bool))
|
||
else:
|
||
v2_prefix_parts.append(torch.zeros(0, dtype=torch.long, device=input_ids.device))
|
||
v2_prefix_masks.append(torch.zeros(0, dtype=torch.bool, device=input_ids.device))
|
||
v2_slots.append(slot_list)
|
||
v2_prefix_length = max((part.numel() for part in v2_prefix_parts), default=0)
|
||
if v2_prefix_length > 0:
|
||
prefix_ids = torch.zeros(
|
||
input_ids.shape[0], v2_prefix_length, dtype=input_ids.dtype, device=input_ids.device
|
||
)
|
||
prefix_mask = torch.zeros(
|
||
input_ids.shape[0], v2_prefix_length, dtype=attention_mask.dtype, device=input_ids.device
|
||
)
|
||
for batch_index, (ids, mask) in enumerate(zip(v2_prefix_parts, v2_prefix_masks)):
|
||
prefix_ids[batch_index, : ids.numel()] = ids.to(input_ids.device)
|
||
prefix_mask[batch_index, : mask.numel()] = mask.to(attention_mask.dtype)
|
||
self.runtime.text_read_slots = torch.tensor(
|
||
[slots + [-1] * max(0, self.memory_config.memory_top_k_records - len(slots)) for slots in v2_slots],
|
||
dtype=torch.long,
|
||
device=input_ids.device,
|
||
)
|
||
self.runtime.text_prefix_used = True
|
||
return prefix_ids, prefix_mask, v2_prefix_length
|
||
|
||
if address is None or self.runtime.text_slot_valid is None:
|
||
return None, None, 0
|
||
|
||
valid = self.runtime.text_slot_valid
|
||
scores = address.masked_fill(~valid, torch.finfo(address.dtype).min)
|
||
top_k = min(self.memory_config.text_memory_top_k, scores.shape[-1])
|
||
top_scores, top_slots = scores.topk(top_k, dim=-1)
|
||
selected = top_scores >= self.memory_config.text_memory_threshold
|
||
# If exactly one stored memory exists, its relevance still has to pass
|
||
# the threshold; this prevents unrelated questions from receiving an
|
||
# arbitrary memory snippet and reduces hallucinated personal facts.
|
||
if relevance is not None:
|
||
selected &= relevance[:, None] >= self.memory_config.text_memory_threshold
|
||
if (
|
||
self.runtime.text_read_overlap is not None
|
||
and not (self.text_retriever is not None and self._text_retriever_ready)
|
||
):
|
||
selected_overlap = self.runtime.text_read_overlap.gather(1, top_slots)
|
||
selected &= selected_overlap >= self.memory_config.text_memory_overlap_threshold
|
||
self.runtime.text_read_slots = top_slots
|
||
if not bool(selected.any()):
|
||
return None, None, 0
|
||
self.runtime.text_prefix_used = True
|
||
|
||
bank_ids = self.runtime.text_token_ids
|
||
bank_mask = self.runtime.text_token_mask
|
||
if bank_ids is None or bank_mask is None:
|
||
return None, None, 0
|
||
prefix_parts: list[Tensor] = []
|
||
prefix_masks: list[Tensor] = []
|
||
for batch_index in range(input_ids.shape[0]):
|
||
ids_parts: list[Tensor] = []
|
||
mask_parts: list[Tensor] = []
|
||
for rank in range(top_k):
|
||
if not bool(selected[batch_index, rank]):
|
||
continue
|
||
slot = int(top_slots[batch_index, rank].item())
|
||
slot_mask = bank_mask[batch_index, slot]
|
||
ids_parts.append(bank_ids[batch_index, slot][slot_mask])
|
||
mask_parts.append(torch.ones_like(ids_parts[-1], dtype=torch.bool))
|
||
if ids_parts:
|
||
prefix_parts.append(torch.cat(ids_parts, dim=0))
|
||
prefix_masks.append(torch.cat(mask_parts, dim=0))
|
||
else:
|
||
prefix_parts.append(torch.zeros(0, dtype=torch.long, device=input_ids.device))
|
||
prefix_masks.append(torch.zeros(0, dtype=torch.bool, device=input_ids.device))
|
||
prefix_length = max((part.numel() for part in prefix_parts), default=0)
|
||
if prefix_length == 0:
|
||
return None, None, 0
|
||
prefix_ids = torch.zeros(
|
||
input_ids.shape[0], prefix_length, dtype=input_ids.dtype, device=input_ids.device
|
||
)
|
||
prefix_mask = torch.zeros(
|
||
input_ids.shape[0], prefix_length, dtype=attention_mask.dtype, device=input_ids.device
|
||
)
|
||
for batch_index, (ids, mask) in enumerate(zip(prefix_parts, prefix_masks)):
|
||
prefix_ids[batch_index, : ids.numel()] = ids.to(input_ids.device)
|
||
prefix_mask[batch_index, : mask.numel()] = mask.to(attention_mask.dtype)
|
||
return prefix_ids, prefix_mask, prefix_length
|
||
|
||
def _find_layer_device(self) -> torch.device:
|
||
for index in self.layer_indices:
|
||
layer = self.base_model.model.language_model.layers[index]
|
||
# In blend mode the adapter owns a trainable scalar. Inspect the
|
||
# original layer first so a CPU-created scalar cannot mislead the
|
||
# device choice for input_ids and the memory module.
|
||
layer_for_device = getattr(layer, "inner", layer)
|
||
for parameter in layer_for_device.parameters():
|
||
if parameter.device.type != "meta":
|
||
return parameter.device
|
||
return self.base_model.get_input_embeddings().weight.device
|
||
|
||
@property
|
||
def trainable_parameters(self):
|
||
parameters = [parameter for parameter in self.memory.parameters() if parameter.requires_grad]
|
||
if self.text_retriever is not None:
|
||
parameters.extend(
|
||
parameter for parameter in self.text_retriever.parameters() if parameter.requires_grad
|
||
)
|
||
if self.memory_policy is not None:
|
||
parameters.extend(
|
||
parameter for parameter in self.memory_policy.parameters() if parameter.requires_grad
|
||
)
|
||
if self.memory_router_v2 is not None:
|
||
parameters.extend(
|
||
parameter for parameter in self.memory_router_v2.parameters() if parameter.requires_grad
|
||
)
|
||
for adapter in self._memory_adapters:
|
||
blend_logit = getattr(adapter, "blend_logit", None)
|
||
if blend_logit is not None and blend_logit.requires_grad:
|
||
parameters.append(blend_logit)
|
||
return iter(parameters)
|
||
|
||
@torch.no_grad()
|
||
def read_hierarchical_memory(
|
||
self,
|
||
query_key: Tensor,
|
||
*,
|
||
query_text: str = "",
|
||
query_token_ids: Optional[Tensor] = None,
|
||
top_k_pages: Optional[int] = None,
|
||
top_k_records: Optional[int] = None,
|
||
max_hops: Optional[int] = None,
|
||
) -> tuple[list[Any], Any]:
|
||
"""Read V2 memory through its bounded page and record router."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
raise RuntimeError("hierarchical memory is disabled in memory_config")
|
||
return self.memory_os_v2.read(
|
||
query_key=query_key,
|
||
query_text=query_text,
|
||
query_token_ids=query_token_ids,
|
||
top_k_pages=top_k_pages,
|
||
top_k_records=top_k_records,
|
||
max_hops=max_hops,
|
||
)
|
||
|
||
@torch.no_grad()
|
||
def write_hierarchical_memory(self, **kwargs: Any) -> tuple[Any, str]:
|
||
"""Write one versioned V2 record owned by the model checkpoint."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
raise RuntimeError("hierarchical memory is disabled in memory_config")
|
||
return self.memory_os_v2.write(**kwargs)
|
||
|
||
def memory_v2_stats(self) -> dict[str, Any]:
|
||
if self.memory_os_v2 is None:
|
||
return {"enabled": False}
|
||
output = dict(self.memory_os_v2.stats())
|
||
output["enabled"] = True
|
||
output["router_ready"] = self._memory_router_v2_ready
|
||
return output
|
||
|
||
def list_memory_records(
|
||
self,
|
||
*,
|
||
query_text: str = "",
|
||
status: str = "active",
|
||
limit: int = 100,
|
||
offset: int = 0,
|
||
) -> list[dict[str, Any]]:
|
||
"""Return a safe, tensor-free view of model-owned memory records."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
return []
|
||
return [
|
||
memory_record_to_dict(record)
|
||
for record in self.memory_os_v2.list_records(
|
||
query_text=query_text,
|
||
status=status,
|
||
limit=limit,
|
||
offset=offset,
|
||
)
|
||
]
|
||
|
||
def get_memory_record(self, record_id: str) -> dict[str, Any]:
|
||
if self.memory_os_v2 is None:
|
||
raise KeyError(record_id)
|
||
record = self.memory_os_v2.bank._resolve_record(record_id)
|
||
return memory_record_to_dict(record)
|
||
|
||
def _clear_legacy_memory_slot(self, slot_index: int) -> None:
|
||
"""Remove a V1 compatibility copy when a V2 record is edited/retracted."""
|
||
|
||
if slot_index < 0:
|
||
return
|
||
fields = (
|
||
("persistent_text_token_ids", "text_token_ids", 0),
|
||
("persistent_text_token_mask", "text_token_mask", 0),
|
||
("persistent_text_slot_valid", "text_slot_valid", 0),
|
||
("persistent_text_slot_keys", "text_slot_keys", 0),
|
||
("persistent_text_slot_age", "text_slot_age", -1),
|
||
("persistent_text_key_token_ids", "text_key_token_ids", 0),
|
||
("persistent_text_key_token_mask", "text_key_token_mask", 0),
|
||
)
|
||
for persistent_name, runtime_name, fill_value in fields:
|
||
persistent = getattr(self, persistent_name, None)
|
||
if isinstance(persistent, Tensor) and persistent.ndim >= 2 and slot_index < persistent.shape[1]:
|
||
persistent[:, slot_index].fill_(fill_value)
|
||
runtime = getattr(self.runtime, runtime_name, None)
|
||
if isinstance(runtime, Tensor) and runtime.ndim >= 2 and slot_index < runtime.shape[1]:
|
||
runtime[:, slot_index].fill_(fill_value)
|
||
|
||
@torch.no_grad()
|
||
def edit_memory_record(
|
||
self,
|
||
record_id: str,
|
||
*,
|
||
text: Optional[str] = None,
|
||
entity: Optional[str] = None,
|
||
attribute: Optional[str] = None,
|
||
value: Optional[str] = None,
|
||
importance: Optional[float] = None,
|
||
confidence: Optional[float] = None,
|
||
evidence: Optional[list[str]] = None,
|
||
token_ids: Optional[Tensor] = None,
|
||
token_mask: Optional[Tensor] = None,
|
||
) -> dict[str, Any]:
|
||
"""Edit memory by creating a versioned successor, never overwriting history."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
raise KeyError(record_id)
|
||
old = self.memory_os_v2.bank._resolve_record(record_id)
|
||
key = None
|
||
if token_ids is not None:
|
||
ids = token_ids.detach().reshape(1, -1).to(self._find_layer_device())
|
||
mask = (
|
||
token_mask.detach().reshape(1, -1).to(self._find_layer_device(), dtype=torch.long)
|
||
if token_mask is not None
|
||
else torch.ones_like(ids)
|
||
)
|
||
key = self._encode_model_key(ids, mask)[0]
|
||
self._clear_legacy_memory_slot(old.slot_index)
|
||
record = self.memory_os_v2.edit_record(
|
||
record_id,
|
||
text=text,
|
||
key=key,
|
||
summary=key,
|
||
entity=entity,
|
||
attribute=attribute,
|
||
value=value,
|
||
importance=importance,
|
||
confidence=confidence,
|
||
evidence=evidence,
|
||
token_ids=token_ids,
|
||
token_mask=token_mask,
|
||
)
|
||
return memory_record_to_dict(record)
|
||
|
||
@torch.no_grad()
|
||
def retract_memory_record(self, record_id: str) -> dict[str, Any]:
|
||
"""Retract a record and its compatibility-slot copy without erasing history."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
raise KeyError(record_id)
|
||
record = self.memory_os_v2.bank._resolve_record(record_id)
|
||
self._clear_legacy_memory_slot(record.slot_index)
|
||
self.memory_os_v2.retract_record(record_id)
|
||
return memory_record_to_dict(record)
|
||
|
||
def audit_memory(self) -> dict[str, Any]:
|
||
if self.memory_os_v2 is None:
|
||
return {"healthy": True, "issues": [], "issue_count": 0, "stats": {"enabled": False}}
|
||
return self.memory_os_v2.audit()
|
||
|
||
def export_memory_records(self, *, limit: int = 10000, offset: int = 0) -> dict[str, Any]:
|
||
records = self.list_memory_records(status="all", limit=limit, offset=offset)
|
||
return {
|
||
"format_version": 1,
|
||
"records": records,
|
||
"offset": offset,
|
||
"limit": limit,
|
||
"returned": len(records),
|
||
"stats": self.memory_v2_stats(),
|
||
}
|
||
|
||
def flush_memory_storage(self) -> None:
|
||
"""Flush the optional warm/cold tier without rewriting model weights."""
|
||
|
||
if self.memory_os_v2 is not None:
|
||
self.memory_os_v2.flush_storage()
|
||
|
||
def close_memory_storage(self) -> None:
|
||
"""Close the optional durable page store before process shutdown."""
|
||
|
||
if self.memory_os_v2 is not None:
|
||
self.memory_os_v2.close_storage()
|
||
|
||
def clear_hierarchical_memory(self) -> None:
|
||
"""Clear the durable V2 page store while keeping router weights."""
|
||
|
||
if self.memory_router_v2 is None:
|
||
return
|
||
if self.memory_os_v2 is not None and self.memory_os_v2.bank.tier_store is not None:
|
||
old_store = self.memory_os_v2.bank.tier_store
|
||
old_store.clear()
|
||
old_store.close()
|
||
self.memory_os_v2 = self._new_memory_os_v2(self.memory.hidden_size)
|
||
|
||
def reset_memory(self, batch_size: Optional[int] = None, *, device: Optional[torch.device] = None) -> None:
|
||
"""Explicitly clear the durable memory state.
|
||
|
||
Use :meth:`reset_runtime_memory` for a temporary conversation or an
|
||
evaluation case. Keeping these operations separate prevents a test
|
||
or a branch from silently deleting a user's embedded checkpoint.
|
||
"""
|
||
|
||
self.runtime.use_persistent_state = True
|
||
self.clear_hierarchical_memory()
|
||
if self.memory_config.persistent_memory:
|
||
device = device or self._find_layer_device()
|
||
batch_size = batch_size or 1
|
||
self.persistent_memory = self.memory.initial_state(batch_size, device=device)
|
||
self._persistent_memory = self.persistent_memory.detach()
|
||
if self.memory_config.natural_language_memory:
|
||
self._ensure_text_memory(batch_size, device=device)
|
||
self._clear_text_memory()
|
||
self.runtime.state = self._persistent_memory
|
||
self.runtime.last_read = None
|
||
self.runtime.auto_memory_probability = None
|
||
self.runtime.raw_memory = None
|
||
self.runtime.reset_mask = None
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask
|
||
self.runtime.text_last_written_slot = None
|
||
return
|
||
if batch_size is None:
|
||
self._persistent_memory = None
|
||
self.runtime.state = None
|
||
self.runtime.last_read = None
|
||
self.runtime.auto_memory_probability = None
|
||
self.runtime.raw_memory = None
|
||
self.runtime.text_token_ids = None
|
||
self.runtime.text_token_mask = None
|
||
self.runtime.text_slot_valid = None
|
||
self.runtime.text_slot_keys = None
|
||
self.runtime.text_slot_age = None
|
||
self.runtime.text_write_counter = None
|
||
self.runtime.text_key_token_ids = None
|
||
self.runtime.text_key_token_mask = None
|
||
self.runtime.text_last_written_slot = None
|
||
return
|
||
device = device or self._find_layer_device()
|
||
self._persistent_memory = self.memory.initial_state(batch_size, device=device)
|
||
if self.memory_config.natural_language_memory:
|
||
self._ensure_text_memory(batch_size, device=device)
|
||
self._clear_text_memory()
|
||
self.runtime.state = self._persistent_memory
|
||
self.runtime.last_read = None
|
||
self.runtime.auto_memory_probability = None
|
||
self.runtime.raw_memory = None
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask
|
||
self.runtime.text_last_written_slot = None
|
||
|
||
def reset_runtime_memory(
|
||
self,
|
||
batch_size: Optional[int] = None,
|
||
*,
|
||
device: Optional[torch.device] = None,
|
||
) -> None:
|
||
"""Reset only the active conversation, preserving durable state.
|
||
|
||
In persistent mode this creates an ephemeral empty bank and prevents
|
||
the generation path from rebinding the embedded user checkpoint. It
|
||
is used by evaluation, memory branches and temporary sessions.
|
||
"""
|
||
|
||
device = device or self._find_layer_device()
|
||
batch_size = batch_size or 1
|
||
self.runtime.use_persistent_state = False
|
||
self.runtime.state = self.memory.initial_state(batch_size, device=device)
|
||
self._persistent_memory = self.runtime.state.detach()
|
||
self.runtime.last_read = None
|
||
self.runtime.auto_memory_probability = None
|
||
self.runtime.raw_memory = None
|
||
self.runtime.v2_query_key = None
|
||
self.runtime.v2_last_decisions = []
|
||
self.runtime.reset_mask = None
|
||
self.runtime.text_last_written_slot = None
|
||
if self.memory_config.natural_language_memory:
|
||
self._bind_text_memory(batch_size, device=device)
|
||
self._clear_text_memory()
|
||
|
||
def _reset_mask(self, input_ids: Optional[Tensor]) -> Optional[Tensor]:
|
||
token_id = self.memory_config.reset_token_id
|
||
if token_id is None or input_ids is None:
|
||
return None
|
||
mask = (input_ids == token_id).any(dim=-1)
|
||
return mask if bool(mask.any()) else None
|
||
|
||
def save_runtime_memory(
|
||
self,
|
||
path: str | Path,
|
||
memory_state: Optional[Tensor] = None,
|
||
) -> Path:
|
||
"""Persist one user's runtime memory, without saving model weights."""
|
||
|
||
state = memory_state if memory_state is not None else self.runtime.state
|
||
if state is None:
|
||
state = self._persistent_memory
|
||
if state is None:
|
||
raise ValueError("no runtime memory state is available to save")
|
||
if state.ndim != 3 or tuple(state.shape[1:]) != (
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.memory_dim,
|
||
):
|
||
raise ValueError(f"unexpected runtime memory shape: {tuple(state.shape)}")
|
||
|
||
raw_memory = self.runtime.raw_memory
|
||
if raw_memory is not None:
|
||
raw_memory = raw_memory.detach().cpu()
|
||
payload = {
|
||
"format_version": 1,
|
||
"memory_state": state.detach().cpu(),
|
||
"raw_memory": raw_memory,
|
||
"model_hidden_size": self.memory.hidden_size,
|
||
"memory_config": asdict(self.memory_config),
|
||
"layer_indices": list(self.layer_indices),
|
||
}
|
||
if self.memory_config.natural_language_memory and self.runtime.text_token_ids is not None:
|
||
payload.update(
|
||
{
|
||
"text_token_ids": self.runtime.text_token_ids.detach().cpu(),
|
||
"text_token_mask": self.runtime.text_token_mask.detach().cpu(),
|
||
"text_slot_valid": self.runtime.text_slot_valid.detach().cpu(),
|
||
"text_slot_keys": self.runtime.text_slot_keys.detach().cpu(),
|
||
"text_slot_age": self.runtime.text_slot_age.detach().cpu(),
|
||
"text_write_counter": self.runtime.text_write_counter.detach().cpu(),
|
||
"text_key_token_ids": self.runtime.text_key_token_ids.detach().cpu(),
|
||
"text_key_token_mask": self.runtime.text_key_token_mask.detach().cpu(),
|
||
}
|
||
)
|
||
output_path = Path(path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
torch.save(payload, output_path)
|
||
return output_path
|
||
|
||
def load_runtime_memory(
|
||
self,
|
||
path: str | Path,
|
||
*,
|
||
device: Optional[torch.device] = None,
|
||
) -> Tensor:
|
||
"""Load a user's runtime memory into this freshly created model."""
|
||
|
||
payload = torch.load(Path(path), map_location="cpu", weights_only=True)
|
||
if isinstance(payload, Tensor):
|
||
state = payload
|
||
raw_memory = None
|
||
elif isinstance(payload, dict):
|
||
state = payload.get("memory_state", payload.get("state"))
|
||
raw_memory = payload.get("raw_memory")
|
||
else:
|
||
state = None
|
||
raw_memory = None
|
||
if not isinstance(state, Tensor):
|
||
raise ValueError("runtime memory file does not contain a tensor memory_state")
|
||
expected_shape = (
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.memory_dim,
|
||
)
|
||
if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape:
|
||
raise ValueError(
|
||
f"runtime memory shape {tuple(state.shape)} does not match {expected_shape}"
|
||
)
|
||
saved_hidden_size = payload.get("model_hidden_size") if isinstance(payload, dict) else None
|
||
if saved_hidden_size is not None and int(saved_hidden_size) != self.memory.hidden_size:
|
||
raise ValueError(
|
||
f"runtime memory hidden size {saved_hidden_size} does not match {self.memory.hidden_size}"
|
||
)
|
||
|
||
target_device = device or self._find_layer_device()
|
||
state = state.to(device=target_device, dtype=self.memory.slot_keys.dtype)
|
||
self._persistent_memory = state.detach()
|
||
if self.memory_config.persistent_memory:
|
||
self.persistent_memory = self._persistent_memory
|
||
self.runtime.state = self._persistent_memory
|
||
self.runtime.last_read = None
|
||
self.runtime.raw_memory = None
|
||
if self.memory_config.natural_language_memory:
|
||
token_ids = payload.get("text_token_ids") if isinstance(payload, dict) else None
|
||
token_mask = payload.get("text_token_mask") if isinstance(payload, dict) else None
|
||
slot_valid = payload.get("text_slot_valid") if isinstance(payload, dict) else None
|
||
slot_keys = payload.get("text_slot_keys") if isinstance(payload, dict) else None
|
||
slot_age = payload.get("text_slot_age") if isinstance(payload, dict) else None
|
||
write_counter = payload.get("text_write_counter") if isinstance(payload, dict) else None
|
||
key_token_ids = payload.get("text_key_token_ids") if isinstance(payload, dict) else None
|
||
key_token_mask = payload.get("text_key_token_mask") if isinstance(payload, dict) else None
|
||
if not all(
|
||
isinstance(item, Tensor)
|
||
for item in (
|
||
token_ids,
|
||
token_mask,
|
||
slot_valid,
|
||
slot_keys,
|
||
slot_age,
|
||
write_counter,
|
||
key_token_ids,
|
||
key_token_mask,
|
||
)
|
||
):
|
||
raise ValueError("natural-language memory file is missing its indexed text bank")
|
||
expected_text_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_tokens,
|
||
)
|
||
if tuple(token_ids.shape) != expected_text_shape or tuple(token_mask.shape) != expected_text_shape:
|
||
raise ValueError("natural-language text bank shape does not match memory config")
|
||
if tuple(slot_valid.shape) != expected_text_shape[:2]:
|
||
raise ValueError("natural-language slot validity shape does not match memory config")
|
||
if tuple(slot_keys.shape) != (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size):
|
||
raise ValueError("natural-language slot key shape does not match memory config")
|
||
if tuple(slot_age.shape) != expected_text_shape[:2] or tuple(write_counter.shape) != (state.shape[0],):
|
||
raise ValueError("natural-language slot age shape does not match memory config")
|
||
expected_key_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_key_tokens,
|
||
)
|
||
if tuple(key_token_ids.shape) != expected_key_shape or tuple(key_token_mask.shape) != expected_key_shape:
|
||
raise ValueError("natural-language key token shape does not match memory config")
|
||
self.persistent_text_token_ids = token_ids.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_token_mask = token_mask.to(device=target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_valid = slot_valid.to(device=target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_keys = slot_keys.to(device=target_device, dtype=torch.float32)
|
||
self.persistent_text_slot_age = slot_age.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_write_counter = write_counter.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_ids = key_token_ids.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_mask = key_token_mask.to(device=target_device, dtype=torch.bool)
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask
|
||
if raw_memory is not None:
|
||
if not isinstance(raw_memory, Tensor):
|
||
raise ValueError("runtime memory raw_memory must be a tensor or null")
|
||
if raw_memory.ndim == 1:
|
||
raw_memory = raw_memory.unsqueeze(0)
|
||
if tuple(raw_memory.shape) != (state.shape[0], self.memory.hidden_size):
|
||
raise ValueError(f"unexpected raw_memory shape: {tuple(raw_memory.shape)}")
|
||
output_weight = self.base_model.get_output_embeddings().weight
|
||
self.runtime.raw_memory = raw_memory.to(
|
||
device=output_weight.device,
|
||
dtype=output_weight.dtype,
|
||
).detach()
|
||
return self._persistent_memory
|
||
|
||
def save_memory_adapter(self, output_dir: str | Path) -> None:
|
||
path = Path(output_dir)
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
torch.save(self.memory.state_dict(), path / "memory.pt")
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
torch.save(self.text_retriever.state_dict(), path / "text_retriever.pt")
|
||
if self.memory_policy is not None and self._memory_policy_ready:
|
||
torch.save(self.memory_policy.state_dict(), path / "memory_policy.pt")
|
||
metadata = {
|
||
"hidden_size": self.memory.hidden_size,
|
||
"memory_config": asdict(self.memory_config),
|
||
"layer_indices": list(self.layer_indices),
|
||
"text_retriever_ready": self._text_retriever_ready,
|
||
"memory_policy_ready": self._memory_policy_ready,
|
||
}
|
||
(path / "memory_config.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
|
||
surgery_state = {
|
||
str(index): adapter.blend_logit.detach().cpu()
|
||
for index, adapter in zip(self.layer_indices, self._memory_adapters)
|
||
if hasattr(adapter, "blend_logit")
|
||
}
|
||
torch.save({"blend_logits": surgery_state}, path / "surgery.pt")
|
||
if self.memory_config.persistent_memory and self.persistent_memory.numel() > 0:
|
||
payload = {
|
||
"format_version": 1,
|
||
"memory_state": self.persistent_memory.detach().cpu(),
|
||
}
|
||
if self.memory_config.natural_language_memory:
|
||
payload.update(
|
||
{
|
||
"text_token_ids": self.persistent_text_token_ids.detach().cpu(),
|
||
"text_token_mask": self.persistent_text_token_mask.detach().cpu(),
|
||
"text_slot_valid": self.persistent_text_slot_valid.detach().cpu(),
|
||
"text_slot_keys": self.persistent_text_slot_keys.detach().cpu(),
|
||
"text_slot_age": self.persistent_text_slot_age.detach().cpu(),
|
||
"text_write_counter": self.persistent_text_write_counter.detach().cpu(),
|
||
"text_key_token_ids": self.persistent_text_key_token_ids.detach().cpu(),
|
||
"text_key_token_mask": self.persistent_text_key_token_mask.detach().cpu(),
|
||
}
|
||
)
|
||
torch.save(
|
||
payload,
|
||
path / "persistent_memory.pt",
|
||
)
|
||
|
||
def save_persistent_memory_checkpoint(self, output_dir: str | Path) -> None:
|
||
"""Save controller weights and the current user memory together.
|
||
|
||
This produces a compact model-owned memory checkpoint. The large
|
||
frozen Qwen shards remain untouched; the learned reader/controller and
|
||
the user's current memory are stored in the adapter package.
|
||
"""
|
||
|
||
state = self.runtime.state if self.runtime.state is not None else self._persistent_memory
|
||
if state is None:
|
||
raise ValueError("no runtime memory state is available to checkpoint")
|
||
was_persistent = self.memory_config.persistent_memory
|
||
self.memory_config.persistent_memory = True
|
||
self.persistent_memory = state.detach().clone()
|
||
self._persistent_memory = self.persistent_memory.detach()
|
||
if self.memory_config.natural_language_memory:
|
||
text_fields = (
|
||
self.runtime.text_token_ids,
|
||
self.runtime.text_token_mask,
|
||
self.runtime.text_slot_valid,
|
||
self.runtime.text_slot_keys,
|
||
self.runtime.text_slot_age,
|
||
self.runtime.text_write_counter,
|
||
self.runtime.text_key_token_ids,
|
||
self.runtime.text_key_token_mask,
|
||
)
|
||
if not all(isinstance(item, Tensor) for item in text_fields):
|
||
raise ValueError("natural-language memory bank is not initialized")
|
||
target_device = self._find_layer_device()
|
||
self.persistent_text_token_ids = text_fields[0].detach().clone().to(target_device, dtype=torch.long)
|
||
self.persistent_text_token_mask = text_fields[1].detach().clone().to(target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_valid = text_fields[2].detach().clone().to(target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_keys = text_fields[3].detach().clone().to(target_device, dtype=torch.float32)
|
||
self.persistent_text_slot_age = text_fields[4].detach().clone().to(target_device, dtype=torch.long)
|
||
self.persistent_text_write_counter = text_fields[5].detach().clone().to(target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_ids = text_fields[6].detach().clone().to(target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_mask = text_fields[7].detach().clone().to(target_device, dtype=torch.bool)
|
||
try:
|
||
self.save_memory_adapter(output_dir)
|
||
path = Path(output_dir)
|
||
metadata_path = path / "memory_config.json"
|
||
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||
metadata["memory_config"]["persistent_memory"] = True
|
||
metadata["memory_config"]["checkpoint_contains_user_memory"] = True
|
||
metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
|
||
finally:
|
||
self.memory_config.persistent_memory = was_persistent
|
||
|
||
def load_memory_adapter(self, adapter_dir: str | Path, *, strict: bool = True) -> None:
|
||
path = Path(adapter_dir)
|
||
state = torch.load(path / "memory.pt", map_location=self._find_layer_device(), weights_only=True)
|
||
self.memory.load_state_dict(state, strict=strict)
|
||
metadata_ready = None
|
||
policy_ready = None
|
||
metadata_path = path / "memory_config.json"
|
||
if metadata_path.exists():
|
||
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||
if isinstance(metadata, dict) and "text_retriever_ready" in metadata:
|
||
metadata_ready = bool(metadata["text_retriever_ready"])
|
||
if isinstance(metadata, dict) and "memory_policy_ready" in metadata:
|
||
policy_ready = bool(metadata["memory_policy_ready"])
|
||
if metadata_ready and self.text_retriever is None:
|
||
raise ValueError(
|
||
"adapter contains a trained text retriever but natural_language_memory is disabled"
|
||
)
|
||
retriever_path = path / "text_retriever.pt"
|
||
if self.text_retriever is not None and metadata_ready and not retriever_path.exists():
|
||
raise FileNotFoundError(
|
||
f"adapter declares a trained text retriever but is missing {retriever_path}"
|
||
)
|
||
should_load_retriever = (
|
||
self.text_retriever is not None
|
||
and retriever_path.exists()
|
||
and metadata_ready is not False
|
||
)
|
||
if should_load_retriever:
|
||
retriever_state = torch.load(
|
||
retriever_path,
|
||
map_location=self._find_layer_device(),
|
||
weights_only=True,
|
||
)
|
||
self.text_retriever.load_state_dict(retriever_state, strict=True)
|
||
self._text_retriever_ready = True
|
||
policy_path = path / "memory_policy.pt"
|
||
if policy_ready and self.memory_policy is None:
|
||
raise ValueError(
|
||
"adapter contains an automatic memory policy but automatic_memory is disabled"
|
||
)
|
||
if self.memory_policy is not None and policy_ready is not None:
|
||
if policy_ready and not policy_path.exists():
|
||
raise FileNotFoundError(
|
||
f"adapter declares an automatic memory policy but is missing {policy_path}"
|
||
)
|
||
if policy_ready:
|
||
policy_state = torch.load(
|
||
policy_path,
|
||
map_location=self._find_layer_device(),
|
||
weights_only=True,
|
||
)
|
||
self.memory_policy.load_state_dict(policy_state, strict=strict)
|
||
self._memory_policy_ready = True
|
||
surgery_path = path / "surgery.pt"
|
||
if surgery_path.exists():
|
||
surgery_state = torch.load(surgery_path, map_location="cpu", weights_only=True)
|
||
blend_logits = surgery_state.get("blend_logits", {})
|
||
for index, adapter in zip(self.layer_indices, self._memory_adapters):
|
||
if hasattr(adapter, "blend_logit") and str(index) in blend_logits:
|
||
adapter.blend_logit.data.copy_(blend_logits[str(index)].to(adapter.blend_logit.device))
|
||
persistent_path = path / "persistent_memory.pt"
|
||
if persistent_path.exists():
|
||
payload = torch.load(persistent_path, map_location="cpu", weights_only=True)
|
||
state = payload.get("memory_state", payload.get("state")) if isinstance(payload, dict) else payload
|
||
if not isinstance(state, Tensor):
|
||
raise ValueError("persistent_memory.pt does not contain a tensor memory_state")
|
||
expected_shape = (self.memory_config.memory_slots, self.memory_config.memory_dim)
|
||
if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape:
|
||
raise ValueError(
|
||
f"persistent memory shape {tuple(state.shape)} does not match {expected_shape}"
|
||
)
|
||
target_device = self._find_layer_device()
|
||
self.memory_config.persistent_memory = True
|
||
self.persistent_memory = state.to(
|
||
device=target_device,
|
||
dtype=self.memory.slot_keys.dtype,
|
||
)
|
||
self._persistent_memory = self.persistent_memory.detach()
|
||
self.runtime.state = self._persistent_memory
|
||
if self.memory_config.natural_language_memory:
|
||
token_ids = payload.get("text_token_ids") if isinstance(payload, dict) else None
|
||
token_mask = payload.get("text_token_mask") if isinstance(payload, dict) else None
|
||
slot_valid = payload.get("text_slot_valid") if isinstance(payload, dict) else None
|
||
slot_keys = payload.get("text_slot_keys") if isinstance(payload, dict) else None
|
||
slot_age = payload.get("text_slot_age") if isinstance(payload, dict) else None
|
||
write_counter = payload.get("text_write_counter") if isinstance(payload, dict) else None
|
||
key_token_ids = payload.get("text_key_token_ids") if isinstance(payload, dict) else None
|
||
key_token_mask = payload.get("text_key_token_mask") if isinstance(payload, dict) else None
|
||
if not all(
|
||
isinstance(item, Tensor)
|
||
for item in (
|
||
token_ids,
|
||
token_mask,
|
||
slot_valid,
|
||
slot_keys,
|
||
slot_age,
|
||
write_counter,
|
||
key_token_ids,
|
||
key_token_mask,
|
||
)
|
||
):
|
||
raise ValueError("persistent adapter is missing its indexed natural-language bank")
|
||
expected_text_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_tokens,
|
||
)
|
||
if tuple(token_ids.shape) != expected_text_shape or tuple(token_mask.shape) != expected_text_shape:
|
||
raise ValueError("persistent text bank shape does not match memory config")
|
||
if tuple(slot_valid.shape) != expected_text_shape[:2]:
|
||
raise ValueError("persistent slot validity shape does not match memory config")
|
||
if tuple(slot_keys.shape) != (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size):
|
||
raise ValueError("persistent slot key shape does not match memory config")
|
||
if tuple(slot_age.shape) != expected_text_shape[:2] or tuple(write_counter.shape) != (state.shape[0],):
|
||
raise ValueError("persistent slot age shape does not match memory config")
|
||
expected_key_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_key_tokens,
|
||
)
|
||
if tuple(key_token_ids.shape) != expected_key_shape or tuple(key_token_mask.shape) != expected_key_shape:
|
||
raise ValueError("persistent key token shape does not match memory config")
|
||
self.persistent_text_token_ids = token_ids.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_token_mask = token_mask.to(device=target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_valid = slot_valid.to(device=target_device, dtype=torch.bool)
|
||
self.persistent_text_slot_keys = slot_keys.to(device=target_device, dtype=torch.float32)
|
||
self.persistent_text_slot_age = slot_age.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_write_counter = write_counter.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_ids = key_token_ids.to(device=target_device, dtype=torch.long)
|
||
self.persistent_text_key_token_mask = key_token_mask.to(device=target_device, dtype=torch.bool)
|
||
self.runtime.text_token_ids = self.persistent_text_token_ids
|
||
self.runtime.text_token_mask = self.persistent_text_token_mask
|
||
self.runtime.text_slot_valid = self.persistent_text_slot_valid
|
||
self.runtime.text_slot_keys = self.persistent_text_slot_keys
|
||
self.runtime.text_slot_age = self.persistent_text_slot_age
|
||
self.runtime.text_write_counter = self.persistent_text_write_counter
|
||
self.runtime.text_key_token_ids = self.persistent_text_key_token_ids
|
||
self.runtime.text_key_token_mask = self.persistent_text_key_token_mask
|
||
|
||
@torch.no_grad()
|
||
def _load_persistent_memory_payload(self, payload: dict[str, Any]) -> None:
|
||
"""Load a runtime/persistent memory payload shared by .pt and safetensors."""
|
||
|
||
state = payload.get("memory_state", payload.get("state"))
|
||
if not isinstance(state, Tensor):
|
||
raise ValueError("persistent memory payload does not contain a tensor memory_state")
|
||
expected_shape = (self.memory_config.memory_slots, self.memory_config.memory_dim)
|
||
if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape:
|
||
raise ValueError(
|
||
f"persistent memory shape {tuple(state.shape)} does not match {expected_shape}"
|
||
)
|
||
target_device = self._find_layer_device()
|
||
self.memory_config.persistent_memory = True
|
||
self.runtime.use_persistent_state = True
|
||
self.persistent_memory = state.to(device=target_device, dtype=self.memory.slot_keys.dtype)
|
||
self._persistent_memory = self.persistent_memory.detach()
|
||
self.runtime.state = self._persistent_memory
|
||
if not self.memory_config.natural_language_memory:
|
||
return
|
||
|
||
fields = {
|
||
"persistent_text_token_ids": ("text_token_ids", torch.long),
|
||
"persistent_text_token_mask": ("text_token_mask", torch.bool),
|
||
"persistent_text_slot_valid": ("text_slot_valid", torch.bool),
|
||
"persistent_text_slot_keys": ("text_slot_keys", torch.float32),
|
||
"persistent_text_slot_age": ("text_slot_age", torch.long),
|
||
"persistent_text_write_counter": ("text_write_counter", torch.long),
|
||
"persistent_text_key_token_ids": ("text_key_token_ids", torch.long),
|
||
"persistent_text_key_token_mask": ("text_key_token_mask", torch.bool),
|
||
}
|
||
if not all(isinstance(payload.get(key), Tensor) for key, _ in fields.values()):
|
||
raise ValueError("persistent memory payload is missing its indexed natural-language bank")
|
||
expected_text_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_tokens,
|
||
)
|
||
expected_key_shape = (
|
||
state.shape[0],
|
||
self.memory_config.memory_slots,
|
||
self.memory_config.text_memory_key_tokens,
|
||
)
|
||
expected_shapes = {
|
||
"text_token_ids": expected_text_shape,
|
||
"text_token_mask": expected_text_shape,
|
||
"text_slot_valid": expected_text_shape[:2],
|
||
"text_slot_keys": (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size),
|
||
"text_slot_age": expected_text_shape[:2],
|
||
"text_write_counter": (state.shape[0],),
|
||
"text_key_token_ids": expected_key_shape,
|
||
"text_key_token_mask": expected_key_shape,
|
||
}
|
||
for key, _ in fields.values():
|
||
if tuple(payload[key].shape) != expected_shapes[key]:
|
||
raise ValueError(f"persistent field {key} shape does not match memory config")
|
||
for attribute, (key, dtype) in fields.items():
|
||
setattr(self, attribute, payload[key].to(device=target_device, dtype=dtype))
|
||
runtime_fields = {
|
||
"text_token_ids": self.persistent_text_token_ids,
|
||
"text_token_mask": self.persistent_text_token_mask,
|
||
"text_slot_valid": self.persistent_text_slot_valid,
|
||
"text_slot_keys": self.persistent_text_slot_keys,
|
||
"text_slot_age": self.persistent_text_slot_age,
|
||
"text_write_counter": self.persistent_text_write_counter,
|
||
"text_key_token_ids": self.persistent_text_key_token_ids,
|
||
"text_key_token_mask": self.persistent_text_key_token_mask,
|
||
}
|
||
for attribute, value in runtime_fields.items():
|
||
setattr(self.runtime, attribute, value)
|
||
|
||
def _export_memory_os_v2_checkpoint(self) -> tuple[dict[str, Tensor], dict[str, Any]]:
|
||
"""Split a V2 bank into tensors plus JSON-safe checkpoint metadata."""
|
||
|
||
if self.memory_os_v2 is None:
|
||
return {}, {}
|
||
payload = self.memory_os_v2.export_payload()
|
||
tensors: dict[str, Tensor] = {}
|
||
metadata = dict(payload)
|
||
|
||
def pack_record(item: dict[str, Any], prefix: str) -> dict[str, Any]:
|
||
item = dict(item)
|
||
for field in ("key", "summary", "token_ids", "token_mask"):
|
||
value = item.pop(field, None)
|
||
if isinstance(value, Tensor):
|
||
name = f"dynamic_memory.v2.{prefix}.{field}"
|
||
tensors[name] = value.detach().cpu().contiguous()
|
||
item[f"{field}_ref"] = name
|
||
return item
|
||
|
||
metadata["records"] = [
|
||
pack_record(item, f"records.{index}")
|
||
for index, item in enumerate(payload.get("records", []))
|
||
]
|
||
metadata["quarantine"] = [
|
||
pack_record(item, f"quarantine.{index}")
|
||
for index, item in enumerate(payload.get("quarantine", []))
|
||
]
|
||
metadata["pages"] = []
|
||
for index, item in enumerate(payload.get("pages", [])):
|
||
item = dict(item)
|
||
for field in ("key", "summary"):
|
||
value = item.pop(field, None)
|
||
if isinstance(value, Tensor):
|
||
name = f"dynamic_memory.v2.pages.{index}.{field}"
|
||
tensors[name] = value.detach().cpu().contiguous()
|
||
item[f"{field}_ref"] = name
|
||
metadata["pages"].append(item)
|
||
return tensors, metadata
|
||
|
||
def _load_memory_os_v2_checkpoint(
|
||
self,
|
||
metadata_text: str,
|
||
weights: dict[str, Tensor],
|
||
) -> None:
|
||
"""Rehydrate the V2 page store from safetensors metadata and tensors."""
|
||
|
||
if self.memory_router_v2 is None:
|
||
return
|
||
payload = json.loads(metadata_text)
|
||
|
||
def unpack_record(item: dict[str, Any]) -> dict[str, Any]:
|
||
item = dict(item)
|
||
for field in ("key", "summary", "token_ids", "token_mask"):
|
||
ref = item.pop(f"{field}_ref", None)
|
||
item[field] = weights.get(ref) if ref else None
|
||
return item
|
||
|
||
payload["records"] = [unpack_record(item) for item in payload.get("records", [])]
|
||
payload["quarantine"] = [unpack_record(item) for item in payload.get("quarantine", [])]
|
||
unpacked_pages = []
|
||
for item in payload.get("pages", []):
|
||
item = dict(item)
|
||
for field in ("key", "summary"):
|
||
ref = item.pop(f"{field}_ref", None)
|
||
item[field] = weights.get(ref) if ref else None
|
||
unpacked_pages.append(item)
|
||
payload["pages"] = unpacked_pages
|
||
tier_store = None
|
||
if self.memory_config.memory_storage_mode == "tiered":
|
||
if not self.memory_config.memory_storage_path:
|
||
raise ValueError("tiered memory requires memory_storage_path")
|
||
tier_store = TieredMemoryStoreV2(
|
||
self.memory_config.memory_storage_path,
|
||
key_dim=self.memory_config.memory_router_dim,
|
||
page_capacity=self.memory_config.memory_page_capacity,
|
||
)
|
||
self.memory_os_v2 = MemoryOSV2.from_payload(
|
||
payload,
|
||
router=self.memory_router_v2,
|
||
tier_store=tier_store,
|
||
max_resident_pages=self.memory_config.memory_resident_pages,
|
||
runtime_device=self._find_layer_device(),
|
||
gpu_cache_records=self.memory_config.memory_gpu_cache_records,
|
||
gpu_cache_tokens=self.memory_config.memory_gpu_cache_tokens,
|
||
gpu_cache_reserve_mb=self.memory_config.memory_gpu_cache_reserve_mb,
|
||
gpu_cache_adaptive=self.memory_config.memory_gpu_cache_adaptive,
|
||
)
|
||
# The package configuration is authoritative for runtime safety and
|
||
# capacity gates. Older snapshots may contain looser values.
|
||
self.memory_os_v2.read_threshold = self.memory_config.memory_v2_read_threshold
|
||
self.memory_os_v2.write_threshold = self.memory_config.memory_v2_write_threshold
|
||
bank = self.memory_os_v2.bank
|
||
bank.page_capacity = self.memory_config.memory_page_capacity
|
||
bank.max_pages = self.memory_config.memory_max_pages
|
||
bank.hot_pages = self.memory_config.memory_hot_pages
|
||
bank.top_k_pages = self.memory_config.memory_top_k_pages
|
||
bank.top_k_records = self.memory_config.memory_top_k_records
|
||
bank.max_hops = self.memory_config.memory_max_hops
|
||
bank._rebuild_coarse_index(self.memory_config.memory_coarse_index_bits)
|
||
# The package configuration is authoritative for runtime safety gates.
|
||
# Older snapshots may contain a looser threshold in their serialized
|
||
# payload; do not let that silently re-enable noisy recalls after an
|
||
# upgrade.
|
||
self.memory_os_v2.read_threshold = self.memory_config.memory_v2_read_threshold
|
||
self.memory_os_v2.write_threshold = self.memory_config.memory_v2_write_threshold
|
||
|
||
@torch.no_grad()
|
||
def load_embedded_memory_weights(self, merged_dir: str | Path) -> None:
|
||
"""Load memory modules and optional user state from one safetensors package."""
|
||
|
||
path = Path(merged_dir)
|
||
manifest = json.loads((path / "memory_merge.json").read_text(encoding="utf-8"))
|
||
from safetensors import safe_open
|
||
|
||
weights_path = path / str(manifest.get("memory_weights", "memory.safetensors"))
|
||
tensor_prefix = str(manifest.get("tensor_prefix", "dynamic_memory."))
|
||
with safe_open(str(weights_path), framework="pt", device="cpu") as handle:
|
||
file_metadata = handle.metadata() or {}
|
||
weights = {
|
||
key: handle.get_tensor(key)
|
||
for key in handle.keys()
|
||
if key.startswith(tensor_prefix)
|
||
}
|
||
|
||
memory_prefix = f"{tensor_prefix}memory."
|
||
memory_state = {
|
||
key.removeprefix(memory_prefix): value
|
||
for key, value in weights.items()
|
||
if key.startswith(memory_prefix)
|
||
}
|
||
if not memory_state:
|
||
raise ValueError(f"merged memory file has no {memory_prefix}* tensors: {weights_path}")
|
||
self.memory.load_state_dict(memory_state, strict=True)
|
||
|
||
retriever_prefix = f"{tensor_prefix}text_retriever."
|
||
retriever_state = {
|
||
key.removeprefix(retriever_prefix): value
|
||
for key, value in weights.items()
|
||
if key.startswith(retriever_prefix)
|
||
}
|
||
if retriever_state:
|
||
if self.text_retriever is None:
|
||
raise ValueError("merged package contains a text retriever but it is disabled")
|
||
self.text_retriever.load_state_dict(retriever_state, strict=True)
|
||
self._text_retriever_ready = True
|
||
|
||
policy_prefix = f"{tensor_prefix}memory_policy."
|
||
policy_state = {
|
||
key.removeprefix(policy_prefix): value
|
||
for key, value in weights.items()
|
||
if key.startswith(policy_prefix)
|
||
}
|
||
if policy_state:
|
||
if self.memory_policy is None:
|
||
raise ValueError("merged package contains a memory policy but it is disabled")
|
||
# A v2 candidate may be loaded on top of an older embedded package:
|
||
# the old shard has only the write head, while the candidate adds
|
||
# a trained forget head. The adapter load immediately afterwards
|
||
# is strict and supplies the complete v2 policy.
|
||
self.memory_policy.load_state_dict(
|
||
policy_state,
|
||
strict=self.memory_config.automatic_memory_policy_version < 2,
|
||
)
|
||
self._memory_policy_ready = True
|
||
|
||
router_prefix = f"{tensor_prefix}memory_router_v2."
|
||
router_state = {
|
||
key.removeprefix(router_prefix): value
|
||
for key, value in weights.items()
|
||
if key.startswith(router_prefix)
|
||
}
|
||
if router_state:
|
||
if self.memory_router_v2 is None:
|
||
raise ValueError("merged package contains a V2 router but hierarchical memory is disabled")
|
||
self.memory_router_v2.load_state_dict(router_state, strict=True)
|
||
self._memory_router_v2_ready = True
|
||
|
||
for index, adapter in zip(self.layer_indices, self._memory_adapters):
|
||
key = f"{tensor_prefix}blend_logits.{index}"
|
||
if hasattr(adapter, "blend_logit") and key in weights:
|
||
adapter.blend_logit.data.copy_(weights[key].to(adapter.blend_logit.device))
|
||
|
||
persistent_prefix = f"{tensor_prefix}persistent."
|
||
persistent_payload = {
|
||
key.removeprefix(persistent_prefix): value
|
||
for key, value in weights.items()
|
||
if key.startswith(persistent_prefix)
|
||
}
|
||
if persistent_payload:
|
||
self._load_persistent_memory_payload(persistent_payload)
|
||
v2_payload = file_metadata.get("memory_os_v2_payload")
|
||
if v2_payload and self.memory_os_v2 is not None:
|
||
self._load_memory_os_v2_checkpoint(v2_payload, weights)
|
||
|
||
@torch.no_grad()
|
||
def save_embedded_memory_weights(self, merged_dir: str | Path) -> None:
|
||
"""Atomically write the current memory snapshot back into the main shard.
|
||
|
||
The merged package keeps the original Qwen shard as the base payload
|
||
and appends ``dynamic_memory.*`` tensors to the same safetensors file.
|
||
This intentionally rewrites the package's second shard, so callers
|
||
should use it as a deliberate weight-backed persistence mode rather
|
||
than on every token.
|
||
"""
|
||
|
||
path = Path(merged_dir)
|
||
manifest = json.loads((path / "memory_merge.json").read_text(encoding="utf-8"))
|
||
from safetensors.torch import save_file
|
||
from .merge_memory_weights import _write_merged_shard
|
||
|
||
memory_weights_name = str(
|
||
manifest.get("memory_weights", "model.safetensors-00002-of-00002.safetensors")
|
||
)
|
||
target_shard = path / memory_weights_name
|
||
if not target_shard.exists():
|
||
raise FileNotFoundError(f"merged package target shard is missing: {target_shard}")
|
||
|
||
tensors: dict[str, Tensor] = {}
|
||
tensors.update(
|
||
{
|
||
f"dynamic_memory.memory.{key}": value.detach().cpu().contiguous()
|
||
for key, value in self.memory.state_dict().items()
|
||
}
|
||
)
|
||
if self.text_retriever is not None and self._text_retriever_ready:
|
||
tensors.update(
|
||
{
|
||
f"dynamic_memory.text_retriever.{key}": value.detach().cpu().contiguous()
|
||
for key, value in self.text_retriever.state_dict().items()
|
||
}
|
||
)
|
||
if self.memory_router_v2 is not None:
|
||
tensors.update(
|
||
{
|
||
f"dynamic_memory.memory_router_v2.{key}": value.detach().cpu().contiguous()
|
||
for key, value in self.memory_router_v2.state_dict().items()
|
||
}
|
||
)
|
||
if self.memory_policy is not None and self._memory_policy_ready:
|
||
tensors.update(
|
||
{
|
||
f"dynamic_memory.memory_policy.{key}": value.detach().cpu().contiguous()
|
||
for key, value in self.memory_policy.state_dict().items()
|
||
}
|
||
)
|
||
for index, adapter in zip(self.layer_indices, self._memory_adapters):
|
||
if hasattr(adapter, "blend_logit"):
|
||
tensors[f"dynamic_memory.blend_logits.{index}"] = (
|
||
adapter.blend_logit.detach().cpu().contiguous()
|
||
)
|
||
|
||
if self.memory_config.persistent_memory and self.persistent_memory.numel() > 0:
|
||
state = self.persistent_memory
|
||
else:
|
||
state = self.runtime.state if self.runtime.state is not None else self._persistent_memory
|
||
if state is None:
|
||
raise ValueError("cannot write embedded memory without a runtime memory state")
|
||
tensors["dynamic_memory.persistent.memory_state"] = state.detach().cpu().contiguous()
|
||
if self.memory_config.natural_language_memory:
|
||
if self.memory_config.persistent_memory and self.persistent_text_token_ids.numel() > 0:
|
||
runtime_fields = {
|
||
"text_token_ids": self.persistent_text_token_ids,
|
||
"text_token_mask": self.persistent_text_token_mask,
|
||
"text_slot_valid": self.persistent_text_slot_valid,
|
||
"text_slot_keys": self.persistent_text_slot_keys,
|
||
"text_slot_age": self.persistent_text_slot_age,
|
||
"text_write_counter": self.persistent_text_write_counter,
|
||
"text_key_token_ids": self.persistent_text_key_token_ids,
|
||
"text_key_token_mask": self.persistent_text_key_token_mask,
|
||
}
|
||
else:
|
||
runtime_fields = {
|
||
"text_token_ids": self.runtime.text_token_ids,
|
||
"text_token_mask": self.runtime.text_token_mask,
|
||
"text_slot_valid": self.runtime.text_slot_valid,
|
||
"text_slot_keys": self.runtime.text_slot_keys,
|
||
"text_slot_age": self.runtime.text_slot_age,
|
||
"text_write_counter": self.runtime.text_write_counter,
|
||
"text_key_token_ids": self.runtime.text_key_token_ids,
|
||
"text_key_token_mask": self.runtime.text_key_token_mask,
|
||
}
|
||
if not all(isinstance(value, Tensor) for value in runtime_fields.values()):
|
||
raise ValueError("cannot write embedded memory without a complete text memory bank")
|
||
tensors.update(
|
||
{
|
||
f"dynamic_memory.persistent.{key}": value.detach().cpu().contiguous()
|
||
for key, value in runtime_fields.items()
|
||
}
|
||
)
|
||
|
||
v2_metadata: dict[str, Any] = {}
|
||
if self.memory_os_v2 is not None:
|
||
v2_tensors, v2_metadata = self._export_memory_os_v2_checkpoint()
|
||
tensors.update(v2_tensors)
|
||
|
||
extra_path = path / f".{memory_weights_name}.memory-update.tmp"
|
||
staging_path = path / f".{memory_weights_name}.merged-update.tmp"
|
||
|
||
def _sync_model_index(shard_name: str) -> None:
|
||
"""Keep the HF shard index aligned with the embedded memory shard."""
|
||
|
||
index_path = path / "model.safetensors.index.json"
|
||
if not index_path.exists():
|
||
return
|
||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||
weight_map = index.get("weight_map")
|
||
if not isinstance(weight_map, dict):
|
||
return
|
||
changed = False
|
||
for key, value in list(weight_map.items()):
|
||
if str(key).startswith("dynamic_memory.") and value != shard_name:
|
||
weight_map[key] = shard_name
|
||
changed = True
|
||
if changed:
|
||
index_path.write_text(
|
||
json.dumps(index, indent=2, ensure_ascii=False), encoding="utf-8"
|
||
)
|
||
|
||
try:
|
||
metadata = {"format": "qwen_dynamic_memory_embedded_v2"}
|
||
if v2_metadata:
|
||
metadata["memory_os_v2_payload"] = json.dumps(
|
||
v2_metadata, ensure_ascii=False, separators=(",", ":")
|
||
)
|
||
save_file(tensors, str(extra_path), metadata=metadata)
|
||
if str(manifest.get("shard_mode", "combined")) == "extra_only":
|
||
# The extra-only layout keeps the official Qwen shards
|
||
# untouched. Windows may keep the currently loaded shard
|
||
# open, so rotate to a new small shard when replacement is
|
||
# denied instead of requiring the chat process to exit.
|
||
try:
|
||
os.replace(extra_path, target_shard)
|
||
except PermissionError:
|
||
rotated_name = f"{target_shard.stem}.runtime-{time.time_ns()}.safetensors"
|
||
rotated_path = path / rotated_name
|
||
os.replace(extra_path, rotated_path)
|
||
manifest["memory_weights"] = rotated_name
|
||
manifest["checkpoint_contains_user_memory"] = True
|
||
(path / "memory_merge.json").write_text(
|
||
json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8"
|
||
)
|
||
_sync_model_index(str(manifest.get("memory_weights", memory_weights_name)))
|
||
return
|
||
base_dir = Path(str(manifest["base_model"]))
|
||
base_shard = base_dir / memory_weights_name
|
||
if not base_shard.exists():
|
||
raise FileNotFoundError(f"merged package base shard is missing: {base_shard}")
|
||
_write_merged_shard(base_shard, extra_path, staging_path)
|
||
os.replace(staging_path, target_shard)
|
||
_sync_model_index(memory_weights_name)
|
||
finally:
|
||
if extra_path.exists():
|
||
extra_path.unlink()
|
||
if staging_path.exists():
|
||
staging_path.unlink()
|
||
|
||
def forward(
|
||
self,
|
||
*args: Any,
|
||
memory_state: Optional[Tensor] = None,
|
||
update_memory: bool = True,
|
||
read_memory: bool = True,
|
||
detach_memory: bool = False,
|
||
return_memory: bool = True,
|
||
memory_text_input_ids: Optional[Tensor] = None,
|
||
memory_text_attention_mask: Optional[Tensor] = None,
|
||
memory_key_input_ids: Optional[Tensor] = None,
|
||
memory_key_attention_mask: Optional[Tensor] = None,
|
||
memory_query_input_ids: Optional[Tensor] = None,
|
||
memory_query_attention_mask: Optional[Tensor] = None,
|
||
memory_storage_input_ids: Optional[Tensor] = None,
|
||
memory_storage_attention_mask: Optional[Tensor] = None,
|
||
force_memory_write: bool = False,
|
||
memory_text: Optional[str] = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
input_ids = kwargs.get("input_ids")
|
||
if input_ids is None and args:
|
||
input_ids = args[0]
|
||
batch_size = input_ids.shape[0] if input_ids is not None else None
|
||
reset_mask = self._reset_mask(input_ids)
|
||
if reset_mask is not None and self.memory_os_v2 is not None:
|
||
# The reset token is a model-level durable-memory clear signal,
|
||
# not merely a transient recurrent-state reset.
|
||
self.clear_hierarchical_memory()
|
||
|
||
if memory_state is not None:
|
||
state = memory_state
|
||
elif self.memory_config.persistent_memory and self.persistent_memory.numel() > 0:
|
||
state = self.persistent_memory
|
||
elif self._persistent_memory is not None:
|
||
state = self._persistent_memory
|
||
elif batch_size is not None:
|
||
state = self.memory.initial_state(batch_size, device=self._find_layer_device())
|
||
else:
|
||
state = None
|
||
if state is not None and batch_size is not None and state.shape[0] != batch_size:
|
||
if self.memory_config.persistent_memory and state.shape[0] == 1:
|
||
state = state.expand(batch_size, -1, -1).clone()
|
||
else:
|
||
raise ValueError(f"memory batch {state.shape[0]} does not match input batch {batch_size}")
|
||
if reset_mask is not None and state is not None:
|
||
state = state.clone()
|
||
state[reset_mask] = 0
|
||
|
||
if batch_size is not None:
|
||
self._bind_text_memory(batch_size, device=self._find_layer_device())
|
||
if reset_mask is not None:
|
||
self._clear_text_memory(reset_mask)
|
||
|
||
self.runtime.state = state
|
||
self.runtime.read_enabled = read_memory
|
||
self.runtime.update_enabled = update_memory
|
||
self.runtime.last_read = None
|
||
self.runtime.reset_mask = reset_mask
|
||
self.runtime.input_ids = input_ids
|
||
self.runtime.attention_mask = kwargs.get("attention_mask")
|
||
kwargs.setdefault("use_cache", False)
|
||
base_output = self.base_model(*args, **kwargs)
|
||
if (
|
||
not self.memory_config.natural_language_memory
|
||
and self.memory_config.direct_logit_scale > 0.0
|
||
and read_memory
|
||
and self.runtime.last_read is not None
|
||
and hasattr(base_output, "logits")
|
||
):
|
||
read = self.runtime.last_read
|
||
output_embeddings = self.base_model.get_output_embeddings()
|
||
memory_logits = output_embeddings(read.to(output_embeddings.weight.device))
|
||
logits = base_output.logits + self.memory_config.direct_logit_scale * memory_logits.to(
|
||
base_output.logits.device
|
||
)
|
||
base_output.logits = logits
|
||
labels = kwargs.get("labels")
|
||
if labels is not None:
|
||
base_output.loss = F.cross_entropy(
|
||
logits[..., :-1, :].contiguous().view(-1, logits.shape[-1]),
|
||
labels[..., 1:].contiguous().view(-1),
|
||
ignore_index=-100,
|
||
)
|
||
if (
|
||
not self.memory_config.natural_language_memory
|
||
and self.memory_config.raw_logit_scale > 0.0
|
||
and read_memory
|
||
and self.runtime.raw_memory is not None
|
||
and hasattr(base_output, "logits")
|
||
):
|
||
output_embeddings = self.base_model.get_output_embeddings()
|
||
raw_logits = output_embeddings(self.runtime.raw_memory.to(output_embeddings.weight.device))
|
||
raw_logits = raw_logits[:, None, :].expand(-1, base_output.logits.shape[1], -1)
|
||
base_output.logits = base_output.logits + self.memory_config.raw_logit_scale * raw_logits.to(
|
||
base_output.logits.device
|
||
)
|
||
labels = kwargs.get("labels")
|
||
if labels is not None:
|
||
base_output.loss = F.cross_entropy(
|
||
base_output.logits[..., :-1, :].contiguous().view(-1, base_output.logits.shape[-1]),
|
||
labels[..., 1:].contiguous().view(-1),
|
||
ignore_index=-100,
|
||
)
|
||
new_state = self.runtime.state
|
||
if new_state is not None and detach_memory:
|
||
new_state = new_state.detach()
|
||
self._persistent_memory = new_state.detach() if new_state is not None else None
|
||
if self.memory_config.persistent_memory and self.runtime.use_persistent_state and new_state is not None:
|
||
self.persistent_memory = new_state.detach()
|
||
if self.memory_config.natural_language_memory and update_memory:
|
||
self.runtime.auto_memory_probability = None
|
||
self.runtime.auto_memory_forget_probability = None
|
||
if (
|
||
self.memory_policy is not None
|
||
and self.memory_config.automatic_memory
|
||
and self._memory_policy_ready
|
||
and isinstance(
|
||
getattr(
|
||
self.memory,
|
||
"last_write_summary",
|
||
getattr(self.memory, "last_write_representation", None),
|
||
),
|
||
Tensor,
|
||
)
|
||
):
|
||
policy_input = getattr(
|
||
self.memory,
|
||
"last_write_summary",
|
||
self.memory.last_write_representation,
|
||
)
|
||
policy_logits = self.memory_policy(policy_input)
|
||
self.runtime.auto_memory_probability = torch.sigmoid(policy_logits).detach()
|
||
if hasattr(self.memory_policy, "forget_logits"):
|
||
forget_logits = self.memory_policy.forget_logits(policy_input)
|
||
self.runtime.auto_memory_forget_probability = torch.sigmoid(forget_logits).detach()
|
||
self._write_text_memory(
|
||
input_ids,
|
||
kwargs.get("attention_mask"),
|
||
memory_text_input_ids,
|
||
memory_text_attention_mask,
|
||
memory_key_input_ids,
|
||
memory_key_attention_mask,
|
||
memory_storage_input_ids,
|
||
memory_storage_attention_mask,
|
||
force_memory_write,
|
||
memory_text,
|
||
)
|
||
self.runtime.state = new_state
|
||
return QwenDynamicMemoryOutput(base_output, new_state) if return_memory else base_output
|
||
|
||
@torch.no_grad()
|
||
def generate(
|
||
self,
|
||
*args: Any,
|
||
memory_state: Optional[Tensor] = None,
|
||
update_memory: bool = True,
|
||
memory_query_input_ids: Optional[Tensor] = None,
|
||
memory_query_attention_mask: Optional[Tensor] = None,
|
||
memory_query_text: Optional[str] = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
"""Use the native HF generation loop with memory-aware decoder layers."""
|
||
input_ids = kwargs.get("input_ids")
|
||
if input_ids is None and args:
|
||
input_ids = args[0]
|
||
self.runtime.context_compaction = None
|
||
reset_mask = self._reset_mask(input_ids)
|
||
if reset_mask is not None and self.memory_os_v2 is not None:
|
||
# Clear before any context compaction so the reset marker can
|
||
# never cause the old bank to be archived again in this call.
|
||
self.clear_hierarchical_memory()
|
||
if input_ids is not None and kwargs.get("attention_mask") is None:
|
||
kwargs["attention_mask"] = torch.ones_like(input_ids)
|
||
if input_ids is not None and reset_mask is None:
|
||
compacted_ids, compacted_mask, compaction_info = self.compact_context_for_kv(
|
||
input_ids,
|
||
kwargs.get("attention_mask"),
|
||
)
|
||
input_ids = compacted_ids
|
||
kwargs["input_ids"] = compacted_ids
|
||
kwargs["attention_mask"] = compacted_mask
|
||
self.runtime.context_compaction = compaction_info
|
||
elif input_ids is not None:
|
||
self.runtime.context_compaction = {
|
||
"compacted": False,
|
||
"reason": "reset_token",
|
||
"archived_records": 0,
|
||
}
|
||
# This is the public prompt returned by generate. When compaction
|
||
# occurred it intentionally contains the bounded hot window, not a
|
||
# second copy of the potentially enormous source prompt.
|
||
original_input_ids = input_ids
|
||
original_prompt_length = int(input_ids.shape[1]) if input_ids is not None else 0
|
||
if memory_state is not None:
|
||
self.runtime.state = memory_state
|
||
elif (
|
||
self.memory_config.persistent_memory
|
||
and self.runtime.use_persistent_state
|
||
and self.persistent_memory.numel() > 0
|
||
):
|
||
self.runtime.state = self.persistent_memory
|
||
elif self._persistent_memory is not None and input_ids is not None and self._persistent_memory.shape[0] == input_ids.shape[0]:
|
||
self.runtime.state = self._persistent_memory
|
||
elif input_ids is not None:
|
||
self.runtime.state = self.memory.initial_state(input_ids.shape[0], device=self._find_layer_device())
|
||
if reset_mask is not None and self.runtime.state is not None:
|
||
state = self.runtime.state.clone()
|
||
if state.shape[0] == 1 and reset_mask.shape[0] > 1:
|
||
state = state.expand(reset_mask.shape[0], -1, -1).clone()
|
||
state[reset_mask] = 0
|
||
self.runtime.state = state
|
||
if input_ids is not None:
|
||
if self.memory_config.natural_language_memory:
|
||
self._bind_text_memory(input_ids.shape[0], device=self._find_layer_device())
|
||
if reset_mask is not None:
|
||
self._clear_text_memory(reset_mask)
|
||
if "attention_mask" not in kwargs or kwargs["attention_mask"] is None:
|
||
kwargs["attention_mask"] = torch.ones_like(input_ids)
|
||
prefix_ids: Optional[Tensor] = None
|
||
prefix_mask: Optional[Tensor] = None
|
||
prefix_length = 0
|
||
retrieval_ids = memory_query_input_ids if memory_query_input_ids is not None else input_ids
|
||
retrieval_mask = (
|
||
memory_query_attention_mask
|
||
if memory_query_attention_mask is not None
|
||
else kwargs["attention_mask"]
|
||
)
|
||
if (
|
||
self.memory_config.natural_language_memory
|
||
and self.runtime.text_slot_valid is not None
|
||
and retrieval_ids is not None
|
||
and retrieval_mask is not None
|
||
):
|
||
prefix_ids, prefix_mask, prefix_length = self._build_text_prefix(
|
||
retrieval_ids.to(device=input_ids.device),
|
||
retrieval_mask.to(device=input_ids.device),
|
||
query_text=memory_query_text or "",
|
||
)
|
||
if prefix_ids is not None and prefix_mask is not None and prefix_length > 0:
|
||
input_ids = torch.cat((prefix_ids, input_ids), dim=1)
|
||
kwargs["input_ids"] = input_ids
|
||
kwargs["attention_mask"] = torch.cat(
|
||
(prefix_mask, kwargs["attention_mask"]), dim=1
|
||
)
|
||
# Internal memory tokens are hidden from the public wrapper
|
||
# API, but remain visible to the actual Qwen forward pass.
|
||
self.runtime.input_ids = input_ids
|
||
# In natural-language mode, the retrieved text prefix is the complete
|
||
# read path. Running the older continuous residual at the same time
|
||
# would distort Qwen's native chat control tokens. With no prefix the
|
||
# memory path is an exact no-op and the base model remains untouched.
|
||
self.runtime.read_enabled = not self.memory_config.natural_language_memory
|
||
self.runtime.update_enabled = update_memory
|
||
self.runtime.last_read = None
|
||
self.runtime.reset_mask = reset_mask
|
||
self.runtime.input_ids = input_ids
|
||
if self.memory_config.kv_offload or self.memory_config.kv_cache_implementation is not None:
|
||
# Transformers 5.x owns the offload lifecycle. Supplying an
|
||
# explicit DynamicCache lets us honor the non-sliding-layer
|
||
# setting, which the shorthand ``cache_implementation`` does not
|
||
# expose for dynamic caches.
|
||
kwargs["use_cache"] = True
|
||
cache_implementation = self.memory_config.kv_cache_implementation
|
||
if cache_implementation is None:
|
||
cache_implementation = "offloaded"
|
||
if (
|
||
cache_implementation == "offloaded"
|
||
and kwargs.get("past_key_values") is None
|
||
):
|
||
try:
|
||
from transformers.cache_utils import DynamicCache, LinearAttentionCacheLayerMixin
|
||
|
||
class _NaturalMemoryOffloadedCache(DynamicCache):
|
||
"""Prefetch the current layer before every cache update.
|
||
|
||
Transformers 5.9 prefetches the next layer as part of
|
||
``Cache.update``. Qwen3.5's hybrid linear/full stack
|
||
can enter the same layer again on the next decoding
|
||
step before that next-layer prefetch runs, leaving a
|
||
CPU tensor to concatenate with CUDA states. The
|
||
current-layer prefetch makes the ownership explicit
|
||
for both attention and recurrent cache updates.
|
||
"""
|
||
|
||
def update(self, key_states, value_states, layer_idx, *args, **extra):
|
||
if self.offloading:
|
||
self.prefetch(layer_idx, self.only_non_sliding)
|
||
return super().update(key_states, value_states, layer_idx, *args, **extra)
|
||
|
||
def update_conv_state(self, conv_states, layer_idx, **extra):
|
||
if self.offloading:
|
||
self.prefetch(layer_idx, self.only_non_sliding)
|
||
value = super().update_conv_state(conv_states, layer_idx, **extra)
|
||
if self.offloading:
|
||
self.offload(layer_idx, self.only_non_sliding)
|
||
return value
|
||
|
||
def update_recurrent_state(self, recurrent_states, layer_idx, **extra):
|
||
if self.offloading:
|
||
self.prefetch(layer_idx, self.only_non_sliding)
|
||
value = super().update_recurrent_state(recurrent_states, layer_idx, **extra)
|
||
if self.offloading:
|
||
self.offload(layer_idx, self.only_non_sliding)
|
||
return value
|
||
|
||
def offload(self, layer_idx, only_non_sliding=True):
|
||
# Qwen3.5 linear-attention state is tiny and is
|
||
# read before update. Keep it on the execution
|
||
# device; offload only the full-attention KV that
|
||
# dominates memory usage.
|
||
if isinstance(self.layers[layer_idx], LinearAttentionCacheLayerMixin):
|
||
return
|
||
return super().offload(layer_idx, only_non_sliding)
|
||
|
||
decoder_config = (
|
||
self.base_model.config.get_text_config(decoder=True)
|
||
if hasattr(self.base_model.config, "get_text_config")
|
||
else getattr(self.base_model.config, "text_config", self.base_model.config)
|
||
)
|
||
kwargs["past_key_values"] = _NaturalMemoryOffloadedCache(
|
||
config=decoder_config,
|
||
offloading=True,
|
||
offload_only_non_sliding=self.memory_config.kv_offload_only_non_sliding,
|
||
)
|
||
except (ImportError, TypeError, AttributeError):
|
||
# Older Transformers releases can still understand the
|
||
# public generation option even if their DynamicCache
|
||
# constructor has a different signature.
|
||
kwargs["cache_implementation"] = cache_implementation
|
||
elif kwargs.get("past_key_values") is None:
|
||
kwargs["cache_implementation"] = cache_implementation
|
||
# Natural-language memory is consumed as an internal text prefix. Do
|
||
# not also add a constant token-level logit bias, because that would
|
||
# repeat the first value token at every generation position.
|
||
if (
|
||
not self.memory_config.natural_language_memory
|
||
and (self.memory_config.direct_logit_scale > 0.0 or self.memory_config.raw_logit_scale > 0.0)
|
||
):
|
||
from transformers import LogitsProcessor, LogitsProcessorList
|
||
|
||
runtime = self.runtime
|
||
output_embeddings = self.base_model.get_output_embeddings()
|
||
prompt_length = int(input_ids.shape[1]) if input_ids is not None else None
|
||
|
||
class MemoryLogitsProcessor(LogitsProcessor):
|
||
def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor:
|
||
first_generated_token = prompt_length is None or input_ids.shape[1] == prompt_length
|
||
if first_generated_token and runtime.raw_memory is not None and raw_scale > 0.0:
|
||
raw = runtime.raw_memory.to(output_embeddings.weight.device)
|
||
scores = scores + raw_scale * output_embeddings(raw).to(scores.device)
|
||
if runtime.last_read is None:
|
||
return scores
|
||
last_read = runtime.last_read[:, -1].to(output_embeddings.weight.device)
|
||
memory_logits = output_embeddings(last_read).to(scores.device)
|
||
return scores + self_scale * memory_logits
|
||
|
||
self_scale = self.memory_config.direct_logit_scale
|
||
raw_scale = self.memory_config.raw_logit_scale
|
||
existing = kwargs.get("logits_processor")
|
||
if existing is None:
|
||
kwargs["logits_processor"] = LogitsProcessorList([MemoryLogitsProcessor()])
|
||
else:
|
||
kwargs["logits_processor"] = LogitsProcessorList(
|
||
list(existing) + [MemoryLogitsProcessor()]
|
||
)
|
||
if input_ids is not None and args:
|
||
args = ()
|
||
generated = self.base_model.generate(*args, **kwargs)
|
||
if (
|
||
original_input_ids is not None
|
||
and input_ids is not None
|
||
and input_ids.shape[1] != original_prompt_length
|
||
):
|
||
generated = torch.cat(
|
||
(
|
||
original_input_ids,
|
||
generated[:, input_ids.shape[1] :],
|
||
),
|
||
dim=1,
|
||
)
|
||
if self.runtime.state is not None:
|
||
self._persistent_memory = self.runtime.state.detach()
|
||
if self.memory_config.persistent_memory and self.runtime.use_persistent_state:
|
||
self.persistent_memory = self._persistent_memory
|
||
return generated
|
||
|
||
def train(self, mode: bool = True) -> "QwenDynamicMemoryModel":
|
||
self.memory.train(mode)
|
||
if self.memory_router_v2 is not None:
|
||
self.memory_router_v2.train(mode)
|
||
self.base_model.eval()
|
||
return self
|
||
|
||
|
||
def load_qwen_base(
|
||
model_path: str | Path,
|
||
*,
|
||
load_in_4bit: bool = True,
|
||
device_map: str | dict[str, Any] = "auto",
|
||
):
|
||
"""Load the unmodified local Qwen checkpoint with the benchmark defaults."""
|
||
|
||
import transformers
|
||
|
||
load_kwargs: dict[str, Any] = {
|
||
"device_map": device_map,
|
||
"low_cpu_mem_usage": True,
|
||
"dtype": torch.bfloat16,
|
||
}
|
||
if load_in_4bit:
|
||
load_kwargs["quantization_config"] = transformers.BitsAndBytesConfig(
|
||
load_in_4bit=True,
|
||
bnb_4bit_quant_type="nf4",
|
||
bnb_4bit_use_double_quant=True,
|
||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||
)
|
||
return transformers.AutoModelForImageTextToText.from_pretrained(str(model_path), **load_kwargs)
|
||
|
||
|
||
def load_qwen_dynamic(
|
||
model_path: str | Path,
|
||
*,
|
||
memory_config: Optional[QwenMemoryConfig] = None,
|
||
load_in_4bit: bool = True,
|
||
device_map: str | dict[str, Any] = "auto",
|
||
freeze_backbone: bool = True,
|
||
) -> QwenDynamicMemoryModel:
|
||
"""Load Qwen and attach either an explicit adapter or an embedded merge package."""
|
||
|
||
model_path = Path(model_path)
|
||
embedded_manifest = model_path / "memory_merge.json"
|
||
if memory_config is None and embedded_manifest.exists():
|
||
embedded_config = model_path / "memory_config.json"
|
||
if not embedded_config.exists():
|
||
raise FileNotFoundError(
|
||
f"embedded memory package is missing {embedded_config}"
|
||
)
|
||
memory_config = load_memory_config(model_path)
|
||
if memory_config is not None and memory_config.memory_storage_mode == "tiered":
|
||
storage_path = Path(memory_config.memory_storage_path) if memory_config.memory_storage_path else None
|
||
if storage_path is None:
|
||
storage_path = model_path / "memory_pages.sqlite"
|
||
elif not storage_path.is_absolute():
|
||
storage_path = model_path / storage_path
|
||
memory_config.memory_storage_path = str(storage_path)
|
||
|
||
base_model = load_qwen_base(
|
||
model_path,
|
||
load_in_4bit=load_in_4bit,
|
||
device_map=device_map,
|
||
)
|
||
model = QwenDynamicMemoryModel(
|
||
base_model,
|
||
memory_config=memory_config,
|
||
freeze_backbone=freeze_backbone,
|
||
)
|
||
if embedded_manifest.exists():
|
||
model.load_embedded_memory_weights(model_path)
|
||
return model
|
||
|
||
|
||
def load_tokenizer(model_path: str | Path):
|
||
import transformers
|
||
|
||
tokenizer = transformers.AutoTokenizer.from_pretrained(str(model_path), use_fast=True)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
return tokenizer
|