Каждый год выпускают новый iPhone, который, судя по заявлениям, быстрее и лучше прошлого по всем параметрам. И да, эти новые модели компьютерного зрения и новые датчики изображений могут использовать возможности телефона по максимуму. Однако и десять лет назад на iPhone можно было делать хорошие снимки. Такие улучшения оказываются инкрементными.

Инкрементные запросы требуют только инкрементных улучшений. Впрочем, раз в несколько лет появляются программы, которые едва шевелятся даже на самых мощных вычислительных устройствах. Однако эти новые программы с новыми возможностями настолько великолепны, что люди готовы смириться со страданиями.

В последний раз такое происходило с глубокими нейронными сетями, а до этого — с 3D-графикой. Наверно, сейчас настал третий раз. На самом деле, я настолько в этом уверен, что создал приложение, чтобы подтвердить свою точку зрения.

За три недели я создал приложение, позволяющее генерировать изображения на основе текстовых описаний, а затем редактировать их нужным вам образом. На самом последнем и мощном iPhone 14 Pro создание изображения заняло около минуты, потребовало около 2 ГиБ памяти приложения, а для начала работы ему необходимо скачать примерно 2 ГиБ данных. И хотя само приложение работает абсолютно надёжно, учитывая эти требования, вряд ли его можно считать пригодным к практическому применению.

Даже несмотря на то, что для создания одного изображения требовалась целая минута, альбом «Фотоплёнка» теперь заполнен рисунками из этого приложения. Это очень увлекательное занятие. Более того, я осваиваю его всё лучше. Если лицо обрезано, то я знаю, как использовать модель inpainting, чтобы дополнить его. Если модель inpainting не справляется, всегда можно использовать кисть, чтобы нарисовать его, а затем выполнить генерацию image-to-image, сфокусировавшись на этой области.

А теперь давайте поговорим о том, как мне это удалось.

Как оказалось, запустить Stable Diffusion на iPhone проще, чем я думал, и при этом по-прежнему можно улучшить производительность в два раза. Вопрос лишь в куче аспектов. Основная сложность заключалась в запуске приложения на устройствах iPhone с 6 ГиБ ОЗУ. Кажется, что 6 ГиБ — это много, однако iOS начинает закрывать приложение, если вы используете больше 2,8 ГиБ на устройстве с 6 ГиБ и больше 2 ГиБ на устройстве с 4 ГиБ.

Исходная опенсорсная версия Stable Diffusion не может работать даже на карте с 8 ГиБ, а ведь все эти 8 ГиБ можно использовать. Но давайте сначала разберёмся с основами. Сколько конкретно нужно памяти модели Stable Diffusion для создания результатов?

Модель состоит из 4 частей: 1. Текстового кодировщика, генерирующего векторы текстовых признаков для управления генерацией изображения. 2. Опционального кодировщика изображений для кодирования изображения в скрытое пространство (latent space) (для генерации image-to-image). 3. Модели устранения шума, постепенно очищающего от шума описание изображения в скрытом пространстве. 4. Декодера изображений, декодирующего изображение из описания в скрытом пространстве. Части 1, 2 и 4 при создании результатов должны выполняться только один раз. Они относительно малозатратны (максимум около 1 ГиБ). Веса модели устранения шума занимают 3,2 ГиБ (в полных числах с плавающей запятой) от исходных весов модели объёмом 4,2 ГиБ. Кроме того, при исполнении они выполняются многократно, поэтому должны находиться в ОЗУ дольше.

Так почему исходная модель Stable Diffusion требует для создания одного изображения почти 10 ГиБ? Наряду с другими весами, которые мы не выгружаем (примерно 1 ГиБ в полных числах с плавающей запятой), требуется множество промежуточных распределений. Между одним вводом (2x4x64x64) и одним выводом (2x4x64x64) существует множество слоёв вывода. Не все слои вывода можно сразу же использовать на следующем этапе. Некоторые из них из-за сетевых структур необходимо хранить для дальнейшего использования (остаточные сети). Кроме того, PyTorch использует библиотеки NVIDIA CUDNN и CUBLAS. Эти библиотеки тоже хранят собственное рабочее пространство. После публикации модели было внесено множество оптимизаций в версию PyTorch для Stable Diffusion с целью снижения объёма необходимой памяти, чтобы её можно было запускать на картах всего с 4 ГиБ.

