Blame for src/train_clean.py

#!/usr/bin/env python3
"""Обучение DiT Mossling на закэшированном датасете — обычный, простой цикл.

Что здесь есть, и только это: веса fp32, forward/backward под autocast(bf16), AdamW8bit по
fp32-весам, градиенты bf16, накопление градиента, клип нормы, cosine с warmup. Никаких блочных
обновлений, свопов весов между bf16 и fp32, выгрузки входов блоков в RAM, paged-аллокатора
и ручной возни с градиентами — это всё было в `train.py` ради 16 ГиБ.

Почему fp32-веса, а не bf16-копия: при lr ~1e-5 шаг `w -= lr*g` — это изменение ~1e-5 от |w|,
а квант bf16 (7 мантийных бит) ~2e-3 от |w|. В bf16 шаг не представим: либо ноль, либо ступенька
решётки (замер: 96.7 % параметров не двигались). Веса fp32 (13.9 ГиБ на 3.72 B), арифметика — bf16.

Почему градиенты bf16: fp32-градиенты (13.9 ГиБ) вместе с fp32-весами и состояниями Adam на 31 ГиБ
не влезают. `p.grad_dtype = torch.bfloat16` — autograd копит градиент сразу в bf16 (6.9 ГиБ),
а шаг `w -= lr*m/(sqrt(v)+eps)` считается в fp32 по fp32-весам, поэтому обновление не тонет
в решётке. `--accum > 1` копит в bf16, для больших accum точнее `--grad-dtype fp32`.

Память на batch 1 (3.72 B): веса 13.9 + градиенты 6.9 + состояния AdamW8bit ~7.0 = 27.8 ГиБ,
остальное — активации. Активации без пересчёта дорогие: замерено на 720x1280 — forward с
градиентом держит ~1.31 ГиБ на каждый из 30 блоков (39 ГиБ), и шаг не влезает даже в 48 ГиБ.
Поэтому блоки пересчитываются на backward (`torch.utils.checkpoint`, флаг `--no-grad-ckpt`
выключает): вместо всех активаций живёт вход блока (~22 МБ) плюс один блок на пересчёт.
Цена — ~1.3x времени на шаг. Замерено на 1152x1152 (самый тяжёлый размер датасета): batch 1 —
3.0 с и 32.1 ГиБ, batch 4 — 10.25 с и 35.9 ГиБ, batch 8 — 20.5 с и 40.5 ГиБ. Скорость
в пересчёте на сэмпл одинаковая начиная с batch 4 (2.56 с/сэмпл против 3.0 у batch 1):
фиксированный оверхед шага (`AdamW8bit.step` + клип по 3.72 B) ~0.5 с, дальше его не размазать.
`peak` печатается каждый шаг и пишется в log.csv: верить этому числу, а не арифметике на бумаге;
`--probe` меряет ровно тем же режимом, что и обучение.

Сигма по умолчанию logit-normal со сдвигом -0.3: лёгкий перекос в сторону низких сигма (детали)
относительно равномерного, как учили базу. При старте печатается фактическое распределение.

Чекпойнт раз в ~час (`--save-every 180`), держим две последние копии (`--keep 2`), пишем в bf16
(`--save-dtype`): 6.9 ГиБ и быстрее, точность шага на диске не нужна — из двух последних
выбирается лучший по сэмплам (карточки каждые полчаса, `--sample-every 1800`).

Данные: датасет `make_dataset.py` без `--te` (латент + капшен, колонок text_embeds нет). Условие
считается на лету адаптерным энкодером Mossling — тем же путём, что в `make_dataset.py`.

Запуск:
    python3 src/train_clean.py --dataset datasets/ds_testg_filtered --epochs 1
    python3 src/train_clean.py --dataset ... --accum 2 --compile
    python3 src/train_clean.py --dataset ... --init out/step_01000 --start-step 1000
"""
import argparse
import gc
import math
import os
import shutil
import sys
import time

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

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

MODEL = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if MODEL not in sys.path:
    sys.path.insert(0, MODEL)

from pipeline_moss_image import MOSSImageScheduler, MOSSImageTransformer  # noqa: E402


