wpy/natural-memory-nm21
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据

Browse Source 
- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线
- 修复 prepare_memory_router_dataset 候选池重建缺陷（mega 家族 3568x 加速，输出逐字节相同）
- 修复 v5 被破坏的拒答与多跳标签（train 未知样本 319 -> 16319，multi_hop 平均正例 1.00 -> 2.00）
- 同存储预算下 V2-128 v6 逐轴 22/22 通过：Top-1 41.12% -> 94.62%，未知拒答 0.00% -> 100.00%
- 记录三条被实测推翻的显然优化（logits_to_keep=1 反而慢 55%、XL 容量未带来收益）
- 记忆手术跨架构可移植性 14/14，读写关闭时与原生模型逐位相同
This commit is contained in:
WpyQwq
2026-09-19 11:11:31 +08:00
commit 643e22ecb9
484 changed files with 306821 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
probe_nm2_portability.py
+328
View File
@@ -0,0 +1,328 @@
"""Measure how far the NM2 memory surgery generalizes across model architectures.
The NM2 embedding is not architecture-neutral. Reading the integration code
shows four concrete assumptions, and this tool tests each one on real checkpoints
instead of arguing about them:
1. ``base_model.config.text_config`` must exist (line 1073);
2. the backbone must be loadable as a conditional-generation / image-text model
(line 4000);
3. the decoder layers must be reachable as
``base_model.model.language_model.layers`` (lines 1155/2393/3961/3970);
4. ``MemoryLayerAdapter.forward`` declares ``position_embeddings`` as its **second
positional parameter** and forwards it positionally to the wrapped layer
(lines 960-999), which only matches layers whose own second positional
parameter is ``position_embeddings`` (Qwen3.5/Gemma-style), not Llama/Qwen2
style layers that expect ``attention_mask`` there.
For every model the probe reports which stages pass, applies the *minimal* shims
(``config.text_config = config`` and ``model.model.language_model = model.model``)
to see how much work a port actually needs, and then runs the strongest possible
capability-preservation check: with memory read/write disabled, the wrapped model
must reproduce the original model's logits **exactly**.
Usage::
python -m V2_dpskw.probe_nm2_portability --models list.json --output nm2_portability.json
python -m V2_dpskw.probe_nm2_portability # built-in architecture list
"""
from __future__ import annotations
import argparse
import inspect
import json
import os
import sys
import traceback
from pathlib import Path
from typing import Any
import torch
if __package__ in {None, ""}:
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from V2_dpskw.qwen_integration import QwenDynamicMemoryModel, QwenMemoryConfig
#: Layer-container attribute paths tried, in the order the integration would need.
CONTAINER_PATHS = (
"model.language_model.layers",
"model.layers",
"language_model.layers",
"layers",
"transformer.h",
"gpt_neox.layers",
"model.decoder.layers",
"decoder.layers",
"model.layers.0.model.layers", # multimodal nesting
)
DEFAULT_MODELS = [
"hf-internal-testing/tiny-random-LlamaForCausalLM",
"hf-internal-testing/tiny-random-MistralForCausalLM",
"hf-internal-testing/tiny-random-Qwen2ForCausalLM",
"hf-internal-testing/tiny-random-Qwen3ForCausalLM",
"hf-internal-testing/tiny-random-Phi3ForCausalLM",
"hf-internal-testing/tiny-random-Gemma2ForCausalLM",
"hf-internal-testing/tiny-random-Gemma3ForCausalLM",
"hf-internal-testing/tiny-random-Starcoder2ForCausalLM",
"hf-internal-testing/tiny-random-OlmoeForCausalLM",
"hf-internal-testing/tiny-random-GraniteMoeForCausalLM",
"hf-internal-testing/tiny-random-MixtralForCausalLM",
"hf-internal-testing/tiny-random-CohereForCausalLM",
"hf-internal-testing/tiny-random-FalconForCausalLM",
"akreal/tiny-random-BloomForCausalLM",
"hf-tiny-model-private/tiny-random-OPTForCausalLM",
"hf-tiny-model-private/tiny-random-GPTNeoXForCausalLM",
"hf-tiny-model-private/tiny-random-GPT2LMHeadModel",
"hf-tiny-model-private/tiny-random-BartForCausalLM",
"trl-internal-testing/tiny-random-LlamaForCausalLM",
]
def _resolve_path(root: Any, path: str) -> Any:
current = root
for part in path.split("."):
if current is None:
return None
if part.isdigit():
current = current[int(part)] if hasattr(current, "__getitem__") else None
else:
current = getattr(current, part, None)
return current
def _find_container(model: Any) -> tuple[str | None, Any]:
for path in CONTAINER_PATHS:
value = _resolve_path(model, path)
if value is not None and hasattr(value, "__len__") and len(value) > 0 and hasattr(value[0], "forward"):
return path, value
return None, None
def _minimal_config(layer_indices: tuple[int, ...], *, full_stack: bool) -> QwenMemoryConfig:
if not full_stack:
return QwenMemoryConfig(
memory_slots=4,
memory_dim=64,
layer_indices=list(layer_indices),
mode="residual",
blend_init=0.0,
native_mode=False,
persistent_memory=False,
summary_pooling=False,
natural_language_memory=False,
automatic_memory=False,
memory_version=1,
hierarchical_memory=False,
kv_offload=False,
auto_compact_context=False,
)
return QwenMemoryConfig(
memory_slots=8,
memory_dim=128,
layer_indices=list(layer_indices),
mode="residual",
blend_init=0.0,
native_mode=True,
persistent_memory=True,
summary_pooling=True,
natural_language_memory=True,
automatic_memory=True,
automatic_memory_policy_version=2,
memory_version=2,
hierarchical_memory=True,
kv_offload=False,
auto_compact_context=False,
)
def probe(model_id: str, *, full_stack: bool = False, dtype: str = "float32") -> dict[str, Any]:
record: dict[str, Any] = {"model": model_id, "full_stack": full_stack, "dtype": dtype, "stages": {}, "patches": []}
stages = record["stages"]
torch_dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}.get(dtype, torch.float32)
def fail(stage: str, exc: BaseException) -> None:
stages[stage] = f"FAIL {type(exc).__name__}: {str(exc).splitlines()[0][:180]}"
# ---- 1. config ---------------------------------------------------------
try:
config = AutoConfig.from_pretrained(model_id)
text_config = getattr(config, "text_config", None)
record["architecture"] = (getattr(config, "architectures", None) or ["?"])[0]
record["model_type"] = getattr(config, "model_type", "?")
record["hidden_size"] = getattr(text_config or config, "hidden_size", None)
record["num_layers"] = getattr(text_config or config, "num_hidden_layers", None)
record["vocab_size"] = getattr(text_config or config, "vocab_size", None)
record["text_config_present"] = text_config is not None
stages["1_config"] = (
"ok (config.text_config)" if text_config is not None
else "resolved by integration (text-only model has no text_config)"
)
except BaseException as exc: # noqa: BLE001
fail("1_config", exc)
return record
# ---- 2. load backbone --------------------------------------------------
try:
try:
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch_dtype)
except (ValueError, KeyError, OSError) as exc:
# Multimodal checkpoints (Qwen3.5, Gemma3-VL) are image-text-to-text
# models, which is what the integration's own loader uses.
from transformers import AutoModelForImageTextToText
record["causal_lm_fallback"] = str(exc).splitlines()[0][:120]
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch_dtype)
model.eval()
record["backbone_class"] = type(model).__name__
stages["2_load_causal_lm"] = "ok"
except BaseException as exc: # noqa: BLE001
fail("2_load_causal_lm", exc)
return record
# ---- 3. layer container -------------------------------------------------
container_path, container = _find_container(model)
record["container_path"] = container_path
record["repo_layer_path_ok"] = container_path == "model.language_model.layers"
stages["3_container"] = (
"ok (repo path model.language_model.layers)" if container_path == "model.language_model.layers"
else f"resolved by integration (found {container_path})"
)
if container is None:
return record
# ---- 4. layer calling convention ---------------------------------------
layer_signature = list(inspect.signature(container[0].forward).parameters)
record["layer_forward_params"] = layer_signature
second = layer_signature[1] if len(layer_signature) > 1 else None
record["second_positional_is_position_embeddings"] = second == "position_embeddings"
stages["4_convention"] = (
"ok (position_embeddings second -> positional call)"
if second == "position_embeddings"
else f"handled by declared-argument form (second positional is {second!r})"
)
# ---- reference output BEFORE any surgery -------------------------------
try:
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("记住：项目代号是 P-123。", return_tensors="pt")
except BaseException:
inputs = {"input_ids": torch.randint(1, 128, (1, 8)), "attention_mask": torch.ones(1, 8, dtype=torch.long)}
with torch.no_grad():
reference = model(**inputs).logits.clone()
# ---- 5. surgery (no shims; the integration resolves architecture itself) --
num_layers = int(record["num_layers"] or len(container))
try:
# Tiny probe models can have very few layers, so mirror the integration's
# own stride rule but clamp it to a valid range for this depth.
stride = max(1, num_layers // 4)
layer_indices = tuple(sorted({max(0, min(num_layers - 1, stride * i - 1)) for i in range(1, 5)}))
record["resolved_layer_indices"] = list(layer_indices)
memory_config = _minimal_config(layer_indices, full_stack=full_stack)
# No shims: the integration resolves text_config and the decoder stack
# itself now, so a successful surgery here means a genuinely portable port.
wrapped = QwenDynamicMemoryModel(model, memory_config, freeze_backbone=True)
record["shims_needed"] = []
stages["5_surgery"] = "ok (no shims)"
except BaseException as exc: # noqa: BLE001
wrapped_layers = [
index for index, layer in enumerate(container)
if type(layer).__name__ == "MemoryLayerAdapter"
]
record["partially_wrapped_layers"] = wrapped_layers
stages["5_surgery"] = f"FAIL {type(exc).__name__}: {str(exc).splitlines()[0][:180]}"
record["traceback"] = traceback.format_exc()[-800:]
return record
# ---- 6. identity check (the capability-preservation proof) -------------
try:
with torch.no_grad():
output = wrapped(**inputs, read_memory=False, update_memory=False)
logits = output.logits if hasattr(output, "logits") else output[0]
delta = float((logits - reference).abs().max())
record["identity_max_abs_logit_delta"] = delta
stages["6_identity"] = "identical" if delta == 0.0 else f"DIFFERS by {delta:.3e}"
except BaseException as exc: # noqa: BLE001
stages["6_identity"] = f"FAIL {type(exc).__name__}: {str(exc).splitlines()[0][:180]}"
record["traceback"] = traceback.format_exc()[-800:]
return record
# ---- 7. memory write/read smoke ---------------------------------------
try:
with torch.no_grad():
written = wrapped(**inputs, read_memory=False, update_memory=True, return_memory=True)
state = getattr(written, "memory_state", None)
changed = None
if state is not None and getattr(wrapped, "persistent_memory", None) is not None:
changed = float((state - wrapped.persistent_memory).abs().max())
record["memory_state_changed_by"] = changed
with torch.no_grad():
wrapped(**inputs, read_memory=True, update_memory=False)
stages["7_write_read"] = "ok" if changed is None or changed > 0 else "state unchanged"
except BaseException as exc: # noqa: BLE001
stages["7_write_read"] = f"FAIL {type(exc).__name__}: {str(exc).splitlines()[0][:180]}"
return record
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--models", default="", help="json file with a list of model ids")
parser.add_argument("--model", action="append", default=None)
parser.add_argument("--full-stack", action="store_true", help="also exercise the V2/native memory stack")
parser.add_argument("--dtype", default="float32", choices=("float32", "bfloat16", "float16"))
parser.add_argument("--output", default="nm2_portability.json")
parser.add_argument("--markdown", default="")
args = parser.parse_args()
if args.model:
models = args.model
elif args.models:
models = json.loads(Path(args.models).read_text(encoding="utf-8"))
else:
models = DEFAULT_MODELS
records = []
for model_id in models:
print(f"--- {model_id}", flush=True)
try:
record = probe(model_id, full_stack=args.full_stack, dtype=args.dtype)
except BaseException as exc: # noqa: BLE001
record = {"model": model_id, "stages": {"0_probe": f"FAIL {type(exc).__name__}: {exc}"}}
records.append(record)
print(json.dumps(record, ensure_ascii=False)[:400], flush=True)
Path(args.output).write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")
if args.markdown:
lines = [
"| model | model_type | hidden | layers | container | 2nd positional | identity | surgery |",
"|---|---|---|---|---|---|---|---|",
]
for r in records:
stages = r.get("stages", {})
lines.append("| {m} | {t} | {h} | {l} | {c} | {s} | {i} | {su} |".format(
m=r.get("model", "?").split("/")[-1],
t=r.get("model_type", "-"),
h=r.get("hidden_size", "-"),
l=r.get("num_layers", "-"),
c=r.get("container_path", "-"),
s="position_embeddings" if r.get("second_positional_is_position_embeddings") else "OTHER",
i=stages.get("6_identity", "-"),
su=stages.get("5_surgery", "-"),
))
Path(args.markdown).write_text("\n".join(lines) + "\n", encoding="utf-8")
print("\n".join(lines), flush=True)
print(f"wrote {args.output}", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())