#!/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()