| |
| """Обучение 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 |
|
|
|
|
| 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): |
| 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() |
|
|