Но это всё равно чуть больше того, что мы можем себе позволить. Далее мы рассмотрим оборудование Apple и оптимизацию.

3,2 ГиБ или 1,6 ГиБ в половинных числах с плавающей запятой — вот начальная начальная точка, с которой мы будем работать. У нас есть примерно 500 МиБ пространства для работы, если мы не хотим приближаться к той границе, после которой OOM killer Apple может «убить» приложение.

Возникает первый вопрос: каков точный размер каждого из промежуточных выводов?

Оказывается, большинство из них относительно мало, каждый меньше 6 МиБ (2x320x64x64). Используемый мной фреймворк (s4nnc) вполне неплохо справляется с их bin-упаковкой в менее чем 50 МиБ с учётом повторного использования и тому подобного. Далее идёт очень интересный момент. Модель устранения шума имеет механизм self-attention, при котором в качестве ввода используется его собственное описание изображения в скрытом пространстве. При вычислении self-attention используется сгруппированная матрица размером (16x4096x4096). Это примерно равно 500 МиБ в половинных числах с плавающей запятой (FP16). Позже мы применяем к этой матрице функцию softmax. Это ещё 500 МиБ в FP16. Аккуратную реализацию softmax можно выполнить «с замещением», то есть она может без повреждения данных перезаписать свой ввод. К счастью, в низкоуровневых библиотеках Apple и NVIDIA есть реализация softmax с замещением. К сожалению, в высокоуровневых библиотеках наподобие PyTorch её нет.

Что ж, мы едва умещаемся, но вроде бы можем получить что-то около 550 МиБ + 1,6 ГиБ?

На оборудовании Apple популярным способом реализации бэкенда нейросетей является фреймворк MPSGraph. Это достаточно многофункциональный фреймворк, имеющий статический механизм построения графа вычислений (также называемый «TensorFlow»). Людям он нравится за свою достаточную эргономичность и производительность (он имеет все удобные штуки наподобие broadcast semantics). В поддержке PyTorch новой архитектуры M1 есть большая доля кода, реализованного при помощи MPSGraph.

Для первого прохода я реализовал все операции нейронных сетей при помощи MPSGraph. На пике он использует около 6 ГиБ (!!!) с точностью FP16. Что происходит?

В начале я должен признаться, что использую MPSGraph не совсем обычным способом (то есть с TensorFlow). Вероятно, MPSGraph ожидает, что вы будете кодировать весь граф вычислений, а затем подавать ему тензоры ввода/вывода. Затем он обрабатывает все внутренние распределения за пользователя и позволяет передать весь граф для исполнения. Я использую MPSGraph почти так же, как его использует PyTorch: как движок исполнения операций. Поэтому для создания результатов множество скомпилированных MPSGraphExecutable исполняется в очереди команд Metal, так как каждый из них может содержать какие-то промежуточные распределения. Если отправить их всех одновременно, все они будут до завершения исполнения хранить распределения на момент отправки.

Простой способ решения этой проблемы заключается в пошаговой отправке. Нет никаких причин отправлять их все за раз; кроме того, у Metal есть ограничение в 64 одновременных передач на очередь. Я ограничил передачи восьмью одновременными операциями, и это снизило пиковое использование памяти до 4 ГиБ.

Но это всё равно на 2 ГиБ больше, чем мы можем позволить себе на iPhone. Что делать дальше? Сначала нужно рассказать одну историю: при вычислении self-attention с помощью CUDA применяется стандартный трюк, используемый и в коде исходной модели Stable Diffusion — это перестановка вместо транспозиций (подробнее о том, что я имею в виду, можно прочитать в Transformers from the Scratch). Это помогает, поскольку CUBLAS может напрямую работать с периодически заполненными тензорами с перестановками, что позволяет избежать использования одного трафика памяти для транспозиции тензора.

