Mossling: MOSS-Image pruned to a 16 GB card (alpha)

2e26e31 verified 13 days ago

19.2 kB

#!/usr/bin/env python3
"""Слияние смежных блоков: учим LoRA блока s отдавать то же, что отдавали блоки s..s+k.
Идея: блоки в середине меняют поток мало и частично друг друга отменяют. Значит из пары-тройки
смежных блоков можно оставить один, добавив ему низкоранговую поправку, обученную на состояниях
учителя (исходной модели). Так глубина падает без переобучения сети.
    python3 src/window.py cache --start 3 --drop 2 --samples 120 --size 1024
    python3 src/window.py train --start 3 --drop 2 --samples 120 --size 1024 --rank 32 --steps 3000
    python3 src/window.py eval  --start 3 --drop 2
Кэш: для каждого сэмпла (промпт, латент, t) пишем вход блока s и выход блока s+k учителя.
Память тренировки — один блок (226 МБ) + LoRA, поэтому шаг дешёвый, а учитель нужен только на кэше.
"""
import argparse
import json
import time
import os
import sys

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from lab import (EVAL, LAB, Metrics, build, conditions, grid, sample,   # noqa: E402
                 drop_blocks)

CACHE_DIR = os.path.join(LAB, "cache")
TRAIN_PROMPTS = None


def prompts_for_training(n, seed=0):
    """Строки из пула ~/dataset (по одной на файл), как в обучении адаптера ТЕ."""
    global TRAIN_PROMPTS
    if TRAIN_PROMPTS is not None:
        return TRAIN_PROMPTS
    files = sorted(f for f in __import__("glob").glob(os.path.expanduser("~/dataset/civit/*.txt"))
                   if os.path.getsize(f) > 40)
    lines = []
    for f in files:
        with open(f, errors="ignore") as fh:
            for line in fh:
                line = line.strip()
                if 20 < len(line) < 400:
                    lines.append(line)
    rng = np.random.default_rng(seed)
    TRAIN_PROMPTS = [lines[i] for i in rng.permutation(len(lines))[:n]]
    return TRAIN_PROMPTS


# ---------------------------------------------------------------- LoRA
class LoRALinear(torch.nn.Module):
    """y = W x + (B A x) * scale; A и B — единственные обучаемые тензоры."""

    def __init__(self, base, rank, alpha=None):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad_(False)
        self.a = torch.nn.Parameter(torch.randn(rank, base.in_features) / base.in_features ** 0.5)
        self.b = torch.nn.Parameter(torch.zeros(base.out_features, rank))
        self.scale = (alpha or rank) / rank

    # _Block читает .weight.dtype у проекций
    @property
    def weight(self):
        return self.base.weight

    @property
    def bias(self):
        return self.base.bias

    def forward(self, x):
        # a/b держим в fp32 (точность мелких поправок), считаем в dtype входа
        a, b = self.a.to(x.dtype), self.b.to(x.dtype)
        return self.base(x) + F.linear(F.linear(x, a), b) * self.scale


def inject_lora(block, rank, targets=("q", "k", "v", "o", "ffn_up_gate", "ffn_down")):
    for name in targets:
        lin = block.proj[name]
        if not isinstance(lin, LoRALinear):
            block.proj[name] = LoRALinear(lin, rank)


def merge_lora(block):
    """Влить LoRA в базовые веса: W' = W + B A * scale. На инференсе поправка ничего не стоит."""
    for name, lin in list(block.proj.items()):
        if isinstance(lin, LoRALinear):
            delta = (lin.b.float() @ lin.a.float()) * lin.scale
            base = lin.base
            base.weight.data = (base.weight.data.float() + delta).to(base.weight.dtype)
            block.proj[name] = base
    return block


def lora_state(block):
    return {k: v.detach().cpu() for k, v in block.state_dict().items() if k.endswith((".a", ".b"))}


