| # src — чем собран трансформер в этой папке |
|
|
| Здесь сжатие DiT и обучение на своём датасете: **30 блоков вместо 42, 3.723 B вместо 5.082 B |
| (−26.7 %)** — качество на уровне исходной модели на прозе, людях и анатомии. Итоговые веса лежат в |
| `../transformer/`. |
|
|
| Чтобы повторить прогон, нужен исходный 42-блочный трансформер: он удалён из папки модели, скачать |
| заново из [`OpenMOSS-Team/MOSS-Image-Base`](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) |
| в `transformer/` (скрипты читают именно этот путь). |
|
|
| | файл | что делает | |
| |---|---| |
| | `window.py` | всё сжатие: `cache` (состояния учителя) → `train` (LoRA окна) → `eval` (замер окна) → `combo` (собрать модель из слияний) | |
| | `lab.py` | харнесс: сборка пайплайна, текстовые условия, сэмплинг, LPIPS + CLIP, удаление блоков | |
| | `run_hi.sh` | точная последовательность, которой собран трансформер в этой папке | |
| | `make_dataset.py` | картинки `*.jpg` + парные `*.txt` → латенты VAE, условие адаптера и маска; кладёт чанки `chunk_*/sub_*` в `--save` (+ `empty_condition.npz` — негатив CFG) | |
| | `train.py` | файнтюн DiT на этом датасете под 16 ГБ: чекпойнтинг по блокам, Adam по блоку, paged-оптимизатор; читает папку с чанками целиком | |
| | `work/` | кэши, чекпойнты LoRA, прогоны и метрики (пересоздаётся, в артефакт не входит) | |
|
|
| ## Методика |
|
|
| **Идея.** Блоки в середине DiT меняют поток мало и частично отменяют друг друга. Значит из двух-трёх |
| смежных блоков можно оставить один, добавив ему низкоранговую поправку, обученную повторять |
| поведение исходной модели. Глубина падает, переобучение всей сети не нужно. |
|
|
| **Схема окна.** Блок `s` остаётся и получает LoRA (rank 32 на `q,k,v,o,ffn_up_gate,ffn_down`, |
| 1.78 M параметров); блоки `s+1, s+2` удаляются. LoRA учится так, чтобы выход блока `s` совпал с |
| выходом блока `s+2` **исходной** модели. |
|
|
| Точный протокол (`window.py cache` + `window.py train`): |
|
|
| 1. Кэш: для N сэмплов (промпт из `~/dataset/civit/*.txt`, случайный латент, случайный `t`) делаем |
| пару грубых шагов Эйлера исходной моделью, чтобы латент был похож на точку реальной траектории, |
| затем один forward с хуками. Пишем вход блока `s` (выход блока `s-1`) и выход блока `s+2`, |
| плюс `rope` и `modulation` (без них блок не посчитать). |
| 2. Обучение: собранный из весов один блок (226 МБ) + LoRA, всё остальное не грузится вообще. |
| Лосс — относительная MSE по токенам `‖pred−target‖²/‖target‖²` плюс `0.1·(1−cos)`, чтобы |
| держать и направление. Батчи собираются по одинаковой длине последовательности (у промптов |
| разное число токенов, паддинг менял бы поведение блока). AdamW8bit, lr 1e-3, клип 1.0, |
| 3000 шагов, batch 2. |
| 3. Приёмка окна: `held`-ошибка на отложенных сэмплах (10 %), затем `eval` — генерация сетки и |
| LPIPS/CLIP против базы. |
| 4. Слияние в веса: `W' = W + B·A·scale`, поправка исчезает как отдельный объект, инференс не |
| замедляется. |
|
|
| **Окна.** Семь слияний по два блока: |
|
|
| ``` |
| 4,5 → 3 7,8 → 6 10,11 → 9 13,14 → 12 16,17 → 15 19,20 → 18 22,23 → 21 |
| ``` |
|
|
| В артефакт пошли первые шесть (42 − 12 = 30 блоков). Седьмое окно даёт 28 блоков (combo7): CLIP |
| 29.34 против 30.14 у combo6, и на глаз хуже — не используется. Восьмое (25,26 → 24) роняет CLIP |
| до 25.4 — зона «громкая», слияние там не работает. |
|
|
| **Почему поправки обучались в 1024 px.** Первый прогон учил слияния на 512 px: CLIP на 512 был |
| 30.14 против 29.49 у базы, но на 1024 падал до 28.92. Переобучение тех же окон на 1024 px дало |
| 29.12 на 1024 (`run_hi.sh` — именно этот прогон); каждая сборка лучше «на своём» разрешении, |
| в артефакт пошла версия для 1024. |
|
|
| ## Проверить, что всё работает |
|
|
| ```bash |
| cd .. # папка модели (там должен лежать исходный transformer/) |
| python3 src/window.py cache --start 3 --drop 2 --samples 2 --size 1024 # ~20 с |
| python3 src/window.py train --start 3 --drop 2 --samples 2 --size 1024 --steps 40 # ~10 с |
| python3 src/window.py combo --tags win3_5_1024,win6_8_1024 --tag check --size 1024 --sample-steps 10 |
| ``` |
|
|
| `combo` соберёт модель из указанных слияний, сгенерит сетку по 10 промптам и посчитает LPIPS/CLIP |
| против `work/runs/base/` (её надо один раз создать: `python3 src/lab.py baseline`). |
|
|
| ## Что не сработало (чтобы не повторять) |
|
|
| * **Переподгонка на собранной модели** (дообучить те же LoRA на собственных траекториях собранной |
| модели, цели — исходная модель). Дрейф по границам блоков оказался 0.15–8 %, переподгонка его |
| дожимала до 0.1–0.2 %, но CLIP падал 29.1 → 23.3 и **лица разрушались**. Метод лечил то, что |
| почти не болело. |
| * **Текст в кадре** слиянием не лечится: вывески и JSON-layout у 30-блочной модели деградируют |
| (у базы «SALE 50% / only today», у сжатой — нечитаемая каша). На уровне состояний это не дрейф |
| (2–8 %), значит потеря способности; нужен flow-matching на парах с надписями. |
| * **Сужение FFN** (8192 → 6144, ещё −13 % параметров через закрытый пересчёт `ffn_down`) — |
| не проверено, код в этой папке намеренно не оставлен. |
|
|