def parse_args():
    p = argparse.ArgumentParser(description="Mossling DiT: обычный цикл обучения (fp32 + AdamW8bit)")
    p.add_argument("--dataset", default=os.path.join(MODEL, "datasets"))
    p.add_argument("--out", default=os.path.join(MODEL, "src", "work", "train_clean"))
    p.add_argument("--epochs", type=float, default=1.0)
    p.add_argument("--steps", type=int, default=0, help="точное число шагов, перебивает --epochs")
    p.add_argument("--start-step", type=int, default=0,
                   help="продолжить с этого шага (веса — через --init): lr и лог продолжаются, "
                        "состояние Adam начинается заново")
    p.add_argument("--batch", type=int, default=8,
                   help="сэмплов на шаг; с 4 до 8 скорость на сэмпл не меняется, только память")
    p.add_argument("--accum", type=int, default=1, help="шагов накопления градиента")
    p.add_argument("--lr", type=float, default=1e-5)
    p.add_argument("--lr-min", type=float, default=0.0, help="конец cosine; 0 — lr/10")
    p.add_argument("--warmup", type=int, default=500)
    p.add_argument("--grad-dtype", choices=("bf16", "fp32"), default="bf16",
                   help="bf16 — градиенты вдвое легче и копятся сразу в bf16 (p.grad_dtype); "
                        "fp32 — точнее, но с активациями на 31 ГиБ не влезает")
    p.add_argument("--clip", type=float, default=1.0, help="клип нормы градиента; 0 — выключить")
    p.add_argument("--no-grad-ckpt", action="store_true",
                   help="хранить активации блоков вместо пересчёта: быстрее, но +39 ГиБ на 720x1280")
    p.add_argument("--no-compile", action="store_true", help="выключить torch.compile модели")
    p.add_argument("--save-dtype", choices=("bf16", "fp32"), default="bf16",
                   help="bf16 — вдвое легче и быстрее писать; fp32 нужен только для точного resume")
    p.add_argument("--text-len", type=int, default=1024)
    p.add_argument("--text-dropout", type=float, default=0.05,
                   help="доля сэмплов с пустым промптом (CFG-ветка)")
    p.add_argument("--word-dropout", type=float, default=0.05,
                   help="доля сэмплов, где из промпта выкидывается одно случайное слово")
    p.add_argument("--sigma", choices=("uniform", "logit"), default="logit",
                   help="uniform — U(0,1), как учили базу; logit — sigmoid(N(shift, scale)), "
                        "сдвиг в низкие сигма даёт больше шагов на детали")
    p.add_argument("--sigma-shift", type=float, default=-0.3)
    p.add_argument("--sigma-scale", type=float, default=1.0)
    p.add_argument("--sizes", default="", help="только эти размеры, напр. 768x1152,1024x768")
    p.add_argument("--limit", type=int, default=0, help="первые N строк (0 — все)")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--log-every", type=int, default=25, help="0 — не писать лог")
    p.add_argument("--save-every", type=int, default=180, help="шагов; ~1 час при 20 с/шаг")
    p.add_argument("--keep", type=int, default=2, help="сколько чекпойнтов держать")
    p.add_argument("--init", default="", help="стартовать с чекпойнта (только веса)")
    p.add_argument("--sample-every", type=float, default=1800.0, help="секунд; 0 — выключить")
    p.add_argument("--sample-n", type=int, default=4)
    p.add_argument("--sample-steps", type=int, default=30)
    p.add_argument("--sample-size", type=int, default=1024)
    p.add_argument("--sample-keep", type=int, default=4)
    p.add_argument("--sample-seed", type=int, default=42)
    p.add_argument("--guidance", type=float, default=3.0)
    p.add_argument("--probe", action="store_true",
                   help="только проверить, какой batch влезает: шаги 1..N с ростом батча, без записи")
    return p.parse_args()


# ---------------------------------------------------------------- данные
def find_chunks(root):
    """Все arrow-датасеты под root (папки с dataset_info.json) — так их пишет make_dataset.py."""
    out = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if not d.startswith(".")]
        if "dataset_info.json" in filenames or "state.json" in filenames:
            out.append(dirpath)
    return out


