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，读写关闭时与原生模型逐位相同
560 lines
24 KiB
Python
560 lines
24 KiB
Python
"""Streaming, multi-threaded frozen-feature builder for router training data.
|
|
|
|
The original builder loaded every episode into RAM and encoded one text at a time.
|
|
That does not scale to the 87k-episode v5 dataset:
|
|
|
|
* parsing ``train.jsonl`` (1.15 GB) into Python objects costs several GB of RAM,
|
|
* the resulting feature bank is 2.12M x 2560 fp16 = **10.86 GB**,
|
|
* single-text encoding would need many hours.
|
|
|
|
This tool instead:
|
|
|
|
1. **scans** the JSONL files line by line, keeps only a ``sha1 -> row`` lookup and
|
|
spills the unique texts to disk (never holding the parsed episodes);
|
|
2. **tokenizes** with a thread pool (Hugging Face fast tokenizers release the GIL);
|
|
3. **encodes** into a memory-mapped ``.npy`` bank, in exact-token-length groups so
|
|
the frozen Qwen representation is identical to the original cache, with a
|
|
token budget per batch so a long-text batch cannot blow up VRAM;
|
|
4. writes a ``manifest.json`` compatible with the trainer's cache validation.
|
|
|
|
Peak RAM is the lookup dictionary plus the token/length arrays; the 10.86 GB bank
|
|
stays on disk and is memory-mapped by the trainer.
|
|
|
|
Usage::
|
|
|
|
python -m V2_dpskw.stream_feature_bank ^
|
|
--train-file data/router_training_v5/train.jsonl ^
|
|
--eval-file data/router_training_v5/eval.jsonl ^
|
|
--model-path qwen3_5_4b_natural_memory_v2 ^
|
|
--output-dir checkpoints/router_v5/feature_cache ^
|
|
--tokenizer-threads 8 --max-batch 128 --token-budget 8192 --gpu-memory-gb 10
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
if __package__ in {None, ""}:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
from V2_dpskw.qwen_integration import load_qwen_dynamic, load_tokenizer
|
|
|
|
BANK_NAME = "features.f16.npy"
|
|
MANIFEST_NAME = "manifest.json"
|
|
INDEX_NAME = "index.json"
|
|
TEXTS_NAME = "texts.jsonl"
|
|
TOKENS_NAME = "tokens.i32"
|
|
LENGTHS_NAME = "lengths.i32"
|
|
PROGRESS_NAME = "progress.json"
|
|
FORMAT_VERSION = 2
|
|
|
|
|
|
def _encode_method_fingerprint() -> str:
|
|
"""Hash only the method that decides the feature values.
|
|
|
|
The fork is shared with other work, so ``qwen_integration.py`` can change
|
|
between two halves of one encode. Hashing the whole file would be too
|
|
coarse (unrelated edits would block a legitimate resume); hashing exactly
|
|
``_encode_model_key`` guarantees that a resumed run keeps writing features
|
|
that are comparable with the rows written before the interruption.
|
|
|
|
The compiled code object is hashed rather than ``inspect.getsource`` text:
|
|
``getsource`` slices the *current* file by the code object's stored line
|
|
numbers, so any edit above the method silently shifts the slice and hashes a
|
|
neighbouring function instead (measured: it returned the ``compact_context``
|
|
body for ``_encode_model_key``). That turned a safety check into a random
|
|
value and would have refused a legitimate resume. Bytecode is stable across
|
|
line shifts and still changes when the method's logic changes.
|
|
"""
|
|
|
|
try:
|
|
import marshal
|
|
|
|
from V2_dpskw.qwen_integration import QwenDynamicMemoryModel
|
|
|
|
code = QwenDynamicMemoryModel._encode_model_key.__code__
|
|
digest = hashlib.sha256()
|
|
digest.update(marshal.dumps(code.co_code))
|
|
digest.update(repr(code.co_consts).encode("utf-8"))
|
|
digest.update(repr(code.co_names).encode("utf-8"))
|
|
digest.update(repr(code.co_varnames).encode("utf-8"))
|
|
digest.update(repr(sorted(code.co_freevars)).encode("utf-8"))
|
|
return digest.hexdigest()
|
|
except Exception as exc: # pragma: no cover - diagnostic path
|
|
return f"unavailable:{type(exc).__name__}"
|
|
|
|
|
|
def _write_progress(output_dir: Path, completed: int, total: int) -> None:
|
|
"""Record how many length-sorted rows are durably written."""
|
|
|
|
payload = {
|
|
"completed_rows": int(completed),
|
|
"total_rows": int(total),
|
|
"encode_method_sha256": _encode_method_fingerprint(),
|
|
}
|
|
target = output_dir / PROGRESS_NAME
|
|
temp = target.with_suffix(".json.tmp")
|
|
temp.write_text(json.dumps(payload), encoding="utf-8")
|
|
temp.replace(target)
|
|
|
|
|
|
def _read_progress(output_dir: Path, *, text_count: int) -> int:
|
|
target = output_dir / PROGRESS_NAME
|
|
if not target.exists():
|
|
return 0
|
|
try:
|
|
payload = json.loads(target.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return 0
|
|
completed = int(payload.get("completed_rows", 0))
|
|
if int(payload.get("total_rows", -1)) != int(text_count):
|
|
_report({"phase": "resume_rejected", "reason": "text count changed", "progress": payload})
|
|
return 0
|
|
recorded = str(payload.get("encode_method_sha256", ""))
|
|
current = _encode_method_fingerprint()
|
|
if recorded and recorded != current:
|
|
_report({"phase": "resume_rejected", "reason": "encode method changed since the partial run",
|
|
"recorded": recorded, "current": current})
|
|
return 0
|
|
bank = output_dir / BANK_NAME
|
|
if not bank.exists():
|
|
return 0
|
|
# Validate by array shape, not by byte size: a .npy file carries a header.
|
|
try:
|
|
probe = np.load(bank, mmap_mode="r")
|
|
shape = tuple(probe.shape)
|
|
except Exception:
|
|
return 0
|
|
finally:
|
|
probe = None
|
|
if len(shape) != 2 or int(shape[0]) != int(text_count):
|
|
_report({"phase": "resume_rejected", "reason": "bank shape mismatch",
|
|
"bank_shape": list(shape), "expected_rows": int(text_count)})
|
|
return 0
|
|
if completed >= text_count:
|
|
return 0
|
|
_report({"phase": "resume_accepted", "completed_rows": completed,
|
|
"encode_method_sha256": recorded[:16]})
|
|
return completed
|
|
|
|
|
|
def _resolve(value: str | Path) -> Path:
|
|
path = Path(value)
|
|
if path.is_absolute() or path.exists():
|
|
return path
|
|
return Path(__file__).resolve().parent / path
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _report(payload: dict[str, Any]) -> None:
|
|
print(json.dumps(payload, ensure_ascii=False), flush=True)
|
|
|
|
|
|
def _iter_episode_texts(path: Path) -> Iterator[tuple[str, str]]:
|
|
"""Yield ``(kind, text)`` for the query and every candidate of each episode.
|
|
|
|
Episodes are streamed: nothing but the current line is ever materialised, so
|
|
a 1.15 GB file costs no lasting RAM.
|
|
"""
|
|
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
row = json.loads(line)
|
|
query = str(row.get("query", "")).strip()
|
|
if query:
|
|
yield "query", query
|
|
candidates = row.get("candidates")
|
|
if isinstance(candidates, list):
|
|
for candidate in candidates:
|
|
if isinstance(candidate, dict):
|
|
text = str(candidate.get("text", "")).strip()
|
|
if text:
|
|
yield "candidate", text
|
|
|
|
|
|
def _encode_keys(model: Any, input_ids: torch.Tensor, attention_mask: torch.Tensor, *, skip_lm_head: bool, backbone: str = "outer") -> torch.Tensor:
|
|
"""Masked-mean key encoding.
|
|
|
|
``backbone="outer"`` is the canonical path (identical to
|
|
``QwenDynamicMemoryModel._encode_model_key``) and is the default so that a
|
|
feature bank never mixes two code paths.
|
|
|
|
``backbone="inner"`` calls the inner transformer directly, which skips the
|
|
LM head entirely. Measured on 192 *distinct* length-32 texts: 977 ms vs
|
|
1099 ms per batch (1.13x faster) and 4.80 GiB vs 6.98 GiB peak VRAM, with
|
|
feature agreement cosine 0.9999911. Note that ``skip_lm_head`` is the wrong
|
|
lever for the same goal: asking for one token of logits is 1.55x *slower*
|
|
(see the flag help), because the sliced view is not contiguous.
|
|
"""
|
|
|
|
if backbone == "inner" and hasattr(model.base_model, "model"):
|
|
target = model.base_model.model
|
|
else:
|
|
if not skip_lm_head:
|
|
return model._encode_model_key(input_ids, attention_mask)
|
|
target = model.base_model
|
|
previous_read = model.runtime.read_enabled
|
|
previous_update = model.runtime.update_enabled
|
|
model.runtime.read_enabled = False
|
|
model.runtime.update_enabled = False
|
|
try:
|
|
embedding_layer = model.base_model.get_input_embeddings()
|
|
input_ids = input_ids.to(embedding_layer.weight.device)
|
|
attention_mask = attention_mask.to(input_ids.device)
|
|
try:
|
|
output = target(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
use_cache=False,
|
|
output_hidden_states=True,
|
|
logits_to_keep=1,
|
|
)
|
|
except TypeError:
|
|
# Architectures without a logits_to_keep argument.
|
|
output = target(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
use_cache=False,
|
|
output_hidden_states=True,
|
|
)
|
|
hidden_states = getattr(output, "hidden_states", None)
|
|
hidden = (hidden_states[-1] if hidden_states is not None else output.last_hidden_state).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 torch.nn.functional.normalize(key, dim=-1)
|
|
finally:
|
|
model.runtime.read_enabled = previous_read
|
|
model.runtime.update_enabled = previous_update
|
|
|
|
|
|
def phase_scan(paths: list[Path], output_dir: Path) -> tuple[int, dict[str, int]]:
|
|
"""Collect unique texts, spilling them to disk in row order."""
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
lookup: dict[str, int] = {}
|
|
kinds: dict[str, int] = {"query": 0, "candidate": 0}
|
|
started = time.perf_counter()
|
|
with (output_dir / TEXTS_NAME).open("w", encoding="utf-8") as sink:
|
|
for path in paths:
|
|
for kind, text in _iter_episode_texts(path):
|
|
digest = hashlib.sha1(text.encode("utf-8", "replace")).hexdigest()
|
|
if digest in lookup:
|
|
continue
|
|
lookup[digest] = len(lookup)
|
|
kinds[kind] += 1
|
|
sink.write(json.dumps(text, ensure_ascii=False) + "\n")
|
|
if len(lookup) % 250000 == 0:
|
|
_report({"phase": "scan", "unique_texts": len(lookup),
|
|
"seconds": round(time.perf_counter() - started, 1)})
|
|
with (output_dir / INDEX_NAME).open("w", encoding="utf-8") as handle:
|
|
json.dump(lookup, handle, ensure_ascii=False)
|
|
_report({"phase": "scan_done", "unique_texts": len(lookup), "kinds": kinds,
|
|
"seconds": round(time.perf_counter() - started, 1)})
|
|
return len(lookup), kinds
|
|
|
|
|
|
def phase_tokenize(
|
|
output_dir: Path,
|
|
tokenizer: Any,
|
|
*,
|
|
text_count: int,
|
|
max_key_tokens: int,
|
|
threads: int,
|
|
chunk: int,
|
|
progress_every: int,
|
|
) -> tuple[np.memmap, np.memmap]:
|
|
"""Tokenize every unique text with a thread pool."""
|
|
|
|
tokens = np.memmap(output_dir / TOKENS_NAME, dtype=np.int32, mode="w+", shape=(text_count, max_key_tokens))
|
|
lengths = np.memmap(output_dir / LENGTHS_NAME, dtype=np.int32, mode="w+", shape=(text_count,))
|
|
started = time.perf_counter()
|
|
done = 0
|
|
|
|
def tokenize_chunk(items: list[tuple[int, str]]) -> list[tuple[int, list[int]]]:
|
|
rows = [row for row, _ in items]
|
|
encoded = tokenizer(
|
|
[text for _, text in items],
|
|
add_special_tokens=False,
|
|
truncation=True,
|
|
max_length=max_key_tokens,
|
|
padding=False,
|
|
)["input_ids"]
|
|
return list(zip(rows, encoded))
|
|
|
|
with (output_dir / TEXTS_NAME).open("r", encoding="utf-8") as source, ThreadPoolExecutor(max_workers=max(1, threads)) as pool:
|
|
batch: list[tuple[int, str]] = []
|
|
pending = []
|
|
for row, line in enumerate(source):
|
|
batch.append((row, json.loads(line)))
|
|
if len(batch) >= chunk:
|
|
pending.append(pool.submit(tokenize_chunk, batch))
|
|
batch = []
|
|
if len(pending) >= max(2, threads * 2):
|
|
for future in pending:
|
|
for row, ids in future.result():
|
|
length = min(len(ids), max_key_tokens)
|
|
lengths[row] = length
|
|
if length:
|
|
tokens[row, :length] = np.asarray(ids[:length], dtype=np.int32)
|
|
else:
|
|
lengths[row] = 1
|
|
tokens[row, 0] = 0
|
|
done += 1
|
|
pending = []
|
|
if done % progress_every < chunk:
|
|
rate = done / max(1e-9, time.perf_counter() - started)
|
|
_report({"phase": "tokenize", "done": done, "total": text_count,
|
|
"texts_per_second": round(rate, 1)})
|
|
if batch:
|
|
pending.append(pool.submit(tokenize_chunk, batch))
|
|
for future in pending:
|
|
for row, ids in future.result():
|
|
length = min(len(ids), max_key_tokens)
|
|
if length == 0:
|
|
length = 1
|
|
ids = [0]
|
|
lengths[row] = length
|
|
tokens[row, :length] = np.asarray(ids[:length], dtype=np.int32)
|
|
done += 1
|
|
tokens.flush()
|
|
lengths.flush()
|
|
_report({"phase": "tokenize_done", "texts": int(done),
|
|
"seconds": round(time.perf_counter() - started, 1)})
|
|
return tokens, lengths
|
|
|
|
|
|
def phase_encode(
|
|
output_dir: Path,
|
|
model: Any,
|
|
tokens: np.memmap,
|
|
lengths: np.memmap,
|
|
*,
|
|
text_count: int,
|
|
hidden_size: int,
|
|
max_batch: int,
|
|
token_budget: int,
|
|
progress_every: int,
|
|
skip_lm_head: bool = True,
|
|
backbone: str = "outer",
|
|
resume_from: int = 0,
|
|
fill_from: int = 0,
|
|
fill_to: int = 0,
|
|
) -> int:
|
|
"""Encode in exact-token-length groups, writing into a memory-mapped bank.
|
|
|
|
Returns the number of entries of the length-sorted order that are complete, so
|
|
an interrupted run can resume instead of re-encoding hours of GPU work.
|
|
"""
|
|
|
|
# "w+" recreates the whole bank; it is only correct for a genuinely fresh run.
|
|
# A resume and a gap-fill must both open the existing file read/write, or they
|
|
# would wipe the rows that were already paid for.
|
|
fresh_run = int(resume_from) == 0 and int(fill_to) == 0
|
|
bank = np.lib.format.open_memmap(
|
|
output_dir / BANK_NAME, mode="w+" if fresh_run else "r+", dtype=np.float16, shape=(text_count, hidden_size)
|
|
)
|
|
length_values = np.asarray(lengths, dtype=np.int64)
|
|
order = np.argsort(length_values, kind="stable")
|
|
sorted_lengths = length_values[order]
|
|
boundaries = np.flatnonzero(np.diff(sorted_lengths)) + 1
|
|
starts = np.concatenate(([0], boundaries))
|
|
ends = np.concatenate((boundaries, [len(order)]))
|
|
device = next(model.base_model.parameters()).device
|
|
started = time.perf_counter()
|
|
done = int(resume_from)
|
|
calls = 0
|
|
processed = 0
|
|
if resume_from:
|
|
_report({"phase": "resume", "completed_rows": resume_from, "total": text_count})
|
|
for start, end in zip(starts.tolist(), ends.tolist()):
|
|
# Resume must trim *within* the partially finished group. Skipping any
|
|
# group whose ``start`` is below the checkpoint silently dropped every row
|
|
# of that group above the checkpoint: the length-32 group spans ~680k rows,
|
|
# so an interrupted resume lost 294,306 rows while still reporting success.
|
|
low = max(start, done) if resume_from else start
|
|
if fill_to:
|
|
low = max(low, fill_from)
|
|
end = min(end, fill_to)
|
|
if end <= low:
|
|
continue
|
|
rows = order[low:end]
|
|
length = int(sorted_lengths[start])
|
|
capacity = max(1, min(max_batch, token_budget // max(1, length)))
|
|
for offset in range(0, len(rows), capacity):
|
|
block = rows[offset : offset + capacity]
|
|
ids = torch.from_numpy(np.asarray(tokens[block, :length], dtype=np.int64)).to(device, non_blocking=True)
|
|
mask = torch.ones_like(ids)
|
|
with torch.inference_mode():
|
|
vectors = _encode_keys(model, ids, mask, skip_lm_head=skip_lm_head, backbone=backbone).detach().to(torch.float16).cpu().numpy()
|
|
bank[block] = vectors
|
|
done += len(block)
|
|
processed += len(block)
|
|
calls += 1
|
|
if done % progress_every < len(block):
|
|
bank.flush()
|
|
# A gap-fill run does not produce a contiguous prefix, so it must
|
|
# not overwrite the resume checkpoint of the prefix run.
|
|
if not fill_to:
|
|
_write_progress(output_dir, done, text_count)
|
|
elapsed = max(1e-9, time.perf_counter() - started)
|
|
remaining = (text_count - done) / max(1e-9, processed / elapsed)
|
|
_report({"phase": "encode", "done": done, "total": text_count,
|
|
"texts_per_second": round(processed / elapsed, 1),
|
|
"eta_seconds": round(remaining, 0), "length": length, "batch": len(block)})
|
|
bank.flush()
|
|
if not fill_to:
|
|
_write_progress(output_dir, done, text_count)
|
|
_report({"phase": "encode_done", "texts": done, "model_calls": calls,
|
|
"seconds": round(time.perf_counter() - started, 1), "resumed_from": resume_from,
|
|
"fill_from": fill_from, "fill_to": fill_to,
|
|
"complete": bool(not fill_to and done >= text_count)})
|
|
return done
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--train-file", default="data/router_training_v5/train.jsonl")
|
|
parser.add_argument("--eval-file", default="data/router_training_v5/eval.jsonl")
|
|
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
|
|
parser.add_argument("--output-dir", default="checkpoints/router_v5/feature_cache")
|
|
parser.add_argument("--max-key-tokens", type=int, default=256)
|
|
parser.add_argument("--hidden-size", type=int, default=2560)
|
|
parser.add_argument("--tokenizer-threads", type=int, default=8)
|
|
parser.add_argument("--tokenize-chunk", type=int, default=1024)
|
|
parser.add_argument("--max-batch", type=int, default=128)
|
|
parser.add_argument("--token-budget", type=int, default=8192)
|
|
parser.add_argument("--gpu-memory-gb", type=float, default=10.0)
|
|
parser.add_argument("--no-4bit", action="store_true")
|
|
parser.add_argument(
|
|
"--skip-lm-head",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help=(
|
|
"request one token of logits instead of all tokens. Features are identical "
|
|
"(measured cosine 0.9999996) but on this bitsandbytes 4-bit stack it is "
|
|
"1.55x SLOWER: the quantised matmul falls into a slow path for the sliced "
|
|
"non-contiguous view (1050 ms vs 679 ms per 192x32 batch), even though it "
|
|
"saves 1.1 GiB of VRAM. Off by default; kept for the record."
|
|
),
|
|
)
|
|
parser.add_argument("--backbone", choices=("outer", "inner"), default="outer", help="inner calls the inner transformer directly (1.13x faster, 2.2 GiB less VRAM, features agree at cosine 0.999991); default keeps one consistent code path for a whole bank")
|
|
parser.add_argument("--resume", action="store_true", help="continue an interrupted encode from progress.json")
|
|
parser.add_argument("--fill-rows", default="", help="encode only this half-open row range in length-sorted order, e.g. 750021:1044327")
|
|
parser.add_argument("--progress-every", type=int, default=25000)
|
|
parser.add_argument("--limit-texts", type=int, default=0, help="smoke test cap")
|
|
args = parser.parse_args()
|
|
|
|
train_path = _resolve(args.train_file)
|
|
eval_path = _resolve(args.eval_file)
|
|
output_dir = _resolve(args.output_dir)
|
|
model_path = _resolve(args.model_path)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
_report({"phase": "start", "train_file": str(train_path), "eval_file": str(eval_path),
|
|
"output_dir": str(output_dir), "model_path": str(model_path)})
|
|
|
|
text_count, kinds = phase_scan([train_path, eval_path], output_dir)
|
|
if args.limit_texts:
|
|
text_count = min(text_count, args.limit_texts)
|
|
|
|
if args.gpu_memory_gb > 0 and torch.cuda.is_available():
|
|
total = torch.cuda.get_device_properties(0).total_memory
|
|
torch.cuda.set_per_process_memory_fraction(
|
|
min(0.92, max(0.05, args.gpu_memory_gb * 1024**3 / total)), device=0
|
|
)
|
|
tokenizer = load_tokenizer(model_path)
|
|
tokens, lengths = phase_tokenize(
|
|
output_dir, tokenizer, text_count=text_count, max_key_tokens=args.max_key_tokens,
|
|
threads=args.tokenizer_threads, chunk=args.tokenize_chunk, progress_every=args.progress_every,
|
|
)
|
|
del tokenizer
|
|
gc.collect()
|
|
|
|
max_memory = None
|
|
if args.gpu_memory_gb > 0 and torch.cuda.is_available():
|
|
max_memory = {0: f"{args.gpu_memory_gb:.1f}GiB", "cpu": "48GiB"}
|
|
model = load_qwen_dynamic(model_path, load_in_4bit=not args.no_4bit, max_memory=max_memory)
|
|
model.eval()
|
|
hidden_size = int(model.memory.hidden_size)
|
|
if hidden_size != args.hidden_size:
|
|
_report({"phase": "hidden_size_override", "expected": args.hidden_size, "actual": hidden_size})
|
|
fill_from = fill_to = 0
|
|
if args.fill_rows:
|
|
try:
|
|
fill_from, fill_to = (int(part) for part in args.fill_rows.split(":"))
|
|
except Exception:
|
|
raise SystemExit("--fill-rows expects START:END")
|
|
if not (0 <= fill_from < fill_to <= text_count):
|
|
raise SystemExit(f"--fill-rows out of range for {text_count} texts")
|
|
encoded = phase_encode(
|
|
output_dir, model, tokens, lengths, text_count=text_count, hidden_size=hidden_size,
|
|
max_batch=args.max_batch, token_budget=args.token_budget, progress_every=args.progress_every,
|
|
skip_lm_head=args.skip_lm_head,
|
|
backbone=args.backbone,
|
|
resume_from=_read_progress(output_dir, text_count=text_count) if (args.resume and not fill_to) else 0,
|
|
fill_from=fill_from,
|
|
fill_to=fill_to,
|
|
)
|
|
del model
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
manifest = {
|
|
"format_version": FORMAT_VERSION,
|
|
"train_sha256": _sha256(train_path),
|
|
"eval_sha256": _sha256(eval_path),
|
|
"model_path": str(args.model_path),
|
|
"max_key_tokens": int(args.max_key_tokens),
|
|
"hidden_size": int(hidden_size),
|
|
"text_count": int(text_count),
|
|
"dtype": "float16_memmap",
|
|
"bank": BANK_NAME,
|
|
"index": INDEX_NAME,
|
|
"texts": TEXTS_NAME,
|
|
"kinds": kinds,
|
|
# A manifest is only written with complete=true when this process encoded
|
|
# every row it was asked to; the trainer refuses an incomplete bank. This
|
|
# flag exists because an interrupted encode used to still emit a manifest
|
|
# that looked complete, and a training run was started on zero rows.
|
|
"complete": bool(not fill_to and encoded >= text_count),
|
|
"encoded_rows": int(encoded),
|
|
"encoding": {
|
|
"tokenizer_threads": args.tokenizer_threads,
|
|
"max_batch": args.max_batch,
|
|
"token_budget": args.token_budget,
|
|
"grouping": "exact_token_length",
|
|
"skip_lm_head": bool(args.skip_lm_head),
|
|
"backbone": args.backbone,
|
|
"fill_rows": args.fill_rows or None,
|
|
},
|
|
}
|
|
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
_report({"phase": "done", **manifest})
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|