Однако у MPSGraph нет поддержки периодически заполненных тензоров. Поэтому тензор с перестановками всё равно внутри будет подвергнут транспозиции, и для этого требуется ещё одно промежуточное распределение. При явной транспозиции распределение обрабатывается слоем более высокого уровня, что позволяет избежать неэффективности внутри MPSGraph. Этот трюк снижает потребление памяти примерно до 3 ГиБ.

Осталось ещё 1 ГиБ, и настало время для ещё одной истории! До MPSGraph были Metal Performance Shaders. Это коллекция фиксированных примитивов Metal для операций с нейронными сетями. Можно считать MPSGraph более современными скомпилированными just-in-time шейдерами, а Metal Performance Shaders — более старой и надёжной альтернативой.

Оказывается, MPSGraph в iOS 16.0 не принимает оптимального решения по распределениям для softmax. Даже если тензоры ввода и вывода указывают на одинаковые данные, MPSGraph распределяет дополнительный тензор вывода, а затем копирует результат в место, на которое мы указали. Из-за вышеупомянутого тензора на 500 МиБ это не очень аккуратно с точки зрения использования памяти. Альтернатива в виде Metal Performance Shaders выполняет именно то, что нам нужно, и это снижает объём используемой памяти до 2,5 ГиБ совершенно без уменьшения производительности.

Та же история произошла с ядром GEMM MPSGraph: некоторые GEMM требуют внутренних транспозиций, а для них нужны внутренние распределения (вместо того, чтобы просто использовать периодически заполненный тензор напрямую для умножения, как поступали GEMM из Metal Performance Shaders или CUBLAS. Однако любопытно, что в слое MLIR GEMM внутри MPSGraph, похоже, и в самом деле поддерживает параметры транспозиции (без дополнительного распределения), как большинство других ядер GEMM). Перемещение этих транспозиций явным образом тоже не помогает, поскольку транспозиции не являются операциями «с замещением» для слоя более высокого уровня, поэтому конкретно для этого тензора размером 500 МиБ дополнительного распределения избежать невозможно. Перейдя обратно на Metal Performance Shaders, мы вернули себе ещё 500 МиБ, потеряв всего около 1% производительности. Так мы наконец получили размер 2 ГиБ, к которому стремились изначально.

В систему по-прежнему можно внести множество улучшений производительности. Я так и не перешёл на ANE, хотя наконец-то разобрался в том, как это сделать (для этого требуется конкретная форма/периодичность входящих данных свёртки, а для них можно включить загадочный флаг OptimizationLevel1). Использование Int8 для свёртки кажется надёжным решением (я посмотрел на величины этих весов, ни один из них не превосходит волшебного значения 6) и может уменьшить размер модели и занимаемую память ещё примерно на 200 МиБ. Мне нужно было перенести модуль attention в созданный самостоятельно, аналогично FlashAttention или XFormer на стороне CUDA. Всё вместе это бы, вероятно, снизило размер на 30%, а занимаемую память ещё на 15%. Ну, оставим это на потом.

Скачать Draw Things можно отсюда: https://draw.nnc.ai/

Вот несколько ссылок по теме:

1. How to Draw Anything (перевод на Хабре) — эта статья оказала на меня в своё время очень большое влияние, рекомендую прочитать её каждому, кого интересует тема моего поста. В ней описывается процесс, при котором модели text-to-image можно использовать не просто как развлечение. Это реальный инструмент повышения продуктивности (с тех пор появилось множество альтернативных процессов, люди всё ещё исследуют возможности моделей).
2. Maple Diffusion — когда я работал над swift-diffusion, @madebyollin вёл параллельный проект по реализации Stable Diffusion непосредственно в MPSGraph. Из его эксперимента я понял, что структура NHWC может быть более эффективной на оборудовании M1, поэтому внёс соответствующие изменения.