| |
| """Файнтюн MOSS-Image DiT на закешированном датасете (`src/make_dataset.py`), цель — 16 ГБ. |
| |
| Что грузится: только сам DiT и тензоры из датасета. Текстовый энкодер не нужен вообще — |
| условие и маска уже посчитаны, негатив CFG лежит рядом с датасетом в `empty_condition.npz`. |
| VAE нужен ровно для одного — посмотреть прогресс-сэмпл; по умолчанию он выключен |
| (`--sample-every 0`), и тогда процесс не трогает ничего, кроме датасета и DiT. |
| |
| Почему влезает (замерено на этой машине, batch 1): |
| * bf16-веса 6.94 ГиБ — база, от неё никуда не деться; |
| * активации не держим: градиентный чекпойнтинг по блокам (`use_reentrant=False`), а входы |
| блоков лежат в RAM и поднимаются на карту на пересчёт (~0.15 с/шаг шины); |
| * градиенты живут на один блок: Adam применяется к блоку сразу после его backward |
| (`BlockUpdater`), градиенты блока сразу освобождаются. Обычный Adam требует градиенты всех |
| 3.72 B одновременно (6.94 ГиБ) и на 1024x1024 падает с OOM при пике 14.7 ГиБ; |
| * состояния Adam (2 байта × 2 на параметр = 7.45 ГиБ) живут в RAM — `PagedAdamW8bit`. |
| Итог: пик 8.28 ГиБ на 768x1024 (2.2 с/шаг), 8.60 ГиБ на 1024x1024 (3.6 с/шаг). |
| batch 2 нестабилен даже при 8 ГиБ запаса: managed-память paged-оптимизатора не успевает |
| вытесняться, и torch получает OOM на 15.0 ГиБ — проверено, поэтому batch 1. |
| |
| Про «fused forward/backward»: склеить forward и backward в один шаг нельзя — веса всё равно |
| нужны целиком. Память освобождает только «обновлять параметры сразу, как для них готов |
| градиент», и это здесь и сделано (BlockUpdater): пик падает с 14.7 до 8.6 ГиБ. Отличие от |
| обычного Adam одно — клип нормы считается по блоку, а не по всей модели. |
| |
| Лосс — flow matching, как в планировщике модели: |
| x_t = (1 - sigma) * x0 + sigma * noise, таргет = noise - x0, timestep = sigma * 1000. |
| Сетка обучения планировщика при `training_shift=1.0` равномерна по sigma, поэтому sigma — U(0,1). |
| |
| Запуск (по умолчанию — одна эпоха по датасету): |
| python3 src/train.py --steps 40 # смоук: 2 минуты |
| python3 src/train.py --epochs 1 --save-every 2000 # прогон |
| python3 src/train.py --epochs 2 # все датасеты в папке datasets |
| |
| LR: `--lr` — базовая скорость; и прогрев, и затухание идут от `--lr/10`: линейный прогрев |
| lr/10 -> lr за `--warmup` шагов, затем cosine от lr до lr/10 к последнему шагу. Постоянного lr |
| нет — на сошедшейся базе он просто держит лосс. |
| |
| Прогресс-сэмплы: раз в `--sample-every` секунд (по умолчанию час) из условий самого датасета |
| (сид фиксирован) -> `<out>/progress_<номер сэмпла>.png`; VAE на время сэмпла загружается и сразу |
| освобождается, на диске живут последние `--sample-keep` картинок. |
| |
| Куда пишется результат: только в `--out` (по умолчанию `src/work/train`). Веса модели |
| (`transformer/`) скрипт только читает; если `--out` указывает внутрь компонентов модели |
| (`transformer/`, `vae/`, `text_encoder/`, `adapter/`, `tokenizer/`, `scheduler/`) или на саму папку |
| модели — падаем до загрузки весов, чтобы случайно не перезаписать модель. Один чекпойнт — ~7.4 ГБ. |
| """ |
| import argparse |
| import gc |
| import glob |
| import math |
| import os |
| import random as random_mod |
| import re |
| 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 |
| from torch.utils.checkpoint import checkpoint |
|
|
| 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="MOSS-Image: файнтюн DiT под 16 ГБ") |
| p.add_argument("--dataset", default=os.path.join(MODEL, "src", "work", "ds_testg_filtered"), |
| help="папка с датасетами и чанками из make_dataset.py; всё найденное внутри " |
| "склеивается в один набор") |
| p.add_argument("--out", default=os.path.join(MODEL, "src", "work", "train_base"), |
| help="каталог логов, чекпойнтов и прогресс-сэмплов") |
| p.add_argument("--epochs", type=float, default=1.0, |
| help="эпох по выбранным строкам; шаги = ceil(epochs * строк / (batch * accum))") |
| p.add_argument("--steps", type=int, default=0, |
| help="точное число шагов, перебивает --epochs; 0 — считать из --epochs") |
| p.add_argument("--lr", type=float, default=1e-5) |
| p.add_argument("--warmup", type=int, default=200, |
| help="шагов линейного прогрева lr/10 -> lr") |
| p.add_argument("--batch-size", type=int, default=1, |
| help="сэмплов на шаг (только одного размера); на 16 ГиБ влезает только 1 — " |
| "batch 2 доходит до 14.85 ГиБ и падает с OOM") |
| p.add_argument("--accum", type=int, default=1, help="шагов накопления градиента") |
| p.add_argument("--clip", type=float, default=1.0, help="клип нормы градиента; 0 — выключить") |
| p.add_argument("--upd", choices=("block", "whole"), default="block", |
| help="block — Adam по блоку сразу после его backward (пик ~8 ГиБ); " |
| "whole — обычный Adam на всём наборе (нужно ~14.7 ГиБ, на 1024x1024 не влезает)") |
| p.add_argument("--master-dtype", choices=("fp32", "bf16"), default="fp32", |
| help="в чём лежат мастер-веса для шага оптимизатора: fp32 — bf16-копия идёт в " |
| "forward/backward, а шаг делается по fp32 (иначе при lr~1e-5 обновление " |
| "теряется в bf16-решётке); bf16 — как было, шаг прямо по bf16-весам") |
| p.add_argument("--master-device", choices=("cuda", "cpu"), default="cuda", |
| help="где живёт fp32-мастер: cuda — на карте (+13.9 ГиБ, шаг без копий через шину); " |
| "cpu — в RAM (на карте только bf16-веса, но градиенты блока едет по PCIe)") |
| p.add_argument("--te", action="store_true", default=True, |
| help="считать условие на лету живым текстовым энкодером. Нужно, когда в датасете " |
| "нет колонок text_embeds/text_mask (make_dataset.py без --te); с колонками " |
| "энкодер не грузится") |
| p.add_argument("--text-len", type=int, default=1024, |
| help="макс. токенов условия при подсчёте на лету (как в make_dataset.py)") |
| p.add_argument("--word-dropout", type=float, default=0.05, |
| help="доля шагов, где из промпта выкидывается одно случайное слово (устойчивость " |
| "к ошибкам в описании); работает только когда энкодер загружен (--te или нет " |
| "колонок TE). 0 — выключить") |
| p.add_argument("--text-dropout", type=float, default=0.05, |
| help="доля шагов с пустым промптом (CFG-ветка); без empty_condition.npz " |
| "молча выключается") |
| p.add_argument("--sigma", choices=("uniform", "logit"), default="logit", |
| help="откуда берётся sigma: uniform — U(0,1), как учили базу; logit — " |
| "sigmoid(N(--sigma-shift, --sigma-scale)), смещает шаги к нижим сигмам (детали)") |
| p.add_argument("--sigma-shift", type=float, default=-0.5, |
| help="logit-normal: среднее в logit-шкале; <0 — меньше шагов на высоком шуме, " |
| "больше на середине (детали)") |
| p.add_argument("--sigma-scale", type=float, default=1.0, help="logit-normal: разброс") |
| p.add_argument("--limit", type=int, default=0, help="первые N строк датасета (0 — все)") |
| p.add_argument("--sizes", default="", help="только эти размеры, напр. 768x1024,1024x768") |
| p.add_argument("--seed", type=int, default=42) |
| p.add_argument("--log-every", type=int, default=50) |
| p.add_argument("--ema", type=float, default=0.99, |
| help="сглаживание лосса в логе: ema = ema*--ema + loss*(1---ema); 0 — без сглаживания") |
| p.add_argument("--save-every", type=int, default=2000, |
| help="шагов между чекпойнтами; 0 — только в конце") |
| p.add_argument("--keep", type=int, default=4, |
| help="сколько step_* чекпойнтов держать на диске (старые удаляются)") |
| p.add_argument("--init", default="", help="стартовать с чекпойнта (step_0005000 или final)") |
| p.add_argument("--start-step", type=int, default=0, |
| help="с какого шага продолжать расписание lr при --init (иначе прогрев и cosine " |
| "пойдут заново с нуля — неверно при продолжении)") |
| p.add_argument("--sample-every", type=float, default=1800.0, |
| help="секунд между прогресс-сэмплами (1800 — раз в полчаса); 0 — выключить " |
| "(тогда VAE не грузится)") |
| p.add_argument("--sample-keep", type=int, default=8, |
| help="сколько progress_*.png держать на диске (старые удаляются); 0 — не удалять") |
| p.add_argument("--sample-n", type=int, default=6, help="картинок в прогресс-сэмпле") |
| p.add_argument("--sample-steps", type=int, default=30, help="шагов денойзинга в сэмпле") |
| p.add_argument("--sample-size", type=int, default=1024, help="длинная сторона сэмпла") |
| p.add_argument("--guidance", type=float, default=3.0) |
| p.add_argument("--sample-rows", default="", |
| help="индексы строк для прогресс-сэмплов через запятую, напр. 0,30000; " |
| "по умолчанию — равномерно по датасету") |
| p.add_argument("--sample-seed", type=int, default=42) |
| return p.parse_args() |
|
|
|
|
| |
| def find_dataset_chunks(root): |
| """Все arrow-датасеты под `root` (папки с dataset_info.json). |
| |
| make_dataset.py пишет не один датасет, а кучу под-датасетов `chunk_*/sub_*` — их и находим. |
| """ |
| chunks = [] |
| 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: |
| chunks.append(dirpath) |
| dirnames.clear() |
| return chunks |
|
|
|
|
| def load_any_dataset(path): |
| """Один датасет или папка с чанками: конкатенация ленивая, arrow-файлы не переписываются.""" |
| from datasets import concatenate_datasets, load_from_disk |
| chunks = find_dataset_chunks(path) |
| if not chunks: |
| raise FileNotFoundError(f"в {path} нет arrow-датасета (ни dataset_info.json, ни state.json)") |
| if len(chunks) == 1: |
| print(f"датасет: {chunks[0]}", flush=True) |
| return load_from_disk(chunks[0]) |
| print(f"датасет: {len(chunks)} чанков в {path}", flush=True) |
| return concatenate_datasets([load_from_disk(c) for c in chunks]) |
|
|
|
|
| def open_dataset(path, limit, sizes): |
| """Датасет не грузим в RAM: латент на строку — мегабайты, а 6k строк — это десятки ГБ. |
| |
| Возвращаем (ds, rows): `rows` — только метаданные (индекс, размер), сами тензоры |
| читаются по индексу в `to_batch` (27 мс на строку при 768x1024, незаметно на фоне шага). |
| """ |
| ds = load_any_dataset(path) |
| meta = ds.select_columns(["width", "height"]) |
| ds = ds.with_format("numpy") |
| wanted = {x.strip() for x in sizes.split(",") if x.strip()} |
| rows = [] |
| for i in range(len(meta)): |
| w, h = int(meta[i]["width"]), int(meta[i]["height"]) |
| if wanted and f"{w}x{h}" not in wanted: |
| continue |
| rows.append({"idx": i, "width": w, "height": h}) |
| if limit > 0 and len(rows) >= limit: |
| break |
| return ds, rows |
|
|
|
|
| def load_empty_condition(path): |
| """Негатив CFG, посчитанный make_dataset.py, или None. |
| |
| Ищем и в корне, и на уровень ниже: при `--dataset datasets` файл лежит в каждом под-датасете |
| (`datasets/civit/empty_condition.npz`), а в корне его нет — без этого сэмплы шли без CFG. |
| """ |
| files = [os.path.join(path, "empty_condition.npz")] |
| files += sorted(glob.glob(os.path.join(path, "*", "empty_condition.npz"))) |
| for file in files: |
| if os.path.exists(file): |
| z = np.load(file) |
| return (torch.from_numpy(z["embeds"].copy()), torch.from_numpy(z["mask"].copy())) |
| return None |
|
|
|
|
| class Buckets: |
| """Батчи только из одинаковых размеров: латент — один тензор, разные HxW не склеить.""" |
|
|
| def __init__(self, rows, generator): |
| self.rows = rows |
| self.generator = generator |
| 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(self.by_size[(w, h)])}" |
| for w, h in self.sizes), flush=True) |
| self._perms, self.cursor = {}, {} |
|
|
| def _shuffle(self): |
| for size, idx in self.by_size.items(): |
| order = list(idx) |
| for i in range(len(order) - 1, 0, -1): |
| j = int(torch.randint(0, i + 1, (1,), generator=self.generator)) |
| order[i], order[j] = order[j], order[i] |
| self._perms[size] = order |
| self.cursor[size] = 0 |
|
|
| def batch(self, batch_size): |
| if not self._perms: |
| self._shuffle() |
| |
| weights = [len(self.by_size[s]) for s in self.sizes] |
| pick = int(torch.multinomial(torch.tensor(weights, dtype=torch.float), 1, |
| generator=self.generator)) |
| size = self.sizes[pick] |
| picked = [] |
| for _ in range(batch_size): |
| if self.cursor[size] >= len(self._perms[size]): |
| self._perms[size] = self._perms[size][::-1] |
| self.cursor[size] = 0 |
| picked.append(self._perms[size][self.cursor[size]]) |
| self.cursor[size] += 1 |
| return [self.rows[i] for i in picked] |
|
|
|
|
| def read_row(ds, row): |
| r = ds[row["idx"]] |
| return {"vae": r["vae"], "text": r["text"], "text_embeds": r.get("text_embeds"), |
| "text_mask": r.get("text_mask"), "source": r["source"], "width": row["width"], |
| "height": row["height"]} |
|
|
|
|
| def to_batch(rows, ds, device): |
| """rows (одинаковый HxW) -> латент [B,128,h,w] bf16, условие [B,L,4096] bf16 + маска [B,L]. |
| |
| Если в датасете нет text_embeds (make_dataset.py без --te), условие собирается живым |
| энкодером в text_encoder.py — возвращаем строки, а не тензоры. |
| """ |
| recs = [read_row(ds, r) for r in rows] |
| lat = torch.stack([torch.from_numpy(r["vae"]) for r in recs]).to(device, torch.bfloat16) |
| if recs[0]["text_embeds"] is None: |
| return lat, None, None |
| length = max(len(r["text_embeds"]) for r in recs) |
| cond = torch.zeros(len(recs), length, 4096, dtype=torch.bfloat16) |
| mask = torch.zeros(len(recs), length, dtype=torch.bool) |
| for i, r in enumerate(recs): |
| n = len(r["text_embeds"]) |
| cond[i, :n] = torch.from_numpy(r["text_embeds"]) |
| mask[i, :n] = torch.from_numpy(r["text_mask"]) |
| return lat, cond.to(device), mask.to(device) |
|
|
|
|
| |
| class LiveTextEncoder: |
| """Текст условия живым ТЕ, когда в датасете нет колонок text_embeds/text_mask. |
| |
| Путь тот же, что в make_dataset.py (`encode_text`/`render`), чтобы условие совпало с тем, |
| что было бы закэшировано: chat-шаблон с add_generation_prompt=True и enable_thinking=True, |
| обрезка по --text-len, отдельный forward на промпт без паддинга (паддинг менял бы условие). |
| Возвращает (embeds float16 [L,4096], mask bool [L]). |
| """ |
|
|
| def __init__(self, model_dir, device, text_len): |
| """`device` — где считать; на карте энкодер держится только на время кодирования. |
| |
| 0.85B стоит ~2.3 ГиБ. Вместе с 20.8 ГиБ fp32-мастера и активациями DiT на 32 ГБ это |
| упирается в потолок, поэтому после encode_batch энкодер уходит в offload (см. `release`). |
| """ |
| 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 = torch.device(device) |
| self.offloaded = False |
| self.text_len = text_len |
|
|
| def release(self): |
| """Убрать энкодер с карты, если он там: память нужна шагу DiT.""" |
| if self.device.type == "cuda" and not self.offloaded: |
| self.te.to("cpu") |
| self.offloaded = True |
| torch.cuda.empty_cache() |
|
|
| def _ensure(self): |
| if self.offloaded: |
| self.te.to(self.device) |
| self.offloaded = False |
|
|
| def _render(self, text): |
| return self.tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=True) |
|
|
| @torch.no_grad() |
| def encode(self, text): |
| self._ensure() |
| ids = self.tok(self._render(text), truncation=True, max_length=self.text_len)["input_ids"] |
| if not ids: |
| raise ValueError("пустой промпт после токенизации") |
| ids = torch.tensor([ids], dtype=torch.long, device=self.device) |
| out = self.te(input_ids=ids, attention_mask=torch.ones_like(ids, dtype=torch.bool), |
| use_cache=False) |
| emb = out.last_hidden_state[0].to(torch.float16) |
| return emb.cpu().numpy(), np.ones(emb.shape[0], dtype=bool) |
|
|
| def encode_batch(self, texts): |
| """Строки промптов -> [B,L,4096] bf16 + маска [B,L] (L = max по батчу, без паддинга в ТЕ).""" |
| pairs = [self.encode(t) for t in texts] |
| length = max(len(e) for e, _ in pairs) |
| cond = torch.zeros(len(pairs), length, 4096, dtype=torch.bfloat16) |
| mask = torch.zeros(len(pairs), length, dtype=torch.bool) |
| for i, (emb, m) in enumerate(pairs): |
| cond[i, :len(emb)] = torch.from_numpy(emb) |
| mask[i, :len(emb)] = torch.from_numpy(m) |
| return cond, mask |
|
|
|
|
| def drop_word(text, gen, rng): |
| """Убрать одно случайное слово из промпта (один бросок на строку). |
| |
| Цель — чтобы модель была устойчива к пропускам и мусору в описании. Если слово одно — |
| остаётся пустая строка (это же и есть ветка пустого промпта). |
| """ |
| words = text.split() |
| if not words: |
| return text |
| del words[int(torch.randint(0, len(words), (1,), generator=gen))] |
| out = " ".join(words) |
| return out or rng.choice(["image", "a picture", "photo"]) |
|
|
|
|
| def _pair(module, hidden_dtype, master_dtype, device): |
| """Завести мастер-тензоры для Linear: (master_w, master_b), скрытые веса — в module.weight.""" |
| if not isinstance(module, torch.nn.Linear): |
| return |
| def empty_like(t): |
| return torch.empty(t.shape, dtype=master_dtype, device=device) |
| pairs = [empty_like(module.weight)] |
| if module.bias is not None: |
| pairs.append(empty_like(module.bias)) |
| module._master = pairs |
|
|
|
|
| class MasterWeights: |
| """fp32-мастер весов: forward/backward по bf16, шаг оптимизатора — по fp32-копии. |
| |
| Зачем: при lr ~1e-5 обновление `w -= lr*g` — это изменение ~1e-5 от |w|, а квант шага bf16 |
| (7 мантийных бит) ~2e-3 от |w|. Шаг не представим: он либо равен нулю, либо округляется |
| вверх до ступеньки решётки. Мастер в fp32 убирает именно это; веса для forward остаются bf16, |
| поэтому VRAM и скорость счёта не меняются. |
| |
| Устройство: `module.weight` всегда остаётся bf16 (`nn.Parameter` не подменяется — иначе |
| autograd и оптимизатор начинают спорить о dtype градиента). Рядом лежит fp32-тензор |
| `module._master[i]`, уже объявленный оптимизатору как параметр. Цикл шага: |
| |
| forward/backward по bf16-весам |
| grads() bf16-градиенты -> fp32-grad мастеров |
| opt.step() Adam в fp32 по мастерам (состояния тоже fp32) |
| sync() мастера -> bf16-веса для следующего forward |
| |
| Предел: fp32-мастер на всю модель — это +13.9 ГиБ (на карте или в RAM при --master-device cpu). |
| """ |
|
|
| def __init__(self, model, device, dtype=torch.float32): |
| self.device = torch.device(device) |
| self.master_dtype = dtype |
| self.slots = [] |
| self._map = None |
| hidden_dtype = next(model.parameters()).dtype |
| for module in model.modules(): |
| _pair(module, hidden_dtype, dtype, self.device) |
| for module in model.modules(): |
| masters = getattr(module, "_master", None) |
| if masters is None: |
| continue |
| hidden = [module.weight] + ([module.bias] if module.bias is not None else []) |
| self.slots += list(zip(hidden, masters)) |
|
|
| @torch.no_grad() |
| def init(self, model): |
| """Один раз, до создания оптимизатора: мастера = стартовые веса в fp32.""" |
| for hidden, master in self.slots: |
| master.copy_(hidden.data.to(self.master_dtype)) |
|
|
| def parameters(self): |
| return [master for _, master in self.slots] |
|
|
| def _master_of(self, hidden): |
| """fp32-двойник bf16-параметра (словарь строится один раз).""" |
| if self._map is None: |
| self._map = {id(h): m for h, m in self.slots} |
| return self._map.get(id(hidden)) |
|
|
| @torch.no_grad() |
| def grads(self): |
| """bf16-градиенты -> fp32-градиенты мастеров (None, если градиента нет).""" |
| self.grads_for([h for h, _ in self.slots]) |
|
|
| @torch.no_grad() |
| def grads_for(self, hidden_params): |
| """То же для конкретной группы (нужно BlockUpdater, который обновляет по блоку).""" |
| for h in hidden_params: |
| m = self._master_of(h) |
| if m is not None: |
| m.grad = None if h.grad is None else h.grad.to(m.device, self.master_dtype) |
|
|
| @torch.no_grad() |
| def sync(self): |
| """fp32-мастер -> bf16-веса для следующего forward.""" |
| for hidden, master in self.slots: |
| hidden.data.copy_(master) |
|
|
| @torch.no_grad() |
| def sync_one(self, hidden): |
| """fp32-мастер одного параметра -> его bf16-вес.""" |
| m = self._master_of(hidden) |
| if m is not None: |
| hidden.data.copy_(m) |
|
|
| |
| def _run_block(block, x_cpu, rope, mod): |
| """Вход блока лежит на CPU, на пересчёт поднимается обратно (см. set_block_checkpointing).""" |
| return block(x_cpu.to("cuda", non_blocking=True), rope, mod) |
|
|
|
|
| def set_block_checkpointing(tr, on): |
| """Пересчёт блоков на backward вместо хранения активаций — без этого 16 ГБ не хватает. |
| |
| Вход каждого блока (25 МиБ на 1024 px) держим в RAM, а не на карте: autograd сохраняет |
| только этот вход, поэтому выигрываем ~0.75-1.0 ГиБ GPU — ровно ту разницу, которой не |
| хватало на квадратных 1024x1024 и на управляемую память PagedAdamW. За шаг это ~0.15 с |
| копий шины, дешевле любого другого способа освободить гигабайт. |
| """ |
| 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(_run_block, _b, x.cpu(), rope, mod, use_reentrant=False)) |
| elif not on and "forward" in layer.__dict__: |
| del layer.forward |
|
|
|
|
| class BlockUpdater: |
| """Adam по каждому блоку сразу после его backward; градиенты блока освобождаются. |
| |
| Зачем: обычному Adam нужны градиенты всех 3.72 B параметров (6.94 ГиБ) одновременно — |
| вместе с весами это 13.88 ГиБ, и 1024x1024 в 16 ГБ уже не влезает (замерено: OOM на 14.7). |
| `register_post_accumulate_grad_hook` срабатывает, когда у параметра появился градиент; как |
| только все параметры блока их получили, блок обновляется и его градиенты обнуляются. |
| Пик: веса + активации одного блока ~ 8 ГиБ. |
| |
| Градиенты при этом те же, что в обычном проходе: backward идёт от последнего блока к первому, |
| и блок обновляется только после того, как посчитан и пересчитан на своих старых весах. |
| Отличие одно — клип локальный по блоку, а не глобальный (это единственное, что нельзя |
| посчитать без всех градиентов сразу). |
| """ |
|
|
| def __init__(self, tr, lr, clip, master=None): |
| import bitsandbytes as bnb |
| block_ids = {id(p) for layer in tr.layers for p in layer.parameters()} |
| groups = [[p for p in layer.parameters()] for layer in tr.layers] |
| groups.append([p for p in tr.parameters() if id(p) not in block_ids]) |
| self.groups = [g for g in groups if g] |
| self.clip = clip |
| |
| self.master = master |
| to_master = {} if master is None else {id(h): m for h, m in master.slots} |
| self.targets = [[to_master.get(id(p), p) for p in g] for g in self.groups] |
| self.opts = [bnb.optim.PagedAdamW8bit(t, lr=lr, betas=(0.9, 0.95), eps=1e-8) |
| for t in self.targets] |
| self.left = [] |
| for gi, params in enumerate(self.groups): |
| self.left.append([len(params)]) |
| for p in params: |
| p.register_post_accumulate_grad_hook(lambda _p, gi=gi: self._release(gi)) |
|
|
| def set_lr(self, lr): |
| for opt in self.opts: |
| for group in opt.param_groups: |
| group["lr"] = lr |
|
|
| def _release(self, gi): |
| self.left[gi][0] -= 1 |
| if self.left[gi][0] > 0: |
| return |
| params = self.targets[gi] |
| if self.master is not None: |
| self.master.grads_for(self.groups[gi]) |
| if self.clip > 0: |
| torch.nn.utils.clip_grad_norm_(params, self.clip) |
| self.opts[gi].step() |
| for p in params: |
| p.grad = None |
| for p in self.groups[gi]: |
| if self.master is not None: |
| self.master.sync_one(p) |
| p.grad = None |
| self.left[gi][0] = len(params) |
|
|
| def finish(self): |
| """Страховка: группа, которой не хватило градиентов, обновляется здесь.""" |
| for gi, left in enumerate(self.left): |
| params = self.groups[gi] |
| if left[0] == len(params): |
| continue |
| print(f"ВНИМАНИЕ: группе {gi} не хватило {left[0]} градиентов — обновляю в конце", |
| flush=True) |
| if any(p.grad is not None for p in params): |
| if self.clip > 0: |
| torch.nn.utils.clip_grad_norm_(params, self.clip) |
| self.opts[gi].step() |
| for p in params: |
| p.grad = None |
| left[0] = len(params) |
|
|
|
|
| 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 name in names[:-keep]: |
| os.remove(os.path.join(out, name)) |
|
|
|
|
| def sample_grid(tr, rows, ds, args, out_path, empty, live=None): |
| """Несколько картинок из условий датасета: VAE грузим на время сэмпла и сразу освобождаем.""" |
| from diffusers import AutoencoderKLFlux2 |
| from pipeline_moss_image import MOSSImagePipeline |
| set_block_checkpointing(tr, False) |
| tr.eval() |
| |
| 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) |
| cfg = args.guidance > 1.0 and (empty is not None or live is not None) |
| if cfg: |
| if live is not None: |
| e, m = live.encode("") |
| nemb = torch.from_numpy(e)[None].to("cuda", torch.bfloat16) |
| nmask = torch.from_numpy(m)[None].cuda() |
| else: |
| nemb, nmask = empty[0][None].cuda(), empty[1][None].cuda() |
| tiles = [] |
| for i, meta in enumerate(rows): |
| r = read_row(ds, meta) |
| long_side = args.sample_size |
| if r["width"] >= r["height"]: |
| w, h = long_side, round(long_side * r["height"] / r["width"] / 32) * 32 |
| else: |
| w, h = round(long_side * r["width"] / r["height"] / 32) * 32, long_side |
| print(f" сэмпл {i + 1}: {w}x{h} {r['text'][:200]}", flush=True) |
| if r["text_embeds"] is not None: |
| cond = torch.from_numpy(r["text_embeds"])[None].to("cuda", torch.bfloat16) |
| cmask = torch.from_numpy(r["text_mask"])[None].cuda() |
| else: |
| cond, cmask = live.encode_batch([r["text"]]) |
| cond, cmask = cond.to("cuda"), cmask.to("cuda") |
| tiles.append(pipe( |
| prompt_embeds=cond, |
| prompt_attention_mask=cmask, |
| negative_prompt_embeds=nemb if cfg else None, |
| negative_prompt_attention_mask=nmask if cfg else None, |
| height=h, width=w, num_inference_steps=args.sample_steps, |
| guidance_scale=args.guidance if cfg 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)) |
| rows_n = (len(tiles) + cols - 1) // cols |
| cell = args.sample_size |
| grid = Image.new("RGB", (cols * cell, rows_n * 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}" + ("" if cfg else " (без CFG: нет empty_condition.npz)"), |
| flush=True) |
| tr.train() |
| set_block_checkpointing(tr, True) |
|
|
|
|
| |
| PROTECTED = ("transformer", "vae", "text_encoder", "adapter", "tokenizer", "scheduler") |
|
|
|
|
| def guard_out(out): |
| """Не даём писать чекпойнты поверх модели: обучение — только в отдельный каталог.""" |
| out_abs = os.path.abspath(out) |
| model_abs = os.path.abspath(MODEL) |
| for name in [""] + list(PROTECTED): |
| protected = os.path.join(model_abs, name) if name else model_abs |
| if out_abs == protected or out_abs.startswith(protected + os.sep) and name: |
| raise SystemExit(f"--out {out_abs} лежит внутри модели ({protected}); " |
| f"выберите отдельный каталог, например src/work/train") |
|
|
|
|
| def main(): |
| args = parse_args() |
| guard_out(args.out) |
| torch.manual_seed(args.seed) |
| os.makedirs(args.out, exist_ok=True) |
| device = "cuda" |
|
|
| |
| |
| |
| |
| torch.cuda.set_per_process_memory_fraction(0.92) |
| print(f"потолок VRAM: {0.92 * torch.cuda.get_device_properties(0).total_memory / 2**30:.2f} ГиБ " |
| f"из {torch.cuda.get_device_properties(0).total_memory / 2**30:.2f}", flush=True) |
|
|
| ds, rows = open_dataset(args.dataset, args.limit, args.sizes) |
| if not rows: |
| raise SystemExit("датасет пуст") |
| print(f"строк: {len(rows)}", flush=True) |
|
|
| has_te_cols = "text_embeds" in ds.column_names |
| live = None |
| if args.te or not has_te_cols: |
| live = LiveTextEncoder(MODEL, device, args.text_len) |
| print(f"условие считается на лету ({'--te' if args.te else 'в датасете нет text_embeds'}), " |
| f"text-len {args.text_len}", flush=True) |
| empty = load_empty_condition(args.dataset) |
| if live is None and empty is not None: |
| print(f"пустой промпт (негатив CFG): {tuple(empty[0].shape)}", flush=True) |
| if live is None and args.text_dropout > 0 and empty is None: |
| print("empty_condition.npz нет -> text-dropout выключен", flush=True) |
| dropout = args.text_dropout if (live is not None or empty is not None) else 0.0 |
| |
| word_dropout = args.word_dropout if live is not None else 0.0 |
| if args.word_dropout > 0 and live is None: |
| print("word-dropout выключен: условие берётся из датасета, энкодер не загружен", flush=True) |
| if word_dropout: |
| print(f"дропауты: пустой промпт {dropout:.0%}, случайное слово {word_dropout:.0%}", |
| flush=True) |
|
|
| n_sample = min(args.sample_n, len(rows)) |
| if args.sample_rows: |
| sample_rows = [rows[int(x)] for x in args.sample_rows.split(",") if x.strip()] |
| else: |
| sample_rows = [rows[int(i)] for i in np.linspace(0, len(rows) - 1, n_sample)] |
| if args.sample_every: |
| text_col = ds.select_columns(["text"]) |
| print("прогресс-сэмплы берутся из строк:", flush=True) |
| for i, meta in enumerate(sample_rows, 1): |
| print(f" {i}) строка {meta['idx']} {meta['width']}x{meta['height']} " |
| f"{text_col[meta['idx']]['text'][:160]}", flush=True) |
|
|
| |
| |
| free_gib = torch.cuda.mem_get_info()[0] / 2**30 |
| if free_gib < 15.0: |
| print(f"ВНИМАНИЕ: свободно {free_gib:.1f} ГиБ — обучению нужно ~14.7. " |
| f"Проверьте, не висит ли другой процесс на GPU (nvidia-smi)", flush=True) |
| tr = MOSSImageTransformer.from_pretrained( |
| os.path.join(MODEL, "transformer"), torch_dtype=torch.bfloat16, |
| local_files_only=True, low_cpu_mem_usage=False).to(device).train() |
| if args.init: |
| state = MOSSImageTransformer.from_pretrained( |
| args.init, torch_dtype=torch.bfloat16, local_files_only=True, |
| low_cpu_mem_usage=False).state_dict() |
| tr.load_training_state_dict(state) |
| print(f"старт с {args.init}", flush=True) |
| tr.requires_grad_(True) |
| params = [p for p in tr.parameters() if p.requires_grad] |
| n_param = sum(p.numel() for p in params) |
| print(f"параметров {n_param / 1e9:.3f}B, веса bf16 {n_param * 2 / 2**30:.2f} ГиБ, " |
| f"градиенты — только текущего блока", flush=True) |
| master = None |
| if args.master_dtype == "fp32": |
| |
| master = MasterWeights(tr, args.master_device) |
| master.init(tr) |
| extra = f"+{n_param * 4 / 2**30:.2f} ГиБ на {args.master_device}" |
| print(f"мастер-веса fp32 ({extra}), forward/backward остаются bf16", flush=True) |
| else: |
| print("мастер-веса bf16: шаг оптимизатора идёт прямо по bf16-весам", flush=True) |
| set_block_checkpointing(tr, True) |
|
|
| if args.upd == "block": |
| if args.accum != 1: |
| raise SystemExit("--upd block требует --accum 1: обновление идёт внутри backward " |
| "и не может ждать следующего микро-батча") |
| updater, opt = BlockUpdater(tr, args.lr, args.clip, master), None |
| print(f"оптимизатор: block-wise PagedAdamW8bit ({len(updater.groups)} групп, " |
| f"градиенты живут только на один блок)", flush=True) |
| else: |
| import bitsandbytes as bnb |
| updater = None |
| targets = master.parameters() if master is not None else params |
| opt = bnb.optim.PagedAdamW8bit(targets, lr=args.lr, betas=(0.9, 0.95), eps=1e-8) |
| print("оптимизатор: общий PagedAdamW8bit " |
| f"({'fp32-мастера' if master is not None else 'bf16-веса'})", flush=True) |
|
|
| gen = torch.Generator().manual_seed(args.seed) |
| buckets = Buckets(rows, gen) |
| if args.sigma == "logit": |
| probe = 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"медиана {probe.median():.2f}, доля шагов sigma<0.2: {(probe < 0.2).float().mean():.0%}, " |
| f"sigma>0.8: {(probe > 0.8).float().mean():.0%}", flush=True) |
| else: |
| print("sigma: uniform U(0,1), как в обучении базы", flush=True) |
| |
| |
| step_rows = args.batch_size * args.accum |
| new_steps = args.steps or max(1, math.ceil(args.epochs * len(rows) / step_rows)) |
| start_step = args.start_step |
| steps = new_steps + start_step |
| min_lr = args.lr / 10 |
| print(f"шагов: {new_steps}" + (f" ({args.epochs:g} эпох x {len(rows)} строк / {step_rows})" |
| if not args.steps else " (--steps)") |
| + (f", продолжаю с {start_step} -> до {steps}" if start_step else "") |
| + f", lr {args.lr:.1e} -> {min_lr:.1e}, warmup {args.warmup}", flush=True) |
| log_path = os.path.join(args.out, "log.csv") |
| with open(log_path, "a" if start_step else "w") as log: |
| if not start_step: |
| log.write("step,loss_ema,loss_raw,lr,sec,peak_gib,eta_h\n") |
|
|
| def lr_at(step): |
| """Прогрев lr/10 -> lr, затем cosine lr -> lr/10. Постоянного lr нет.""" |
| if args.warmup > 0 and step < args.warmup: |
| return min_lr + (args.lr - min_lr) * (step + 1) / args.warmup |
| progress = 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 * progress)) |
|
|
| ema = None |
| sec_sum = 0.0 |
| t_start = time.time() |
| t_last_sample = time.time() |
| sample_n = 0 |
| torch.cuda.reset_peak_memory_stats() |
| if args.sample_every: |
| sample_grid(tr, sample_rows, ds, args, |
| os.path.join(args.out, "progress_00000.png"), empty, live) |
| sample_n += 1 |
| for step in range(start_step + 1, steps + 1): |
| lr = lr_at(step - 1) |
| if updater is not None: |
| updater.set_lr(lr) |
| else: |
| for group in opt.param_groups: |
| group["lr"] = lr |
| t0 = time.time() |
| total = 0.0 |
| for _ in range(args.accum): |
| batch = buckets.batch(args.batch_size) |
| lat, cond, mask = to_batch(batch, ds, device) |
| if live is not None: |
| texts = [read_row(ds, r)["text"] for r in batch] |
| dropped = False |
| |
| if dropout and torch.rand(1, generator=gen).item() < dropout: |
| texts, dropped = [""] * len(batch), True |
| elif word_dropout and torch.rand(1, generator=gen).item() < word_dropout: |
| i = int(torch.randint(0, len(batch), (1,), generator=gen)) |
| texts[i] = drop_word(texts[i], gen, random_mod) |
| dropped = True |
| if dropped or cond is None: |
| cond, mask = live.encode_batch(texts) |
| cond, mask = cond.to(device), mask.to(device) |
| live.release() |
| del texts, dropped |
| noise = torch.randn(lat.shape, dtype=torch.float32, device=device) |
| if args.sigma == "logit": |
| sigma = torch.sigmoid(torch.randn(len(batch), 1, 1, 1, device=device) |
| * args.sigma_scale + args.sigma_shift) |
| else: |
| sigma = torch.rand(len(batch), 1, 1, 1, device=device) |
| x0 = lat.float() |
| xt = ((1 - sigma) * x0 + sigma * noise).to(torch.bfloat16) |
| target = noise - x0 |
| pred = tr(xt, sigma.flatten() * 1000.0, cond, mask).sample.float() |
| loss = F.mse_loss(pred, target) / args.accum |
| loss.backward() |
| total += loss.item() |
| del lat, cond, mask, xt, target, pred, loss |
| if updater is not None: |
| updater.finish() |
| else: |
| if master is not None: |
| master.grads() |
| params_to_clip = master.parameters() |
| else: |
| params_to_clip = params |
| if args.clip > 0: |
| torch.nn.utils.clip_grad_norm_(params_to_clip, args.clip) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| if master is not None: |
| master.sync() |
| step_sec = time.time() - t0 |
| sec_sum += step_sec |
| peak = torch.cuda.max_memory_allocated() / 2**30 |
| raw = total / args.accum |
| ema = raw if ema is None or args.ema <= 0 else args.ema * ema + (1 - args.ema) * raw |
|
|
| if step % args.log_every == 0 or step == 1: |
| done = step - start_step |
| eta = sec_sum / max(done, 1) * (steps - step) |
| eta_h = eta / 3600 |
| print(f"шаг {step}/{steps} loss {ema:.4f} (raw {raw:.3f}) lr {lr:.1e} " |
| f"{step_sec:.1f} c/шаг пик {peak:.2f} ГиБ " |
| f"осталось {int(eta_h)}ч{int(eta % 3600 // 60):02d}м", flush=True) |
| with open(log_path, "a") as log: |
| log.write(f"{step},{ema:.6f},{raw:.6f},{lr:.2e},{step_sec:.2f},{peak:.2f}," |
| f"{eta_h:.2f}\n") |
|
|
| if args.save_every and step % args.save_every == 0: |
| path = os.path.join(args.out, f"step_{step:05d}") |
| tr.save_pretrained(path) |
| 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) |
| print(f"старый чекпойнт удалён: {name}", flush=True) |
|
|
| if args.sample_every and (time.time() - t_last_sample >= args.sample_every |
| or step == steps): |
| sample_grid(tr, sample_rows, ds, args, |
| os.path.join(args.out, f"progress_{sample_n:05d}.png"), empty, live) |
| sample_n += 1 |
| t_last_sample = time.time() |
| prune_samples(args.out, args.sample_keep) |
|
|
| tr.save_pretrained(os.path.join(args.out, "final")) |
| peak = torch.cuda.max_memory_allocated() / 2**30 |
| print(f"готово: {steps} шагов за {time.time() - t_start:.0f} c, пик VRAM {peak:.2f} ГиБ, " |
| f"чекпойнт {os.path.join(args.out, 'final')}, лог {log_path}", flush=True) |
| if peak > 15.5: |
| print("ВНИМАНИЕ: пик близко к 16 ГиБ — на другой карте шаг может не влезть", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