# ---------------------------------------------------------------- cache
def cmd_cache(a):
    os.makedirs(CACHE_DIR, exist_ok=True)
    out = cache_path(a)
    if os.path.exists(out):
        print("cache exists:", out)
        return
    pipe = build()
    tr = pipe.transformer.eval()
    prompts = prompts_for_training(a.samples)
    # по одному-восемь промптов за раз: батч из 240 строк паддится до 4096 токенов и съедает 8 ГБ
    conds = []
    with torch.no_grad():
        for i in range(0, len(prompts), 8):
            e, m = pipe.encode_prompt(prompts[i:i + 8], 4096)
            conds += [(e[j:j + 1].cpu(), m[j:j + 1].cpu()) for j in range(e.shape[0])]
        nemb, nmask = pipe.encode_prompt([""], 4096)
    nemb, nmask = nemb.cpu(), nmask.cpu()
    pipe.text_encoder = None
    torch.cuda.empty_cache()

    start, last = a.start, a.start + a.drop
    rec = {}
    def make_hook(i):
        def hook(m, inp, out):
            if i in (start - 1, last):
                rec.setdefault("states", {})[i] = out.to(torch.bfloat16).cpu()
            if i == start - 1:
                rec["rope"] = inp[1]
            if i == start:
                rec["mod"] = inp[2]
        return hook
    handles = [l.register_forward_hook(make_hook(i)) for i, l in enumerate(tr.layers)]

    g = a.size // 16
    xs, mods, tgts, ropes, meta = [], [], [], [], []
    rng = np.random.default_rng(a.seed)

    @torch.no_grad()
    def collect(si, ei_emb, ei_mask):
        rec.clear()
        lat = torch.randn(1, 128, g, g,
                          generator=torch.Generator("cpu").manual_seed(int(rng.integers(1 << 30)))).cuda()
        lat = lat.bfloat16()
        t0 = float(rng.uniform(0.05, 0.95))
        # пара шагов Эйлера, чтобы латент был «как на реальной траектории», а не чистый шум
        for _ in range(a.warm):
            v = tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda()).sample.float()
            lat = (lat.float() - 0.08 * v).bfloat16()
        t0 = max(0.02, t0 - 0.08 * a.warm)
        tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda())
        return (rec["states"][start - 1], rec["states"][last],
                rec["mod"][0].cpu(), rec["rope"].cpu(), t0)

    for si in range(a.samples):
        x_s, t_s, mod_s, rope_s, t0 = collect(si, *conds[si])
        xs.append(x_s)
        tgts.append(t_s)
        mods.append(mod_s)
        ropes.append(rope_s)
        meta.append({"prompt": si, "t": round(t0, 3)})
        if si == 0:
            print(f"  state norm {x_s.float().norm():.1f} -> target norm {t_s.float().norm():.1f}, "
                  f"rel diff {((t_s.float() - x_s.float()).norm() / x_s.float().norm()):.3f}", flush=True)
        if (si + 1) % 25 == 0:
            print(f"  cached {si + 1}/{a.samples}", flush=True)
    for h in handles:
        h.remove()
    torch.save({"x": xs, "target": tgts,
                "mod": torch.stack(mods), "rope": ropes, "meta": meta,
                "start": start, "drop": a.drop, "size": a.size, "warm": a.warm}, out)
    print("saved", out, f"{os.path.getsize(out) / 2 ** 30:.1f} GiB")


def cache_path(a):
    return os.path.join(CACHE_DIR, f"win{a.start}_{a.drop}_{a.samples}_{a.size}.pt")


