| |
| """Шаг 1: цели учителя (Qwen3.5-9B) на диск. |
| |
| Зачем отдельным шагом: 15.63 ГиБ весов учителя при 16 ГиБ карты не живут вместе ни с чем, а цель |
| (пост-норм `last_hidden_state` слоя 32) считается один раз на промпт и дальше не меняется. |
| Кэш позволяет перебирать гиперпараметры адаптера часами без единого прогона учителя. |
| |
| Формат кэша (shard_XXXXX.pt, torch.save, грузится с weights_only=True): |
| {"items": [{"ids": int32[L], "y": bf16[L, 4096]}, ...], "texts": [...], "kinds": [...]} |
| Пробы и val — те же поля отдельными файлами, плюс meta.json с параметрами и статистикой норм. |
| |
| python3 src/cache_targets.py --prompts-dir ~/dataset --out cache/teacher --len 256 |
| # продолжить прерванный прогон: те же аргументы, готовые шарды пропускаются |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from moss_lib import MOSS_DIR, encode, load_teacher, load_tokenizer |
| from prompts import NEGATIVES, PROBE, load_prompts |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
|
|
|
|
| def save_items(path, items, texts, kinds): |
| torch.save({"items": items, "texts": texts, "kinds": kinds}, path) |
|
|
|
|
| def load_items(path): |
| return torch.load(path, map_location="cpu", weights_only=True) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="цели учителя MOSS (Qwen3.5-9B) -> кэш на диске") |
| ap.add_argument("--moss-dir", default=MOSS_DIR) |
| ap.add_argument("--teacher-dir", default=None, help="папка учителя (по умолчанию <moss-dir>/text_encoder)") |
| ap.add_argument("--prompts-dir", default=os.path.expanduser("~/dataset")) |
| ap.add_argument("--out", default=os.path.join(ROOT, "cache", "teacher")) |
| ap.add_argument("--limit", type=int, default=0, help="0 — весь пул") |
| ap.add_argument("--len", type=int, default=256, help="макс. длина последовательности (обрезка хвоста)") |
| ap.add_argument("--shard", type=int, default=500, help="промптов в шарде") |
| ap.add_argument("--val-size", type=int, default=128, help="сколько промптов отложить на val") |
| ap.add_argument("--batch", type=int, default=1, |
| help="промптов за один forward; 1 — бит в бит как в пайплайне, >1 меняет " |
| "цель на ~0.5%% (разные блоки ядер, bf16)") |
| ap.add_argument("--prefer-8bit", action="store_true", help="сразу 8 бит, без попытки bf16") |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
|
|
| tok = load_tokenizer(args.moss_dir) |
| pool = load_prompts(args.prompts_dir, limit=args.limit, seed=args.seed) |
| if len(pool) < args.val_size + 1: |
| raise SystemExit(f"пул слишком мал: {len(pool)}") |
|
|
| |
| special_texts = [""] + list(NEGATIVES) |
| special_kinds = ["empty"] + ["neg"] * len(NEGATIVES) |
| val_texts = pool[:args.val_size] |
| train_texts = pool[args.val_size:] |
| probe_texts = [t for _, t in PROBE] |
| print(f"кэш: {len(train_texts)} обучающих + special {len(special_texts)} + val {len(val_texts)} " |
| f"+ probe {len(probe_texts)}, len={args.len}", flush=True) |
|
|
| t0 = time.time() |
| ids = encode(tok, special_texts + train_texts + val_texts + probe_texts, args.len) |
| n_special, n_train, n_val = len(special_texts), len(train_texts), len(val_texts) |
| print(f"токенизация: {time.time() - t0:.1f} c, медиана длины " |
| f"{int(np.median([len(r) for r in ids[n_special:n_special + n_train]]))} из {args.len}", |
| flush=True) |
|
|
| teacher, mode = load_teacher(prefer_8bit=args.prefer_8bit, path=args.teacher_dir) |
|
|
| def targets(rows): |
| """list[list[int]] -> list[(ids, y)]: батч с правым паддингом, цель обрезается по длине. |
| |
| Паддинг справа + маска: реальные токены от него не зависят (внимание причинное), это |
| проверяется selfcheck-ом ниже. |
| """ |
| out = [] |
| for s in range(0, len(rows), args.batch): |
| chunk = rows[s:s + args.batch] |
| length = max(len(r) for r in chunk) |
| idb = torch.full((len(chunk), length), tok.pad_token_id, dtype=torch.long) |
| mb = torch.zeros((len(chunk), length), dtype=torch.long) |
| for i, r in enumerate(chunk): |
| idb[i, :len(r)] = torch.tensor(r) |
| mb[i, :len(r)] = 1 |
| y = teacher.last_hidden(idb.to("cuda"), mb.to("cuda")) |
| if not torch.isfinite(y).all(): |
| raise SystemExit("в цели NaN/inf — учитель посчитал в fp16 и переполнился") |
| for i, r in enumerate(chunk): |
| out.append((torch.tensor(r, dtype=torch.int32), |
| y[i, :len(r)].to(torch.bfloat16).cpu())) |
| return out |
|
|
| if args.batch > 1: |
| check_rows = ids[:4] |
| single = torch.cat([y for r in check_rows for _, y in targets([r])]).float() |
| batched = torch.cat([y for _, y in targets(check_rows)]).float() |
| rel = (single - batched).abs().max().item() / single.norm(dim=-1).mean().item() |
| print(f" selfcheck: батч меняет цель на {rel * 100:.2f}% (округление bf16, ждём <1%)", flush=True) |
| if rel > 0.01: |
| raise SystemExit("батч с паддингом меняет цель сильнее округления — считай --batch 1") |
|
|
| os.makedirs(args.out, exist_ok=True) |
| special = [{"ids": r, "y": y} for r, y in targets(ids[:n_special])] |
| save_items(os.path.join(args.out, "special.pt"), special, special_texts, special_kinds) |
| probe = [{"ids": r, "y": y} for r, y in targets(ids[n_special + n_train + n_val:])] |
| save_items(os.path.join(args.out, "probe.pt"), probe, probe_texts, [n for n, _ in PROBE]) |
| val = [{"ids": r, "y": y} for r, y in targets(ids[n_special + n_train:n_special + n_train + n_val])] |
| save_items(os.path.join(args.out, "val.pt"), val, val_texts, ["val"] * n_val) |
|
|
| norms, shards, tok_total, done, t0 = [], [], 0, 0, time.time() |
| for s in range(0, n_train, args.shard): |
| path = os.path.join(args.out, f"shard_{s // args.shard:05d}.pt") |
| shards.append(os.path.basename(path)) |
| if os.path.exists(path): |
| data = load_items(path) |
| items = data["items"] |
| done += len(items) |
| else: |
| chunk = list(range(s, min(s + args.shard, n_train))) |
| items = [{"ids": r, "y": y} for r, y in targets([ids[n_special + i] for i in chunk])] |
| save_items(path, items, [train_texts[i] for i in chunk], ["pool"] * len(chunk)) |
| done += len(items) |
| el = time.time() - t0 |
| t0 = time.time() |
| print(f" {done}/{n_train} промптов, {el:.0f} c на шард, " |
| f"осталось ~{el * (n_train - done) / max(1, len(items)) / 60:.0f} мин", flush=True) |
| norms += [it["y"].float().norm(dim=-1).numpy() for it in items] |
| tok_total += int(sum(it["ids"].numel() for it in items)) |
|
|
| norm = np.concatenate(norms) |
| meta = {"moss_dir": args.moss_dir, "prompts_dir": args.prompts_dir, "teacher_mode": mode, |
| "len": args.len, "batch": args.batch, |
| "n_train": n_train, "n_val": n_val, "n_probe": len(probe_texts), |
| "n_special": n_special, "shards": shards, "tokens": tok_total, |
| "norms": {f"p{q}": float(np.percentile(norm, q)) for q in (1, 10, 50, 90, 99)}, |
| "norm_mean": float(norm.mean()), "norm_floor": float(np.median(norm))} |
| with open(os.path.join(args.out, "meta.json"), "w") as f: |
| json.dump(meta, f, indent=1) |
| raw = sum(os.path.getsize(os.path.join(args.out, s)) for s in shards) |
| print(f"готово: {args.out}, {n_train} промптов, {tok_total / 1e6:.1f}M токенов, " |
| f"{raw / 2**30:.1f} ГиБ; норма цели медиана {meta['norm_floor']:.0f} " |
| f"(p1 {meta['norms']['p1']:.0f}, p99 {meta['norms']['p99']:.0f}); учитель {mode}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
