WpyQwq
643e22ecb9
- 引入 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，读写关闭时与原生模型逐位相同
329 lines
14 KiB
Python
329 lines
14 KiB
Python
"""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())
|