# ---------------------------------------------------------------- train
def cmd_train(a):
    data = cache_path(a)
    if not os.path.exists(data):
        raise SystemExit(f"нет кэша {data}; сначала запустите cache")
    d = torch.load(data, map_location="cpu", weights_only=False)
    xs, tgts = d["x"], d["target"]
    mod = d["mod"].to(torch.bfloat16)
    rope_map = {r.shape[1]: r.to("cuda") for r in d["rope"]}      # rope зависит от числа токенов
    n = len(xs)
    # длины последовательностей разные (у промптов разное число токенов) — батчи собираем
    # по одинаковой длине, чтобы не паддить и не менять поведение блока
    buckets = {}
    for i, x in enumerate(xs):
        buckets.setdefault(x.shape[1], []).append(i)
    rng = np.random.default_rng(0)
    for k in buckets:
        buckets[k] = rng.permutation(buckets[k]).tolist()

    pipe = build()
    tr = pipe.transformer.eval()
    idx = d["start"]
    if a.base_tags:                       # refit учится ПОВЕРХ первых слияний, а не на голом блоке
        drops = apply_spec(tr, [t for t in a.base_tags.split(",") if t])
        # после удаления блоков нумерация съезжает: переводим исходный номер в индекс живой модели
        idx = d["start"] - sum(1 for x in drops if x < d["start"])
        print(f"база: {a.base_tags}, блок {d['start']} -> индекс {idx}", flush=True)
    block = tr.layers[idx].to("cuda")
    pipe.transformer = None
    pipe.vae = None
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    if not a.full:
        inject_lora(block, a.rank)
        block.to("cuda")
    params = [p for p in block.parameters() if p.requires_grad]
    print(f"trainable {sum(p.numel() for p in params) / 1e6:.2f}M"
          f" (base {sum(p.numel() for p in block.parameters()) / 1e6:.1f}M)", flush=True)
    import bitsandbytes as bnb
    opt = bnb.optim.AdamW8bit(params, lr=a.lr, betas=(0.9, 0.95), weight_decay=0.0)

    def pack(idx):
        """индексы одной длины -> (x, target, mod) на GPU"""
        x = torch.cat([xs[i] for i in idx]).cuda().bfloat16()
        t = torch.cat([tgts[i] for i in idx]).cuda().bfloat16()
        return x, t, mod[list(idx)].cuda(), rope_map[x.shape[1]]

    @torch.no_grad()
    def quality(idx):
        block.eval()
        rels, coss = [], []
        for i in idx:
            x, t, m, rope = pack([i])
            pred = block(x, rope, m).float()
            rels.append(float((((pred - t.float()) ** 2).mean(-1)
                               / t.float().pow(2).mean(-1).clamp_min(1e-6)).mean()))
            coss.append(float(F.cosine_similarity(pred.flatten(1), t.float().flatten(1), dim=1)))
        return sum(rels) / len(rels), sum(coss) / len(coss)

    held = [int(v) for k in buckets for v in buckets[k][:max(1, len(buckets[k]) // 10)]][:12]
    rel0, cos0 = quality(held)
    print(f"до обучения (просто удалить блоки): rel {rel0:.4f} cos {cos0:.5f}", flush=True)

    tag = a.tag or f"win{d['start']}_{d['start'] + d['drop']}"
    out = os.path.join(LAB, "lora", f"{tag}.pt")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    keys = [k for k in buckets if len(buckets[k]) > 1] or list(buckets)   # бывает и по одному
    cursor = {k: 0 for k in keys}

    def next_batch():
        k = keys[int(rng.integers(len(keys)))]
        b = buckets[k]
        c = cursor[k]
        idx = [b[(c + j) % len(b)] for j in range(min(a.batch, len(b)))]
        cursor[k] = c + a.batch
        return idx

    step, t_start = 0, time.time()
    while step < a.steps:
        idx = next_batch()
        block.train()
        x, t, m, rope = pack(idx)
        pred = block(x, rope, m)
        tf = t.float()
        rel_tok = ((pred.float() - tf) ** 2).mean(-1) / tf.pow(2).mean(-1).clamp_min(1e-6)
        loss = rel_tok.mean() + a.cos_w * (1 - F.cosine_similarity(
            pred.float().flatten(1), tf.flatten(1), dim=1).mean())
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step()
        step += 1
        if step % 500 == 0 or step == a.steps:
            rel, cos = quality(held)
            print(f"step {step:5d} loss {float(loss):.4f} | held rel {rel:.4f} cos {cos:.5f} "
                  f"| {step / (time.time() - t_start):.1f} шаг/с", flush=True)
            save(out, block, a, d)
    rel, cos = quality(held)
    save(out, block, a, d)
    row = {"tag": tag, "start": d["start"], "drop": d["drop"], "rank": a.rank, "steps": step,
           "held_rel_before": round(rel0, 4), "held_cos_before": round(cos0, 5),
           "held_rel_after": round(rel, 4), "held_cos_after": round(cos, 5),
           "trainable_m": round(sum(p.numel() for p in params) / 1e6, 2)}
    with open(os.path.join(LAB, "window_results.jsonl"), "a") as f:
        f.write(json.dumps(row) + "\n")
    print(json.dumps(row), flush=True)


def save(path, block, a, d):
    torch.save({"state": lora_state(block) if not a.full else
                {k: v.cpu() for k, v in block.state_dict().items()},
                "full": a.full, "rank": a.rank, "start": d["start"], "drop": d["drop"]}, path)


# ---------------------------------------------------------------- eval
def cmd_eval(a):
    tag = a.tag or f"win{a.start}_{a.start + a.drop}"
    path = os.path.join(LAB, "lora", f"{tag}.pt")
    ck = torch.load(path, map_location="cpu", weights_only=False)
    pipe = build()
    tr = pipe.transformer.eval()
    cond = conditions(pipe)
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    block = tr.layers[ck["start"]]
    if ck["full"]:
        block.load_state_dict(ck["state"])
    else:
        inject_lora(block, ck["rank"])
        block.to("cuda")
        block.load_state_dict(ck["state"], strict=False)
    if not ck["full"] and not a.no_merge:
        merge_lora(block)
    for l in block.proj.values():
        for p in l.parameters():
            p.requires_grad_(False)
    drop_blocks(tr, list(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1)))
    images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
    run_dir = os.path.join(LAB, "runs", tag)
    os.makedirs(run_dir, exist_ok=True)
    for i, im in enumerate(images):
        im.save(os.path.join(run_dir, f"{i:02d}.png"))
    grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
    res = {"tag": tag, "blocks": len(tr.layers),
           "params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
           "size": a.size, "steps": a.steps, "seed": a.seed,
           "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
    base_dir = os.path.join(LAB, "runs", "base")
    if os.path.isdir(base_dir):
        base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
        m = Metrics()
        lp = m.lpips_vs(images, base_imgs)
        res["lpips_vs_base"] = None if lp is None else round(lp, 4)
        res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
    with open(os.path.join(LAB, "results.jsonl"), "a") as f:
        f.write(json.dumps(res, ensure_ascii=False) + "\n")
    print(json.dumps(res, ensure_ascii=False), flush=True)


def apply_spec(tr, tags):
    """Влить LoRA всех слияний в веса и удалить поглощённые блоки. Возвращает список удалённых."""
    drops = set()
    for tag in tags:
        ck = torch.load(os.path.join(LAB, "lora", f"{tag}.pt"), map_location="cpu", weights_only=False)
        block = tr.layers[ck["start"]]
        if ck["full"]:
            block.load_state_dict(ck["state"])
        else:
            inject_lora(block, ck["rank"])
            block.to("cuda")
            block.load_state_dict(ck["state"], strict=False)
            merge_lora(block)
        drops |= set(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1))
    drop_blocks(tr, sorted(drops))
    return sorted(drops)


def cmd_combo(a):
    """Собрать все слияния в одну модель: каждый LoRA влит в свои веса, удалённые блоки — вон."""
    tags = [t for t in a.tags.split(",") if t]
    pipe = build()
    tr = pipe.transformer.eval()
    cond = conditions(pipe)
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    drops = set(apply_spec(tr, tags))
    tag = a.tag or "combo"
    images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
    run_dir = os.path.join(LAB, "runs", tag)
    os.makedirs(run_dir, exist_ok=True)
    for i, im in enumerate(images):
        im.save(os.path.join(run_dir, f"{i:02d}.png"))
    grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
    res = {"tag": tag, "blocks": len(tr.layers), "dropped": sorted(drops),
           "params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
           "size": a.size, "steps": a.sample_steps, "seed": a.seed,
           "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
    base_dir = os.path.join(LAB, "runs", "base")
    if os.path.isdir(base_dir):
        base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
        m = Metrics()
        lp = m.lpips_vs(images, base_imgs)
        res["lpips_vs_base"] = None if lp is None else round(lp, 4)
        res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
    with open(os.path.join(LAB, "results.jsonl"), "a") as f:
        f.write(json.dumps(res, ensure_ascii=False) + "\n")
    print(json.dumps(res, ensure_ascii=False), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["cache", "train", "eval", "combo"])
    ap.add_argument("--tags", default="", help="для combo: через запятую")
    ap.add_argument("--start", type=int, default=3, help="блок, который остаётся и получает LoRA")
    ap.add_argument("--drop", type=int, default=2, help="сколько блоков после него удалить")
    ap.add_argument("--samples", type=int, default=300)
    ap.add_argument("--warm", type=int, default=2, help="шагов Эйлера перед записью состояния")
    ap.add_argument("--size", type=int, default=512)
    ap.add_argument("--steps", type=int, default=3000, help="шаги обучения")
    ap.add_argument("--sample-steps", type=int, default=20, help="шаги сэмплинга на eval")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--rank", type=int, default=32)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--cos-w", type=float, default=0.1)
    ap.add_argument("--full", action="store_true", help="учить весь блок, не только LoRA")
    ap.add_argument("--tag", default=None)
    ap.add_argument("--base-tags", default="", help="сначала влить эти слияния в веса модели")
    ap.add_argument("--no-merge", action="store_true", help="не вливать LoRA в веса на eval")
    a = ap.parse_args()
    {"cache": cmd_cache, "train": cmd_train, "eval": cmd_eval, "combo": cmd_combo}[a.cmd](a)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Слияние смежных блоков: учим LoRA блока s отдавать то же, что отдавали блоки s..s+k.

	Идея: блоки в середине меняют поток мало и частично друг друга отменяют. Значит из пары-тройки
	смежных блоков можно оставить один, добавив ему низкоранговую поправку, обученную на состояниях
	учителя (исходной модели). Так глубина падает без переобучения сети.

	python3 src/window.py cache --start 3 --drop 2 --samples 120 --size 1024
	python3 src/window.py train --start 3 --drop 2 --samples 120 --size 1024 --rank 32 --steps 3000
	python3 src/window.py eval --start 3 --drop 2

	Кэш: для каждого сэмпла (промпт, латент, t) пишем вход блока s и выход блока s+k учителя.
	Память тренировки — один блок (226 МБ) + LoRA, поэтому шаг дешёвый, а учитель нужен только на кэше.
	"""
	import argparse
	import json
	import time
	import os
	import sys

	import numpy as np
	import torch
	import torch.nn.functional as F
	from PIL import Image

	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	from lab import (EVAL, LAB, Metrics, build, conditions, grid, sample, # noqa: E402
	drop_blocks)

	CACHE_DIR = os.path.join(LAB, "cache")
	TRAIN_PROMPTS = None


	def prompts_for_training(n, seed=0):
	"""Строки из пула ~/dataset (по одной на файл), как в обучении адаптера ТЕ."""
	global TRAIN_PROMPTS
	if TRAIN_PROMPTS is not None:
	return TRAIN_PROMPTS
	files = sorted(f for f in __import__("glob").glob(os.path.expanduser("~/dataset/civit/*.txt"))
	if os.path.getsize(f) > 40)
	lines = []
	for f in files:
	with open(f, errors="ignore") as fh:
	for line in fh:
	line = line.strip()
	if 20 < len(line) < 400:
	lines.append(line)
	rng = np.random.default_rng(seed)
	TRAIN_PROMPTS = [lines[i] for i in rng.permutation(len(lines))[:n]]
	return TRAIN_PROMPTS


	# ---------------------------------------------------------------- LoRA
	class LoRALinear(torch.nn.Module):
	"""y = W x + (B A x) * scale; A и B — единственные обучаемые тензоры."""

	def __init__(self, base, rank, alpha=None):
	super().__init__()
	self.base = base
	for p in self.base.parameters():
	p.requires_grad_(False)
	self.a = torch.nn.Parameter(torch.randn(rank, base.in_features) / base.in_features ** 0.5)
	self.b = torch.nn.Parameter(torch.zeros(base.out_features, rank))
	self.scale = (alpha or rank) / rank

	# _Block читает .weight.dtype у проекций
	@property
	def weight(self):
	return self.base.weight

	@property
	def bias(self):
	return self.base.bias

	def forward(self, x):
	# a/b держим в fp32 (точность мелких поправок), считаем в dtype входа
	a, b = self.a.to(x.dtype), self.b.to(x.dtype)
	return self.base(x) + F.linear(F.linear(x, a), b) * self.scale


	def inject_lora(block, rank, targets=("q", "k", "v", "o", "ffn_up_gate", "ffn_down")):
	for name in targets:
	lin = block.proj[name]
	if not isinstance(lin, LoRALinear):
	block.proj[name] = LoRALinear(lin, rank)


	def merge_lora(block):
	"""Влить LoRA в базовые веса: W' = W + B A * scale. На инференсе поправка ничего не стоит."""
	for name, lin in list(block.proj.items()):
	if isinstance(lin, LoRALinear):
	delta = (lin.b.float() @ lin.a.float()) * lin.scale
	base = lin.base
	base.weight.data = (base.weight.data.float() + delta).to(base.weight.dtype)
	block.proj[name] = base
	return block


	def lora_state(block):
	return {k: v.detach().cpu() for k, v in block.state_dict().items() if k.endswith((".a", ".b"))}


	# ---------------------------------------------------------------- cache
	def cmd_cache(a):
	os.makedirs(CACHE_DIR, exist_ok=True)
	out = cache_path(a)
	if os.path.exists(out):
	print("cache exists:", out)
	return
	pipe = build()
	tr = pipe.transformer.eval()
	prompts = prompts_for_training(a.samples)
	# по одному-восемь промптов за раз: батч из 240 строк паддится до 4096 токенов и съедает 8 ГБ
	conds = []
	with torch.no_grad():
	for i in range(0, len(prompts), 8):
	e, m = pipe.encode_prompt(prompts[i:i + 8], 4096)
	conds += [(e[j:j + 1].cpu(), m[j:j + 1].cpu()) for j in range(e.shape[0])]
	nemb, nmask = pipe.encode_prompt([""], 4096)
	nemb, nmask = nemb.cpu(), nmask.cpu()
	pipe.text_encoder = None
	torch.cuda.empty_cache()

	start, last = a.start, a.start + a.drop
	rec = {}
	def make_hook(i):
	def hook(m, inp, out):
	if i in (start - 1, last):
	rec.setdefault("states", {})[i] = out.to(torch.bfloat16).cpu()
	if i == start - 1:
	rec["rope"] = inp[1]
	if i == start:
	rec["mod"] = inp[2]
	return hook
	handles = [l.register_forward_hook(make_hook(i)) for i, l in enumerate(tr.layers)]

	g = a.size // 16
	xs, mods, tgts, ropes, meta = [], [], [], [], []
	rng = np.random.default_rng(a.seed)

	@torch.no_grad()
	def collect(si, ei_emb, ei_mask):
	rec.clear()
	lat = torch.randn(1, 128, g, g,
	generator=torch.Generator("cpu").manual_seed(int(rng.integers(1 << 30)))).cuda()
	lat = lat.bfloat16()
	t0 = float(rng.uniform(0.05, 0.95))
	# пара шагов Эйлера, чтобы латент был «как на реальной траектории», а не чистый шум
	for _ in range(a.warm):
	v = tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda()).sample.float()
	lat = (lat.float() - 0.08 * v).bfloat16()
	t0 = max(0.02, t0 - 0.08 * a.warm)
	tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda())
	return (rec["states"][start - 1], rec["states"][last],
	rec["mod"][0].cpu(), rec["rope"].cpu(), t0)

	for si in range(a.samples):
	x_s, t_s, mod_s, rope_s, t0 = collect(si, *conds[si])
	xs.append(x_s)
	tgts.append(t_s)
	mods.append(mod_s)
	ropes.append(rope_s)
	meta.append({"prompt": si, "t": round(t0, 3)})
	if si == 0:
	print(f" state norm {x_s.float().norm():.1f} -> target norm {t_s.float().norm():.1f}, "
	f"rel diff {((t_s.float() - x_s.float()).norm() / x_s.float().norm()):.3f}", flush=True)
	if (si + 1) % 25 == 0:
	print(f" cached {si + 1}/{a.samples}", flush=True)
	for h in handles:
	h.remove()
	torch.save({"x": xs, "target": tgts,
	"mod": torch.stack(mods), "rope": ropes, "meta": meta,
	"start": start, "drop": a.drop, "size": a.size, "warm": a.warm}, out)
	print("saved", out, f"{os.path.getsize(out) / 2 ** 30:.1f} GiB")


	def cache_path(a):
	return os.path.join(CACHE_DIR, f"win{a.start}_{a.drop}_{a.samples}_{a.size}.pt")


	# ---------------------------------------------------------------- train
	def cmd_train(a):
	data = cache_path(a)
	if not os.path.exists(data):
	raise SystemExit(f"нет кэша {data}; сначала запустите cache")
	d = torch.load(data, map_location="cpu", weights_only=False)
	xs, tgts = d["x"], d["target"]
	mod = d["mod"].to(torch.bfloat16)
	rope_map = {r.shape[1]: r.to("cuda") for r in d["rope"]} # rope зависит от числа токенов
	n = len(xs)
	# длины последовательностей разные (у промптов разное число токенов) — батчи собираем
	# по одинаковой длине, чтобы не паддить и не менять поведение блока
	buckets = {}
	for i, x in enumerate(xs):
	buckets.setdefault(x.shape[1], []).append(i)
	rng = np.random.default_rng(0)
	for k in buckets:
	buckets[k] = rng.permutation(buckets[k]).tolist()

	pipe = build()
	tr = pipe.transformer.eval()
	idx = d["start"]
	if a.base_tags: # refit учится ПОВЕРХ первых слияний, а не на голом блоке
	drops = apply_spec(tr, [t for t in a.base_tags.split(",") if t])
	# после удаления блоков нумерация съезжает: переводим исходный номер в индекс живой модели
	idx = d["start"] - sum(1 for x in drops if x < d["start"])
	print(f"база: {a.base_tags}, блок {d['start']} -> индекс {idx}", flush=True)
	block = tr.layers[idx].to("cuda")
	pipe.transformer = None
	pipe.vae = None
	pipe.text_encoder = None
	torch.cuda.empty_cache()
	if not a.full:
	inject_lora(block, a.rank)
	block.to("cuda")
	params = [p for p in block.parameters() if p.requires_grad]
	print(f"trainable {sum(p.numel() for p in params) / 1e6:.2f}M"
	f" (base {sum(p.numel() for p in block.parameters()) / 1e6:.1f}M)", flush=True)
	import bitsandbytes as bnb
	opt = bnb.optim.AdamW8bit(params, lr=a.lr, betas=(0.9, 0.95), weight_decay=0.0)

	def pack(idx):
	"""индексы одной длины -> (x, target, mod) на GPU"""
	x = torch.cat([xs[i] for i in idx]).cuda().bfloat16()
	t = torch.cat([tgts[i] for i in idx]).cuda().bfloat16()
	return x, t, mod[list(idx)].cuda(), rope_map[x.shape[1]]

	@torch.no_grad()
	def quality(idx):
	block.eval()
	rels, coss = [], []
	for i in idx:
	x, t, m, rope = pack([i])
	pred = block(x, rope, m).float()
	rels.append(float((((pred - t.float()) ** 2).mean(-1)
	/ t.float().pow(2).mean(-1).clamp_min(1e-6)).mean()))
	coss.append(float(F.cosine_similarity(pred.flatten(1), t.float().flatten(1), dim=1)))
	return sum(rels) / len(rels), sum(coss) / len(coss)

	held = [int(v) for k in buckets for v in buckets[k][:max(1, len(buckets[k]) // 10)]][:12]
	rel0, cos0 = quality(held)
	print(f"до обучения (просто удалить блоки): rel {rel0:.4f} cos {cos0:.5f}", flush=True)

	tag = a.tag or f"win{d['start']}_{d['start'] + d['drop']}"
	out = os.path.join(LAB, "lora", f"{tag}.pt")
	os.makedirs(os.path.dirname(out), exist_ok=True)
	keys = [k for k in buckets if len(buckets[k]) > 1] or list(buckets) # бывает и по одному
	cursor = {k: 0 for k in keys}

	def next_batch():
	k = keys[int(rng.integers(len(keys)))]
	b = buckets[k]
	c = cursor[k]
	idx = [b[(c + j) % len(b)] for j in range(min(a.batch, len(b)))]
	cursor[k] = c + a.batch
	return idx

	step, t_start = 0, time.time()
	while step < a.steps:
	idx = next_batch()
	block.train()
	x, t, m, rope = pack(idx)
	pred = block(x, rope, m)
	tf = t.float()
	rel_tok = ((pred.float() - tf) ** 2).mean(-1) / tf.pow(2).mean(-1).clamp_min(1e-6)
	loss = rel_tok.mean() + a.cos_w * (1 - F.cosine_similarity(
	pred.float().flatten(1), tf.flatten(1), dim=1).mean())
	opt.zero_grad(set_to_none=True)
	loss.backward()
	torch.nn.utils.clip_grad_norm_(params, 1.0)
	opt.step()
	step += 1
	if step % 500 == 0 or step == a.steps:
	rel, cos = quality(held)
	print(f"step {step:5d} loss {float(loss):.4f} \| held rel {rel:.4f} cos {cos:.5f} "
	f"\| {step / (time.time() - t_start):.1f} шаг/с", flush=True)
	save(out, block, a, d)
	rel, cos = quality(held)
	save(out, block, a, d)
	row = {"tag": tag, "start": d["start"], "drop": d["drop"], "rank": a.rank, "steps": step,
	"held_rel_before": round(rel0, 4), "held_cos_before": round(cos0, 5),
	"held_rel_after": round(rel, 4), "held_cos_after": round(cos, 5),
	"trainable_m": round(sum(p.numel() for p in params) / 1e6, 2)}
	with open(os.path.join(LAB, "window_results.jsonl"), "a") as f:
	f.write(json.dumps(row) + "\n")
	print(json.dumps(row), flush=True)


	def save(path, block, a, d):
	torch.save({"state": lora_state(block) if not a.full else
	{k: v.cpu() for k, v in block.state_dict().items()},
	"full": a.full, "rank": a.rank, "start": d["start"], "drop": d["drop"]}, path)


	# ---------------------------------------------------------------- eval
	def cmd_eval(a):
	tag = a.tag or f"win{a.start}_{a.start + a.drop}"
	path = os.path.join(LAB, "lora", f"{tag}.pt")
	ck = torch.load(path, map_location="cpu", weights_only=False)
	pipe = build()
	tr = pipe.transformer.eval()
	cond = conditions(pipe)
	pipe.text_encoder = None
	torch.cuda.empty_cache()
	block = tr.layers[ck["start"]]
	if ck["full"]:
	block.load_state_dict(ck["state"])
	else:
	inject_lora(block, ck["rank"])
	block.to("cuda")
	block.load_state_dict(ck["state"], strict=False)
	if not ck["full"] and not a.no_merge:
	merge_lora(block)
	for l in block.proj.values():
	for p in l.parameters():
	p.requires_grad_(False)
	drop_blocks(tr, list(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1)))
	images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
	run_dir = os.path.join(LAB, "runs", tag)
	os.makedirs(run_dir, exist_ok=True)
	for i, im in enumerate(images):
	im.save(os.path.join(run_dir, f"{i:02d}.png"))
	grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
	res = {"tag": tag, "blocks": len(tr.layers),
	"params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
	"size": a.size, "steps": a.steps, "seed": a.seed,
	"peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
	base_dir = os.path.join(LAB, "runs", "base")
	if os.path.isdir(base_dir):
	base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
	m = Metrics()
	lp = m.lpips_vs(images, base_imgs)
	res["lpips_vs_base"] = None if lp is None else round(lp, 4)
	res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
	with open(os.path.join(LAB, "results.jsonl"), "a") as f:
	f.write(json.dumps(res, ensure_ascii=False) + "\n")
	print(json.dumps(res, ensure_ascii=False), flush=True)


	def apply_spec(tr, tags):
	"""Влить LoRA всех слияний в веса и удалить поглощённые блоки. Возвращает список удалённых."""
	drops = set()
	for tag in tags:
	ck = torch.load(os.path.join(LAB, "lora", f"{tag}.pt"), map_location="cpu", weights_only=False)
	block = tr.layers[ck["start"]]
	if ck["full"]:
	block.load_state_dict(ck["state"])
	else:
	inject_lora(block, ck["rank"])
	block.to("cuda")
	block.load_state_dict(ck["state"], strict=False)
	merge_lora(block)
	drops \|= set(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1))
	drop_blocks(tr, sorted(drops))
	return sorted(drops)


	def cmd_combo(a):
	"""Собрать все слияния в одну модель: каждый LoRA влит в свои веса, удалённые блоки — вон."""
	tags = [t for t in a.tags.split(",") if t]
	pipe = build()
	tr = pipe.transformer.eval()
	cond = conditions(pipe)
	pipe.text_encoder = None
	torch.cuda.empty_cache()
	drops = set(apply_spec(tr, tags))
	tag = a.tag or "combo"
	images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
	run_dir = os.path.join(LAB, "runs", tag)
	os.makedirs(run_dir, exist_ok=True)
	for i, im in enumerate(images):
	im.save(os.path.join(run_dir, f"{i:02d}.png"))
	grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
	res = {"tag": tag, "blocks": len(tr.layers), "dropped": sorted(drops),
	"params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
	"size": a.size, "steps": a.sample_steps, "seed": a.seed,
	"peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
	base_dir = os.path.join(LAB, "runs", "base")
	if os.path.isdir(base_dir):
	base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
	m = Metrics()
	lp = m.lpips_vs(images, base_imgs)
	res["lpips_vs_base"] = None if lp is None else round(lp, 4)
	res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
	with open(os.path.join(LAB, "results.jsonl"), "a") as f:
	f.write(json.dumps(res, ensure_ascii=False) + "\n")
	print(json.dumps(res, ensure_ascii=False), flush=True)


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("cmd", choices=["cache", "train", "eval", "combo"])
	ap.add_argument("--tags", default="", help="для combo: через запятую")
	ap.add_argument("--start", type=int, default=3, help="блок, который остаётся и получает LoRA")
	ap.add_argument("--drop", type=int, default=2, help="сколько блоков после него удалить")
	ap.add_argument("--samples", type=int, default=300)
	ap.add_argument("--warm", type=int, default=2, help="шагов Эйлера перед записью состояния")
	ap.add_argument("--size", type=int, default=512)
	ap.add_argument("--steps", type=int, default=3000, help="шаги обучения")
	ap.add_argument("--sample-steps", type=int, default=20, help="шаги сэмплинга на eval")
	ap.add_argument("--seed", type=int, default=1234)
	ap.add_argument("--rank", type=int, default=32)
	ap.add_argument("--lr", type=float, default=1e-3)
	ap.add_argument("--batch", type=int, default=4)
	ap.add_argument("--cos-w", type=float, default=0.1)
	ap.add_argument("--full", action="store_true", help="учить весь блок, не только LoRA")
	ap.add_argument("--tag", default=None)
	ap.add_argument("--base-tags", default="", help="сначала влить эти слияния в веса модели")
	ap.add_argument("--no-merge", action="store_true", help="не вливать LoRA в веса на eval")
	a = ap.parse_args()
	{"cache": cmd_cache, "train": cmd_train, "eval": cmd_eval, "combo": cmd_combo}[a.cmd](a)


	if __name__ == "__main__":
	main()