def open_dataset(path, limit, sizes):
    from datasets import concatenate_datasets, load_from_disk
    chunks = find_chunks(path)
    if not chunks:
        raise FileNotFoundError(f"в {path} нет arrow-датасета")
    ds = load_from_disk(chunks[0]) if len(chunks) == 1 else \
        concatenate_datasets([load_from_disk(c) for c in chunks])
    for col in ("vae", "text", "width", "height"):
        if col not in ds.column_names:
            raise ValueError(f"в датасете нет колонки {col!r}: нужен формат make_dataset.py без --te, "
                             f"есть {ds.column_names}")
    widths = np.asarray(ds["width"], dtype=np.int64)
    heights = np.asarray(ds["height"], dtype=np.int64)
    wanted = {x.strip() for x in sizes.split(",") if x.strip()}
    if wanted:
        keys = np.char.add(np.char.add(widths.astype("U8"), "x"), heights.astype("U8"))
        keep = np.isin(keys, sorted(wanted))
    else:
        keep = np.ones(len(widths), dtype=bool)
    idx = np.flatnonzero(keep)
    if limit > 0:
        idx = idx[:limit]
    rows = [{"idx": int(i), "width": int(widths[i]), "height": int(heights[i])} for i in idx]
    return ds.with_format("numpy"), rows


class Buckets:
    """Батчи только из одинаковых размеров: латент — один тензор, разные HxW не склеить.

    Размер выбирается пропорционально числу строк, поэтому за `len(rows)` шагов каждая строка
    просматривается примерно один раз: одна эпоха = один проход по датасету.
    """

    def __init__(self, rows, gen):
        self.rows, self.gen = rows, gen
        self.by_size = {}
        for i, r in enumerate(rows):
            self.by_size.setdefault((r["width"], r["height"]), []).append(i)
        self.sizes = sorted(self.by_size)
        print("размеры: " + ", ".join(f"{w}x{h}: {len(v)}" for (w, h), v in self.by_size.items()),
              flush=True)
        self._perm, self.cursor = {}, {}

    def _shuffle(self, size):
        order = list(self.by_size[size])
        for i in range(len(order) - 1, 0, -1):
            j = int(torch.randint(0, i + 1, (1,), generator=self.gen))
            order[i], order[j] = order[j], order[i]
        self._perm[size], self.cursor[size] = order, 0

    def batch(self, bs):
        if not self.cursor:
            for size in self.sizes:
                self._shuffle(size)
        weights = torch.tensor([len(self.by_size[s]) for s in self.sizes], dtype=torch.float)
        size = self.sizes[int(torch.multinomial(weights, 1, generator=self.gen))]
        picked = []
        for _ in range(bs):
            if self.cursor[size] >= len(self._perm[size]):
                self._shuffle(size)  # размер кончился — новый проход в новом порядке
            picked.append(self.rows[self._perm[size][self.cursor[size]]])
            self.cursor[size] += 1
        return picked


def to_batch(rows, ds, device):
    """rows (одинаковый HxW) -> латент [B,128,h,w] bf16 + тексты."""
    recs = [ds[r["idx"]] for r in rows]
    lat = torch.stack([torch.from_numpy(np.asarray(rec["vae"], dtype=np.float32)) for rec in recs])
    if lat.shape[1] != 128:
        raise ValueError(f"латент {tuple(lat.shape)}: ожидается [B,128,h,w] — формат make_dataset.py")
    return lat.to(device, torch.bfloat16), [rec["text"] for rec in recs]


# ---------------------------------------------------------------- условие на лету
class LiveTextEncoder:
    """Капшен -> условие [1,L,4096] тем же путём, что в make_dataset.py: chat-шаблон, без паддинга.

    `MOSSAdapterTextEncoder.from_pretrained` сам отдаёт bf16 и игнорирует torch_dtype — поэтому
    энкодер ничего не стоит по памяти и не требует флагов.
    """

    def __init__(self, model_dir, device, text_len):
        from transformers import AutoTokenizer
        from pipeline_moss_image import MOSSAdapterTextEncoder
        path = os.path.join(model_dir, "text_encoder")
        self.tok = AutoTokenizer.from_pretrained(path)
        self.te = MOSSAdapterTextEncoder.from_pretrained(path).to(device).eval()
        self.te.requires_grad_(False)
        self.device, self.text_len = torch.device(device), text_len

    @torch.no_grad()
    def encode(self, texts):
        ids_list = [
            self.tok(self.tok.apply_chat_template(
                [{"role": "user", "content": t}], tokenize=False,
                add_generation_prompt=True, enable_thinking=True),
                truncation=True, max_length=self.text_len)["input_ids"] for t in texts]
        if any(not x for x in ids_list):
            raise ValueError("пустой промпт после токенизации")
        length = max(len(x) for x in ids_list)
        ids = torch.zeros(len(ids_list), length, dtype=torch.long)
        mask = torch.zeros(len(ids_list), length, dtype=torch.bool)
        for i, x in enumerate(ids_list):  # паддинг только для батча; маска его исключает
            ids[i, :len(x)] = torch.tensor(x)
            mask[i, :len(x)] = True
        out = self.te(input_ids=ids.to(self.device), attention_mask=mask.to(self.device),
                      use_cache=False)
        return out.last_hidden_state.to(torch.bfloat16), mask.to(self.device)


