| |
| """Шаг 2: обучение адаптера на кэше целей (учитель уже не нужен). |
| |
| Чистая регрессия условия: срезы замороженной 0.8B -> пост-норм hidden states учителя 9B |
| (см. cache_targets.py). Ни картинок, ни VAE, ни диффузии. |
| |
| Лосс: MSE с пер-токенным весом 1/max(||y||^2, floor^2), паддинг весом `--pad-weight`. |
| Пустой промпт и негативы подмешиваются отдельными долями — пустой идёт в негативную ветку CFG |
| на каждом шаге генерации, его нельзя «выучить случайно». |
| |
| python3 src/train_adapter.py --cache cache/teacher --selfcheck |
| python3 src/train_adapter.py --cache cache/teacher --epochs 2 --batch 16 --out out/adapter_v1 |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from adapter_lib import build_adapter, save_safetensors |
| from moss_lib import LAYERS, Student |
|
|
| DEVICE = "cuda" |
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
|
|
|
|
| def load_pt(path): |
| return torch.load(path, map_location="cpu", weights_only=True) |
|
|
|
|
| def collate(items): |
| """list[{ids,y}] -> (ids, mask, y) с правым паддингом до максимума в батче.""" |
| length = max(it["ids"].numel() for it in items) |
| ids = torch.full((len(items), length), 0, dtype=torch.long) |
| mask = torch.zeros(len(items), length, dtype=torch.bool) |
| y = torch.zeros(len(items), length, items[0]["y"].shape[-1], dtype=torch.bfloat16) |
| for i, it in enumerate(items): |
| n = it["ids"].numel() |
| ids[i, :n] = it["ids"].long() |
| mask[i, :n] = True |
| y[i, :n] = it["y"] |
| return ids.to(DEVICE), mask.to(DEVICE), y.to(DEVICE) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="адаптер Qwen3.5-0.8B -> текстовый энкодер MOSS (Qwen3.5-9B)") |
| ap.add_argument("--cache", default=os.path.join(ROOT, "cache", "teacher")) |
| ap.add_argument("--student", default=None, |
| help="по умолчанию локальная копия в <moss-dir>/text_encoder") |
| ap.add_argument("--layers", default=",".join(map(str, LAYERS))) |
| ap.add_argument("--hidden", type=int, default=6144) |
| ap.add_argument("--proj-layers", type=int, default=3) |
| ap.add_argument("--norm", default="ln-per-layer") |
| ap.add_argument("--attention", type=int, default=2, help="блоков ветки внимания по токенам") |
| ap.add_argument("--attn-dim", type=int, default=1024) |
| ap.add_argument("--attn-heads", type=int, default=8) |
| ap.add_argument("--attn-max-len", type=int, default=1024, help="таблица позиций ветки внимания") |
| ap.add_argument("--mixer", type=int, default=2, help="блоков внимания ПО ОСИ СРЕЗОВ") |
| ap.add_argument("--mixer-heads", type=int, default=8) |
| ap.add_argument("--mixer-ffn", type=int, default=2) |
| ap.add_argument("--epochs", type=int, default=2) |
| ap.add_argument("--batch", type=int, default=16) |
| ap.add_argument("--lr", type=float, default=3e-4) |
| ap.add_argument("--min-lr", type=float, default=1e-5) |
| ap.add_argument("--attn-lr", type=float, default=None) |
| ap.add_argument("--norm-floor", type=float, default=0.0, help="0 — взять медиану нормы из meta.json") |
| ap.add_argument("--pad-weight", type=float, default=0.1) |
| ap.add_argument("--empty-ratio", type=float, default=0.03, help="доля пустого промпта (негатив CFG)") |
| ap.add_argument("--neg-ratio", type=float, default=0.02) |
| ap.add_argument("--probe-every", type=int, default=200) |
| ap.add_argument("--save-every", type=int, default=500) |
| ap.add_argument("--max-steps", type=int, default=0) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--init-from", default=None, help="старт с готового адаптера (.safetensors)") |
| ap.add_argument("--out", default=os.path.join(ROOT, "out", "adapter_v1")) |
| ap.add_argument("--selfcheck", action="store_true", help="только проверить контракт и выйти") |
| args = ap.parse_args() |
|
|
| torch.manual_seed(args.seed) |
| rng = np.random.default_rng(args.seed) |
| layers = [int(x) for x in args.layers.split(",")] |
|
|
| meta = json.load(open(os.path.join(args.cache, "meta.json"))) |
| probe = load_pt(os.path.join(args.cache, "probe.pt")) |
| val = load_pt(os.path.join(args.cache, "val.pt")) |
| special = load_pt(os.path.join(args.cache, "special.pt")) |
| floor = args.norm_floor or float(meta["norm_floor"]) |
| print(f"кэш: {meta['n_train']} промптов, val {meta['n_val']}, probe {meta['n_probe']}, " |
| f"учитель {meta['teacher_mode']}, len {meta['len']}, norm_floor {floor:.0f}", flush=True) |
|
|
| student = Student(args.student, layers, device=DEVICE) |
| in_dim = student.in_dim |
| out_dim = int(probe["items"][0]["y"].shape[-1]) |
| adapter = build_adapter(in_dim, out_dim, args.hidden, args.proj_layers, args.norm, |
| n_slices=len(layers), attention=args.attention, attn_dim=args.attn_dim, |
| attn_heads=args.attn_heads, max_len=args.attn_max_len, |
| mixer=args.mixer, mixer_heads=args.mixer_heads, |
| mixer_ffn=args.mixer_ffn).to(DEVICE) |
| n_par = sum(p.numel() for p in adapter.parameters()) |
| n_branch = sum(p.numel() for n, p in adapter.named_parameters() if n.startswith(("attn.", "mixer."))) |
| print(f"адаптер: {n_par / 1e6:.1f}M ({in_dim} -> {args.hidden} x{args.proj_layers - 1} -> {out_dim}), " |
| f"нормировка {args.norm}, срезов {len(layers)}, ветки {n_branch / 1e6:.1f}M", flush=True) |
| if args.init_from: |
| from safetensors.torch import load_file |
| sd = load_file(args.init_from) |
| missing, unexpected = adapter.load_state_dict(sd, strict=False) |
| print(f"старт с {args.init_from}: missing {list(missing)[:3]}, unexpected {list(unexpected)[:3]}", |
| flush=True) |
|
|
| def token_weights(y, mask): |
| w = 1.0 / y.pow(2).sum(-1).clamp_min(floor ** 2) |
| return w * torch.where(mask, 1.0, args.pad_weight) |
|
|
| def loss_of(ids, mask, y): |
| pred = adapter(student.hidden(ids, mask), mask).float() |
| yf = y.float() |
| w = token_weights(yf, mask)[..., None].expand_as(yf) |
| return (w * (pred - yf).pow(2)).sum() / w.sum() |
|
|
| @torch.no_grad() |
| def measure(items): |
| """Относительная ошибка и косинус по реальным токенам — по одному примеру, без паддинга.""" |
| adapter.eval() |
| rel, cos = [], [] |
| for it in items: |
| ids = it["ids"].long().to(DEVICE)[None] |
| mask = torch.ones_like(ids, dtype=torch.bool) |
| pred = adapter(student.hidden(ids, mask), mask)[0].float() |
| y = it["y"].to(DEVICE).float() |
| rel.append(((pred - y).norm(dim=-1) / y.norm(dim=-1).clamp_min(1e-6)).mean().item()) |
| cos.append(F.cosine_similarity(pred, y, dim=-1).mean().item()) |
| adapter.train() |
| return float(np.mean(rel)), float(np.mean(cos)) |
|
|
| def report(step, tag): |
| names = probe["kinds"] |
| rel, cos = measure(probe["items"]) |
| vrel, vcos = measure(val["items"]) |
| print(f" {tag}@{step}: probe cos {cos:.3f} / rel {rel * 100:.1f}% | " |
| f"val cos {vcos:.3f} / rel {vrel * 100:.1f}%", flush=True) |
| return {"step": step, "probe_cos": cos, "probe_rel": rel, "val_cos": vcos, "val_rel": vrel, |
| "probe_each": {n: measure([it])[1] for n, it in zip(names, probe["items"])}} |
|
|
| if args.selfcheck: |
| hist = report(0, "probe") |
| print(f" probe по кейсам: " + ", ".join(f"{k}={v:.3f}" for k, v in hist["probe_each"].items())) |
| print(f" VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ, самопроверка пройдена") |
| return |
|
|
| def save(step): |
| save_safetensors(adapter, args.out + ".safetensors", { |
| "student_layers": ",".join(map(str, layers)), "norm": args.norm, "hidden": args.hidden, |
| "proj_layers": args.proj_layers, "attention": args.attention, "attn_dim": args.attn_dim, |
| "attn_heads": args.attn_heads, "attn_max_len": args.attn_max_len, "mixer": args.mixer, |
| "mixer_heads": args.mixer_heads, "n_slices": len(layers), "in_dim": in_dim, |
| "out_dim": out_dim, "norm_floor": floor, "pad_weight": args.pad_weight, |
| "student": student.path, "teacher_mode": meta["teacher_mode"], "len": meta["len"], |
| "prompts_dir": meta["prompts_dir"], "step": step}) |
|
|
| mlp_params = [p for n, p in adapter.named_parameters() if not n.startswith(("attn.", "mixer."))] |
| branch_params = [p for n, p in adapter.named_parameters() if n.startswith(("attn.", "mixer."))] |
| groups = [{"params": mlp_params, "base": args.lr, "floor": args.min_lr}] |
| if branch_params: |
| branch_lr = args.attn_lr if args.attn_lr is not None else args.lr |
| groups.append({"params": branch_params, "base": branch_lr, "floor": args.min_lr}) |
| opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=0.0) |
|
|
| n_train = int(meta["n_train"]) |
| steps_per_epoch = n_train // args.batch |
| total_steps = steps_per_epoch * args.epochs |
| if args.max_steps: |
| total_steps = min(total_steps, args.max_steps) |
| print(f"{steps_per_epoch} шагов/эпоху, {total_steps} всего", flush=True) |
|
|
| adapter.train() |
| history = [report(0, "probe")] |
| step, seen, t_epoch = 0, 0, time.time() |
| for epoch in range(args.epochs): |
| order = rng.permutation(len(meta["shards"])) |
| for si in order: |
| shard = load_pt(os.path.join(args.cache, meta["shards"][si])) |
| items = shard["items"] |
| idx = rng.permutation(len(items)) |
| for b in range(0, len(idx) - args.batch + 1, args.batch): |
| batch = [items[i] for i in idx[b:b + args.batch]] |
| mixed = [] |
| for it in batch: |
| r = rng.random() |
| if r < args.empty_ratio: |
| mixed.append(special["items"][0]) |
| elif r < args.empty_ratio + args.neg_ratio: |
| mixed.append(special["items"][1 + int(rng.integers(len(special["items"]) - 1))]) |
| else: |
| mixed.append(it) |
| ids, mask, y = collate(mixed) |
| frac = 0.5 * (1 + np.cos(np.pi * step / max(1, total_steps))) |
| for g in opt.param_groups: |
| g["lr"] = g["floor"] + (g["base"] - g["floor"]) * frac |
| opt.zero_grad() |
| loss = loss_of(ids, mask, y) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) |
| opt.step() |
| step += 1 |
| seen += ids.numel() |
| if step % 20 == 0: |
| el = time.time() - t_epoch |
| print(f" эпоха {epoch} шаг {step}/{total_steps}: loss {loss.item():.4f}, " |
| f"lr {opt.param_groups[0]['lr']:.1e}, {el / 20:.2f} с/шаг, " |
| f"{seen / el / 1000:.0f}k ток/с", flush=True) |
| t_epoch = time.time() |
| if args.probe_every and step % args.probe_every == 0: |
| history.append(report(step, "probe")) |
| if args.save_every and step % args.save_every == 0: |
| save(step) |
| if args.max_steps and step >= args.max_steps: |
| break |
| if args.max_steps and step >= args.max_steps: |
| print(f"достигнут --max-steps {args.max_steps}, стоп", flush=True) |
| break |
| if args.max_steps and step >= args.max_steps: |
| break |
|
|
| history.append(report(step, "final")) |
| save(step) |
| with open(args.out + ".history.json", "w") as f: |
| json.dump(history, f, indent=1) |
| print(f"готово: {args.out}.safetensors ({n_par / 1e6:.1f}M, шагов {step}, " |
| f"пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ)", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
