iPhone на пределе: 2 Гб модель, которая может рисовать все, что угодно, у вас в кармане

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

Эти дополнительные запросы заслуживают только дополнительных улучшений. Раз в несколько лет появляются программы, которые даже на самых лучших наших вычислительных устройствах едва ли можно использовать. Но эти новые программы с новыми сценариями настолько хороши, что люди готовы терпеть.

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

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

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

Теперь достанем кота из коробки, поговорим о том, как.

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

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

Модель состоит из 4 частей: текстовый кодировщик, который генерирует векторы текстовых признаков для управления генерацией изображения. Дополнительный кодировщик изображений для кодирования изображения в скрытое пространство (для image-to-image преобразования). Модель шумоподавителя, которая медленно удаляет шум из скрытого представления. Декодер изображения для декодирования изображения из этого скрытого представления. 1-я, 2-я и 4-я модели должны запускаться один раз во время вывода. Они относительно дешевы (максимум около 1 Гб). Веса модели шумоподавителя занимают 3.2 Гб (в полных числах с плавающей запятой) от исходных весов модели 4.2 Гб. Она также должна запускаться несколько раз за выполнение, поэтому мы хотим дольше держать его в оперативной памяти.

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

Это все еще немного больше, чем мы можем себе позволить. Но сейчас я сосредоточусь на оборудовании Apple и оптимизации.

3.2 Гб или 1.6 Гб с плавающей запятой — это начальная точка, с которой мы работаем. У нас есть около 500 Мб пространства для работы, если мы не хотим приближаться к тому месту, где нас может убить OOM-убийца Apple.

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

Оказывается, большинство из них относительно небольшие, в диапазоне менее 6 Мб каждый (2x320x64x64). Фреймворк, который я использую ( s4nnc) делает разумную работу по их упаковке в где-то менее 50 Мб общего объема с учетом повторного использования и т. д. Затем есть более интересный. Denoiser имеет механизм внутреннего внимания со своим собственным скрытым представлением изображения в качестве входных данных. Во время вычисления собственного внимания существует пакетная матрица размера (16x4096x4096). Это, если не очевидно, составляет около 500 Мб в половине плавающей запятой (FP16). Позже мы применили softmax к этой матрице. Это еще 500 Мб в FP16. Тщательная реализация softmax может быть выполнена «на месте», что означает, что он может безопасно переписать ввод без повреждения. К счастью, низкоуровневые библиотеки Apple и NVIDIA предоставили реализацию softmax на месте. К сожалению, библиотеки более высокого уровня, такие как PyTorch, этого не сделали.

Итак, это близко к пределу, но похоже, что мы можем сделать это где-то около 550 Мб + 1.6 Гб?

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

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

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

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

Это все еще на 2 Гб больше, чем мы можем себе позволить на iPhone. Что же делать? Перед этим еще немного предыстории: при вычислении собственного внимания с помощью CUDA распространенный трюк, реализованный в исходном коде Stable Diffusion, заключается в использовании перестановки, а не транспонирования (подробнее о том, что я имею в виду, см . в разделе «Трансформеры с нуля» ). Это помогает, потому что CUBLAS может работать с переставленными шаговыми тензорами напрямую, избегая одного выделенного трафика памяти для транспонирования тензора.

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

Осталось 1 Гб и больше предысторий! До MPSGraph были шейдеры Metal Performance. Это набор фиксированных примитивов Metal для некоторых операций нейронной сети. Вы можете думать о MPSGraph как о более блестящих, своевременно скомпилированных шейдерах, в то время как шейдеры Metal Performance — это более старая, но более надежная альтернатива.

Оказывается, 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, мы восстановили еще 500 МБ с потерей производительности примерно на 1%. Там мы, наконец, достигли размера 2 ГБ, к которому стремились ранее.

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

Вы можете скачать Draw Things сегодня по адресу https://draw.nnc.ai/

Вот несколько связанных ссылок по этой теме:

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