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，读写关闭时与原生模型逐位相同
118 lines
4.5 KiB
Python
118 lines
4.5 KiB
Python
"""Fail-fast check that router training can reuse a frozen feature cache.
|
|
|
|
The XL router reuses the exact same frozen Qwen features as the 512-dim
|
|
baseline, so training must never re-encode the corpus. ``_prepare_feature_cache``
|
|
silently loads the 4B model when the cache looks stale, which would cost ~20
|
|
minutes and 9 GiB of VRAM. This tool runs the real cache-validation path with
|
|
the model loader replaced by a hard failure, so a stale cache is reported in
|
|
seconds instead of being discovered halfway through a training launch.
|
|
|
|
Usage (from the fork root, e.g. H:\\Memory\\V2_dpskw)::
|
|
|
|
python -m V2_dpskw.check_router_cache ^
|
|
--train-file data/router_training_v3/train.jsonl ^
|
|
--eval-file data/router_training_v3/eval.jsonl ^
|
|
--feature-cache-dir checkpoints/router_shared/feature_cache ^
|
|
--model-path qwen3_5_4b_natural_memory_v2
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
if __package__ in {None, ""}:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
import V2_dpskw.train_memory_router_large as trainer
|
|
|
|
|
|
def _explode(*_args: object, **_kwargs: object) -> object:
|
|
raise RuntimeError(
|
|
"feature cache is stale: the trainer would now load the 4B Qwen model to "
|
|
"re-encode the corpus. Fix the cache or pass --rebuild-features deliberately."
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--train-file", default="data/router_training_v3/train.jsonl")
|
|
parser.add_argument("--eval-file", default="data/router_training_v3/eval.jsonl")
|
|
parser.add_argument("--feature-cache-dir", default="checkpoints/router_shared/feature_cache")
|
|
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
|
parser.add_argument("--max-key-tokens", type=int, default=256)
|
|
parser.add_argument("--hidden-size", type=int, default=2560)
|
|
args = parser.parse_args()
|
|
|
|
train_path = trainer._resolve_path(args.train_file)
|
|
eval_path = trainer._resolve_path(args.eval_file)
|
|
cache_dir = trainer._resolve_path(args.feature_cache_dir)
|
|
model_path = trainer._resolve_path(args.model_path)
|
|
print(json.dumps({
|
|
"train_file": str(train_path),
|
|
"eval_file": str(eval_path),
|
|
"feature_cache_dir": str(cache_dir),
|
|
"model_path": str(model_path),
|
|
}, ensure_ascii=False, indent=2), flush=True)
|
|
|
|
train = trainer._read_episodes(train_path)
|
|
evaluation = trainer._read_episodes(eval_path)
|
|
texts, _lookup = trainer._collect_texts(train + evaluation)
|
|
expected = {
|
|
"format_version": 1,
|
|
"train_sha256": trainer._sha256(train_path),
|
|
"eval_sha256": trainer._sha256(eval_path),
|
|
"model_path": str(model_path),
|
|
"max_key_tokens": int(args.max_key_tokens),
|
|
"hidden_size": int(args.hidden_size),
|
|
"text_count": len(texts),
|
|
"dtype": "float16_cpu",
|
|
}
|
|
saved_path = trainer._cache_meta_path(cache_dir)
|
|
saved = json.loads(saved_path.read_text(encoding="utf-8")) if saved_path.exists() else {}
|
|
compatible = trainer._cache_is_compatible(cache_dir, expected)
|
|
for key in sorted(expected):
|
|
mark = "==" if saved.get(key) == expected[key] else "!="
|
|
print(f"{mark} {key}: expected={expected[key]!r} saved={saved.get(key)!r}", flush=True)
|
|
if not compatible:
|
|
print("CACHE MISS", flush=True)
|
|
return 1
|
|
|
|
# Run the real code path with the model loader disabled.
|
|
trainer.load_qwen_dynamic = _explode # type: ignore[assignment]
|
|
trainer.load_tokenizer = _explode # type: ignore[assignment]
|
|
vectors, lookup, meta = trainer._prepare_feature_cache(
|
|
argparse.Namespace(
|
|
feature_cache_dir=str(cache_dir),
|
|
model_path=str(model_path),
|
|
max_key_tokens=int(args.max_key_tokens),
|
|
hidden_size=int(args.hidden_size),
|
|
rebuild_features=False,
|
|
precompute_features=True,
|
|
gpu_memory_gb=0.0,
|
|
no_4bit=False,
|
|
encode_batch_size=1,
|
|
precompute_log_every=256,
|
|
),
|
|
train,
|
|
evaluation,
|
|
train_path,
|
|
eval_path,
|
|
)
|
|
print(json.dumps({
|
|
"verdict": "CACHE HIT",
|
|
"feature_shape": list(vectors.shape),
|
|
"dtype": str(vectors.dtype),
|
|
"unique_texts": len(lookup),
|
|
"train_episodes": len(train),
|
|
"eval_episodes": len(evaluation),
|
|
"meta": meta,
|
|
}, ensure_ascii=False, indent=2), flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|