File size: 7,532 Bytes
2e26e31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # 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`) —
не проверено, код в этой папке намеренно не оставлен.
|