def set_grad_dtype(params, dtype):
    """fp32-веса + bf16-градиенты: autograd копит сразу в bf16, fp32-градиенты не существуют.

    Проверено на torch 2.14: `clip_grad_norm_` и `bnb.optim.AdamW8bit` принимают такие градиенты.
    """
    if not hasattr(torch.Tensor, "grad_dtype"):
        raise SystemExit("torch без grad_dtype — обновите torch или запустите с --grad-dtype fp32")
    for p in params:
        p.grad_dtype = dtype


def drop_word(text, gen):
    words = text.split()
    if not words:
        return text
    del words[int(torch.randint(0, len(words), (1,), generator=gen))]
    return " ".join(words) or "image"


# ---------------------------------------------------------------- сэмпл прогресса
def sample_grid(tr, rows, ds, args, live, out_path):
    from diffusers import AutoencoderKLFlux2
    from pipeline_moss_image import MOSSImagePipeline
    gc.collect()
    torch.cuda.empty_cache()
    tr.eval()
    try:
        vae = AutoencoderKLFlux2.from_pretrained(os.path.join(MODEL, "vae"),
                                                torch_dtype=torch.float16).to("cuda").eval()
        pipe = MOSSImagePipeline(tr, MOSSImageScheduler(), None, None, vae)
        pipe.set_progress_bar_config(disable=True)
        neg = live.encode([""]) if args.guidance > 1.0 else (None, None)
        tiles = []
        for i, meta in enumerate(rows):
            rec = ds[meta["idx"]]
            long_side = args.sample_size
            if meta["width"] >= meta["height"]:
                w, h = long_side, round(long_side * meta["height"] / meta["width"] / 32) * 32
            else:
                w, h = round(long_side * meta["width"] / meta["height"] / 32) * 32, long_side
            cond, cmask = live.encode([rec["text"]])
            with torch.autocast("cuda", dtype=torch.bfloat16):  # fp32-веса, счёт в bf16
                tiles.append(pipe(prompt_embeds=cond, prompt_attention_mask=cmask,
                                  negative_prompt_embeds=neg[0],
                                  negative_prompt_attention_mask=neg[1],
                                  height=h, width=w, num_inference_steps=args.sample_steps,
                                  guidance_scale=args.guidance if neg[0] is not None else 1.0,
                                  generator=torch.Generator("cuda").manual_seed(args.sample_seed + i),
                                  output_type="pil").images[0])
        del pipe, vae
        gc.collect()
        torch.cuda.empty_cache()
        cols = min(2, len(tiles))
        n_rows = (len(tiles) + cols - 1) // cols
        cell = args.sample_size
        grid = Image.new("RGB", (cols * cell, n_rows * cell), "white")
        for i, tile in enumerate(tiles):
            grid.paste(tile, ((i % cols) * cell + (cell - tile.width) // 2,
                              (i // cols) * cell + (cell - tile.height) // 2))
        grid.save(out_path)
        print(f"сэмпл -> {out_path}", flush=True)
    finally:
        tr.train()


def try_sample(tr, rows, ds, args, live, out_path):
    """Сэмпл не должен убивать прогон: OOM здесь — пропущенная картинка, а не потеря обучения."""
    try:
        sample_grid(tr, rows, ds, args, live, out_path)
    except torch.OutOfMemoryError:
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        print("сэмпл: OOM — картинка пропущена, обучение продолжается", flush=True)


def prune_samples(out, keep):
    if keep <= 0:
        return
    names = sorted(n for n in os.listdir(out) if n.startswith("progress_") and n.endswith(".png"))
    for n in names[:-keep]:
        os.remove(os.path.join(out, n))


# ---------------------------------------------------------------- обучение
PROTECTED = ("transformer", "vae", "text_encoder", "adapter", "tokenizer", "scheduler")
DTYPES = {"bf16": torch.bfloat16, "fp32": torch.float32}


def guard_out(out, dataset):
    """Чекпойнты не должны лечь в веса модели или в датасет."""
    out_abs, model_abs = os.path.realpath(out), os.path.realpath(MODEL)

    def inside(path, root):
        return path == root or path.startswith(root + os.sep)

    if out_abs == model_abs:
        raise SystemExit(f"--out {out_abs} — это корень модели; выберите отдельный каталог")
    for name in PROTECTED:
        if inside(out_abs, os.path.join(model_abs, name)):
            raise SystemExit(f"--out {out_abs} лежит внутри {name}/ модели; выберите отдельный каталог")
    if inside(out_abs, os.path.realpath(dataset)):
        raise SystemExit(f"--out {out_abs} лежит внутри датасета; выберите отдельный каталог")


def build_model(device, init=""):
    tr = MOSSImageTransformer.from_pretrained(
        os.path.join(MODEL, "transformer"), torch_dtype=torch.float32,
        local_files_only=True, low_cpu_mem_usage=False).to(device).train()
    if init:
        tr.load_training_state_dict(MOSSImageTransformer.from_pretrained(
            init, torch_dtype=torch.float32, local_files_only=True,
            low_cpu_mem_usage=False).state_dict())
        print(f"старт с {init}", flush=True)
    return tr


def set_block_checkpointing(tr, on):
    """Пересчёт блока на backward вместо хранения его активаций.

    Замерено на 720x1280: forward с градиентом держит ~1.31 ГиБ на блок, 39 ГиБ на 30 блоков.
    С пересчётом autograd хранит только вход блока (~22 МБ) и один блок за раз.
    """
    from torch.utils.checkpoint import checkpoint
    for layer in tr.layers:
        if on and "forward" not in layer.__dict__:
            block = layer.forward
            layer.forward = lambda x, rope, mod, _b=block: checkpoint(
                _b, x, rope, mod, use_reentrant=False)
        elif not on and "forward" in layer.__dict__:
            del layer.forward


def raw_model(tr):
    """Компилированный модуль -> исходный: у `_orig_mod` в state_dict правильные имена ключей."""
    return tr._orig_mod if hasattr(tr, "_orig_mod") else tr


def save_ckpt(tr, path, dtype):
    """Чекпойнт bf16: 6.9 ГиБ вместо 13.9 и быстрее запись; копия идёт сразу в RAM.

    Resume с bf16-чекпойнта округлит веса до bf16 (шаг 1e-5 снова утонет в решётке) —
    если нужен точный resume, `--save-dtype fp32`.
    """
    model = raw_model(tr)
    if dtype == torch.float32:
        model.save_pretrained(path)
        return
    from safetensors.torch import save_file
    os.makedirs(path, exist_ok=True)
    tensors = {k: v.detach().to("cpu", dtype).contiguous() for k, v in model.state_dict().items()}
    save_file(tensors, os.path.join(path, "diffusion_pytorch_model.safetensors"))
    model.save_config(path)


def one_step(tr, batch_rows, ds, args, live, gen, device):
    """Один микро-шаг; градиент копит autograd, шаг делает вызывающий."""
    lat, texts = to_batch(batch_rows, ds, device)
    if args.text_dropout or args.word_dropout:
        rnd = torch.rand(len(texts), generator=gen).tolist()
        texts = ["" if r < args.text_dropout
                 else (drop_word(t, gen) if r < args.text_dropout + args.word_dropout else t)
                 for r, t in zip(rnd, texts)]
    with torch.autocast(device, dtype=torch.bfloat16):
        cond, mask = live.encode(texts)
        noise = torch.randn(lat.shape, dtype=torch.float32, device=device)
        if args.sigma == "logit":
            sigma = torch.sigmoid(torch.randn(lat.shape[0], 1, 1, 1, device=device)
                                  * args.sigma_scale + args.sigma_shift)
        else:
            sigma = torch.rand(lat.shape[0], 1, 1, 1, device=device)
        x0 = lat.float()
        xt = ((1 - sigma) * x0 + sigma * noise).to(torch.bfloat16)
        pred = tr(xt, sigma.flatten() * 1000.0, cond, mask).sample.float()
        loss = F.mse_loss(pred, noise - x0)
    return loss


def main():
    args = parse_args()
    guard_out(args.out, args.dataset)
    torch.manual_seed(args.seed)
    device = "cuda"
    torch.cuda.set_per_process_memory_fraction(0.95)

    ds, rows = open_dataset(args.dataset, args.limit, args.sizes)
    if not rows:
        raise SystemExit("датасет пуст")
    print(f"строк: {len(rows)}, устройство {torch.cuda.get_device_name(0)} "
          f"({torch.cuda.get_device_properties(0).total_memory / 2**30:.1f} ГиБ)", flush=True)
    gen = torch.Generator().manual_seed(args.seed)
    buckets = Buckets(rows, gen)
    live = LiveTextEncoder(MODEL, device, args.text_len)

    if args.probe:
        return probe(args, ds, buckets, device, gen, live)

    tr = build_model(device, args.init)
    if not args.no_grad_ckpt:
        set_block_checkpointing(tr, True)
        print("активации блоков: пересчёт на backward (--no-grad-ckpt выключает)", flush=True)
    params = [p for p in tr.parameters() if p.requires_grad]
    if args.grad_dtype == "bf16":
        set_grad_dtype(params, torch.bfloat16)
    if not args.no_compile:
        tr = torch.compile(tr, dynamic=True)  # размеры и длина текста меняются на каждом батче
        print("torch.compile: модель целиком, dynamic=True (первый шаг долгий; чекпойнт из _orig_mod)",
              flush=True)
    n = sum(p.numel() for p in params)
    print(f"параметров {n / 1e9:.3f}B; веса fp32 {n * 4 / 2**30:.2f} ГиБ; "
          f"batch {args.batch} x accum {args.accum}", flush=True)

    import bitsandbytes as bnb
    opt = bnb.optim.AdamW8bit(params, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0)
    print(f"оптимизатор: AdamW8bit по fp32-весам; градиенты {args.grad_dtype}", flush=True)

    os.makedirs(args.out, exist_ok=True)
    log_path = os.path.join(args.out, "log.csv")
    if args.start_step == 0:
        with open(log_path, "w") as fh:
            fh.write("step,loss,lr,sec,peak_gib,eta_h\n")

    step_rows = args.batch * args.accum
    steps = args.steps or max(1, math.ceil(args.epochs * len(rows) / step_rows))
    min_lr = args.lr_min or args.lr / 10
    print(f"шагов {steps} (1 эпоха = {len(rows)} строк / {step_rows} на шаг), lr {args.lr:.1e} -> "
          f"{min_lr:.1e}, warmup {args.warmup}", flush=True)
    if args.warmup >= steps:
        print(f"внимание: warmup {args.warmup} >= шагов {steps} — lr не дойдёт до {args.lr:.1e}",
              flush=True)
    if args.sigma == "logit":  # печатаем, что именно получилось на бумаге
        sigmas = torch.sigmoid(torch.randn(200000) * args.sigma_scale + args.sigma_shift)
        print(f"sigma: logit-normal m={args.sigma_shift:+.2f} s={args.sigma_scale:.2f}, "
              f"медиана {sigmas.median():.2f}, доля шагов sigma<0.2: "
              f"{(sigmas < 0.2).float().mean():.0%}, sigma>0.8: {(sigmas > 0.8).float().mean():.0%}",
              flush=True)
    else:
        print("sigma: uniform U(0,1), как в обучении базы", flush=True)

    def lr_at(step):
        if args.warmup and step < args.warmup:
            return min_lr + (args.lr - min_lr) * (step + 1) / args.warmup
        p = min(1.0, max(0.0, (step - args.warmup) / max(1, steps - args.warmup)))
        return min_lr + (args.lr - min_lr) * 0.5 * (1 + math.cos(math.pi * p))

    def sample_subset():
        idx = np.linspace(0, len(rows) - 1, min(args.sample_n, len(rows)))
        return [rows[int(i)] for i in idx]

    t_start = time.time()
    t_train = 0.0
    peak_all = 0.0
    sample_n = 0
    if args.sample_every:
        try_sample(tr, sample_subset(), ds, args, live, os.path.join(args.out, "progress_00000.png"))
        sample_n += 1
    t_last_sample = time.time()

    for step in range(args.start_step + 1, steps + 1):
        lr = lr_at(step - 1)
        for g in opt.param_groups:
            g["lr"] = lr
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        total = 0.0
        for _ in range(args.accum):
            loss = one_step(tr, buckets.batch(args.batch), ds, args, live, gen, device)
            (loss / args.accum).backward()
            total += float(loss.detach())
        if args.clip > 0:
            torch.nn.utils.clip_grad_norm_(params, args.clip)
        opt.step()
        opt.zero_grad(set_to_none=True)

        sec = time.time() - t0
        t_train += sec
        peak = torch.cuda.max_memory_allocated() / 2**30
        peak_all = max(peak_all, peak)
        loss_avg = total / args.accum
        if args.log_every and (step % args.log_every == 0 or step == 1):
            eta = t_train / (step - args.start_step) * (steps - step) / 3600
            print(f"шаг {step}/{steps}  loss {loss_avg:.4f}  lr {lr:.1e}  {sec:.1f} c/шаг  "
                  f"пик {peak:.2f} ГиБ  осталось {int(eta)}ч{int(eta % 1 * 60):02d}м", flush=True)
            with open(log_path, "a") as fh:
                fh.write(f"{step},{loss_avg:.6f},{lr:.2e},{sec:.2f},{peak:.2f},{eta:.2f}\n")

        if args.save_every and step % args.save_every == 0:
            path = os.path.join(args.out, f"step_{step:05d}")
            save_ckpt(tr, path, DTYPES[args.save_dtype])
            print(f"чекпойнт -> {path}", flush=True)
            old = sorted((d for d in os.listdir(args.out) if d.startswith("step_")),
                         key=lambda d: os.path.getmtime(os.path.join(args.out, d)))
            for name in (old[:-args.keep] if args.keep > 0 else []):
                shutil.rmtree(os.path.join(args.out, name), ignore_errors=True)

        if args.sample_every and (time.time() - t_last_sample >= args.sample_every or step == steps):
            try_sample(tr, sample_subset(), ds, args, live,
                       os.path.join(args.out, f"progress_{sample_n:05d}.png"))
            sample_n += 1
            t_last_sample = time.time()
            prune_samples(args.out, args.sample_keep)

    save_ckpt(tr, os.path.join(args.out, "final"), DTYPES[args.save_dtype])
    print(f"готово: {steps} шагов за {(time.time() - t_start) / 3600:.1f} ч, пик "
          f"{peak_all:.2f} ГиБ -> {args.out}/final", flush=True)


def probe(args, ds, buckets, device, gen, live):
    """Проверить, какой batch влезает: тот же режим, что в обучении (fp32-веса, bf16-градиенты,
    AdamW8bit), шаг на каждый размер батча от 1 до --batch, без записи."""
    tr = build_model(device)
    if not args.no_grad_ckpt:
        set_block_checkpointing(tr, True)
    params = [p for p in tr.parameters() if p.requires_grad]
    if args.grad_dtype == "bf16":
        set_grad_dtype(params, torch.bfloat16)
    import bitsandbytes as bnb
    opt = bnb.optim.AdamW8bit(params, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0)
    # самый тяжёлый размер: где влезет он, там влезут и остальные
    heavy = max(buckets.by_size, key=lambda k: k[0] * k[1])
    pool = [buckets.rows[i] for i in buckets.by_size[heavy]]
    print(f"пробую batch 1..{args.batch} на самом тяжёлом размере {heavy[0]}x{heavy[1]} "
          f"(строк {len(pool)}, градиенты {args.grad_dtype})", flush=True)
    ok = []
    for bs in range(1, args.batch + 1):
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        rows = [pool[int(i)] for i in torch.randint(0, len(pool), (bs,), generator=gen)]
        try:
            for _ in range(2):  # второй шаг важнее: состояния оптимизатора появляются на первом
                one_step(tr, rows, ds, args, live, gen, device).backward()
                if args.clip > 0:
                    torch.nn.utils.clip_grad_norm_(params, args.clip)
                opt.step()
                opt.zero_grad(set_to_none=True)
            free, _ = torch.cuda.mem_get_info()
            print(f"  batch {bs}: OK, пик {torch.cuda.max_memory_allocated() / 2**30:.2f} ГиБ, "
                  f"свободно {free / 2**30:.2f} ГиБ", flush=True)
            ok.append(bs)
        except torch.OutOfMemoryError:
            print(f"  batch {bs}: OOM", flush=True)
            break
    print(f"\nвлезает batch: {ok}; максимум {max(ok) if ok else 0}", flush=True)


if __name__ == "__main__":
    main()