У панядзелак CompVis выклалі вагі Stable Diffusion. У іх README напісана, што ў мадэлі “860M UNet and 123M text encoder” і што яна “runs on a GPU with at least 10GB VRAM”. Я перамножыў у галаве. Прыкладна мільярд лікаў па 4 байты гэта 4 ГБ. Астатнія 6 з 10 я лічыў запасам для CUDA і звычкай акругляць уверх.
Гэта і была ўся мая мадэль відэапамяці: вагі плюс трохі зверху. Я ніколі яе не правяраў, бо ўсё, што я запускаў раней, змяшчалася куды заўгодна. На машыне, дзе я раблю гэтыя доследы, відэакарты няма зусім. Там 4 ядры AMD EPYC і крыху менш за 5 ГіБ, якія я магу аддаць аднаму працэсу. Гэта аказалася зручна: працэс, які просіць зашмат, адразу забіваюць, а не даюць яму тармазіць. Усё ніжэй лічылася на PyTorch 1.12.1 на CPU ў float32.
Лічу вагі
У файлах CompVis/stable-diffusion-v1-4 толькі ў UNet ляжыць 859,5 мільёна лікаў. Разам з тэкставым энкодарам і VAE, часткай, якая ператварае вынік у пікселі, выходзіць 1,07 мільярда. Гэта 4,26 ГБ у float32 і 2,13 ГБ у палавіннай дакладнасці. Вагі я ацаніў правільна. Толькі тлумачыла гэтая ацэнка менш за палову з 10.
Адзін UNet у float32 займае 3,2 ГіБ. Загрузка падняла памяць працэсу да 3,76 ГіБ разам з файлавым кэшам. Пакуль у мяне не было чыстых лікаў, я доўга разбіраўся з забойствамі працэсаў, якія нічым не тлумачыліся. Потым знайшоў, што ў маёй карыстальніцкай сесіі свая столь у 5 ГіБ і ў яе ж лічыцца ўсё астатняе, што ў мяне адкрыта.
Куды сыходзіць астатняе
Адказ ляжыць у ldm/modules/attention.py. UNet глядзіць на малюнак праз слаі ўвагі. Першы з іх працуе на 1/8 памеру малюнка: на 512x512 ён бачыць 64 на 64, гэта значыць 4096 пазіцый. Кожная пазіцыя параўноўваецца з кожнай, атрымліваецца матрыца 4096 на 4096 лікаў. Галоў 8, а сэмплер ганяе промпт і пусты промпт разам, таму такіх матрыц 16. Код лічыць іх усе адным einsum, а потым softmax робіць другі тэнзар таго ж памеру.
16 матрыц па 64 МіБ даюць 1 ГіБ у float32, а разам з копіяй пасля softmax 2 ГіБ, гэта значыць 2048 МіБ. Я выняў гэты блок з UNet і запусціў асобна, з сапраўднымі вагамі. Памяць вырасла на 2241 МіБ звыш базы, так што 2 тэнзары тлумачаць 91% росту. Астатняе прыпадае на іншыя тэнзары блока. Потым я перапісаў тую ж матэматыку цыклам па галовах, па адной матрыцы за раз. Цыкл заняў 387 МіБ. Паўторны прагон даў 405. Выхады супалі біт у біт.
Я чакаў, што цыкл будзе даражэйшы па часе. За 3 прагоны кожнага варыянта ён заняў ад 2,04 да 2,16 секунды, а ўсе галовы адразу ад 1,85 да 2,19. На гэтым CPU розніца тоне ў шуме. На карце можа выйсці інакш, але на любой карце, дзе не хапае памяці, я б на такі абмен пайшоў.
Матрыца расце з квадратам колькасці пазіцый. Вось гэтага я не адчуваў, пакуль не ўбачыў. На 640x640 аднаму слою трэба 2,4 ГіБ, а на 768x768 ужо 5,1 ГіБ, яшчэ без копіі пасля softmax. З усімі галовамі адразу абодва прагоны забіла на ліміце ў 4 ГіБ. Па адной галаве яны занялі 795 МіБ і 1390 МіБ.
Каманда па змаўчанні
scripts/txt2img.py па змаўчанні робіць 3 малюнкі за раз, а сэмплер падвойвае гэта дзеля пустога промпта, так што UNet бачыць батч 6. Вагі ён трымае ў float32 і працуе пад autocast. Запыты і ключы выходзяць з лінейных слаёў у палавіннай дакладнасці, значыць і ацэнкі таксама. softmax стаіць у спісе аперацый, якія autocast вяртае ў float32. На адзін слой пры 512x512 у мяне выходзіць 1,5 ГіБ ацэнак у палавіннай дакладнасці і 3 ГіБ пасля softmax. Першы тэнзар код трымае жывым, пакуль працуе з другім. Наступны einsum ідзе праз bmm, які autocast прыводзіць да палавіннай дакладнасці, так што гэта яшчэ адна копія на 1,5 ГіБ. Разам 6 ГіБ на адзін слой, побач з 3,97 ГіБ вагаў.
Гэта блізка да 10 ГіБ, што нечакана добра сыходзіцца з README. На карце я гэтага не запускаў, так што гэта маё прачытанне кода, але памыліцца яно можа ў абодва бакі. Autocast можа яшчэ трымаць палавінныя копіі вагаў, якія ён прывёў, а колькі бярэ кантэкст CUDA, я не ведаю. Таму большую частку з 10 ГБ я адношу на рахунак батча 6 пад autocast па змаўчанні, а меншую на рахунак самой мадэлі.
Адзін малюнак цалкам
Напрыканцы я зрабіў адзін малюнак, каб пабачыць, што часткі працуюць разам. Узяў промпт “a photograph of an astronaut riding a horse”, 50 крокаў сэмплера PLMS і ўвагу па адной галаве. Промпт і пусты промпт ішлі праз UNet адным батчам з 2, як у кодзе CompVis. Кожная з 3 частак працавала асобным працэсам, так што ў памяці разам яны не былі ніколі. Працэс UNet дайшоў да 3,76 ГіБ. Ужо на загрузцы было 3,70, так што 51 праход сэмплера дадаў зверху толькі 64 МіБ. Астатнія 2 трымаліся ніжэй за 1,7 ГіБ.
На 4 ядрах усё разам заняло 16 хвілін, звычайны крок ішоў 17,4 секунды. Астранаўт на месцы, конь таксама. Колькі тая ж праца ідзе на карце, сказаць не магу: свайго ліку ў мяне няма.
На паперы ніжняя мяжа для 512x512 цяпер значна ніжэй за 10 ГБ. Гэта каля 2 ГіБ вагаў у палавіннай дакладнасці і каля 0,1 ГіБ ацэнак адной галавы, калі лічыць усе копіі, якія робіць autocast. Перш чым паверыць, я хацеў бы пабачыць гэта на сапраўднай карце.
Чаго я не правяраў
- Нічога тут не запускалася на GPU. Ні кантэкст CUDA, ні хуткасць на карце я не мераў.
- 10 ГіБ для каманды па змаўчанні ўзятыя з кода і задакументаваных правілаў autocast, а не з прагону.
- Кожны замер часу тут зроблены за 3 прагоны ці менш, на машыне, дзе ішла і іншая праца.
- Усе галовы адразу я мераў толькі на першым блоку ўвагі. Цэлы UNet я ганяў толькі з увагай па адной галаве. У тым прагоне працаваў код UNet з diffusers, а не мой цыкл.
- У VAE ёсць свая ўвага на 64 на 64 і асобна я яе не глядзеў.