Запускаем Stable Diffusion на Raspberry PI Zero 2 (или на 260 МБ ОЗУ)

Задача — запустить Stable Diffusion ^[1], включающую большую трансформирующую модель c почти 1 миллиардом параметров, на Raspberry Pi Zero 2 ^[2] с 512 МБ RAM, не добавляя дополнительного пространства подкачки и не выгружая промежуточные результаты на диск. Рекомендуемый минимальный объём RAM/VRAM для Stable Diffusion составляет 8 ГБ.

Как правило, ведущие фреймворки машинного обучения и библиотеки фокусируются на минимизации задержки и повышении пропускной способности, всё за счёт потребления RAM. Поэтому я решил написать миниатюрную конфигурируемую библиотеку вывода специально для минимизации потребления памяти: OnnxStream.

OnnxStream основана на идее отделения механизма вывода от компонента, отвечающего за предоставление весов модели, а именно класса, происходящего от WeightsProvider.

WeightsProvider может реализовывать любой вид загрузки, кэширования и предварительного получения параметров модели. Например, кастомные WeightsProvider могут решить загрузить свои данные непосредственно с HTTP-сервера, не загружая и не записывая что-либо на диск (отсюда и слово Stream в OnnxStream). По умолчанию доступны два WeightsProvider:
DiskNoCache и DiskPrefetch.

OnnxStream может потреблять даже в 55 раз меньше памяти, чем OnnxRuntime, работая всего в 0,5-2 раза медленнее (на CPU, читайте раздел «Производительность» ниже).

▍ Stable Diffusion

Эти изображения были сгенерированы примером реализации Stable Diffusion, включённым в этот репозиторий, с использованием OnnxStream при разной точности декодера VAE. Декодер VAE — это единственная модель Stable Diffusion, которая не вместилась в память Raspberry Pi Zero 2 ни с одинарной, ни с половинной точностью. Причиной стало присутствие в модели остаточных соединений, а также очень больших тензоров и свёрток. Единственным решением оказалось статическое квантование (8 бит).

Третье изображение я сгенерировал на RPI Zero 2 примерно за 3 часа. Первое для сравнения было получено на моём ПК с использованием тех же латентов (latents), сгенерированных RPI Zero 2:

Декодер VAE с точностью W16A16

Декодер VAE с точностью W8A32

Декодер VAE с точностью W8A8 (сгенерировано RPI Zero 2 где-то за 3 часа)

▍ Особенности OnnxStream

• Механизм вывода отделён от WeightsProvider.
• WeightsProvider может быть DiskNoCache, DiskPrefetch или кастомным.
• Разделение внимания.
• Динамическое квантование (8 бит без знака, асимметричное, по перцентилю).
• Статическое квантование (W8A8 без знака, асимметричное, по перцентилю).
• Простая калибровка квантованной модели.
• Поддержка FP16 (с арифметикой FP16 или без неё).
• Реализовано 24 оператора ONNX (самые распространённые).
• Операции выполняются последовательно, но все операторы являются многопоточными.
• Один файл реализации + заголовочный файл.
• Вызовы XNNPACK обёрнуты в класс XnnPack (для дальнейшей замены).

OnnxStream зависит от XNNPACK для некоторых (ускоренных) примитивов: MatMul, Convolution, поэлементные Add/Sub/Mul/Div, Sigmoid и Softmax.

▍ Производительность

Stable Diffusion состоит из трёх моделей: кодировщика текста (672 операции и 123 миллиона параметров), UNET (2050 операций и 854 миллиона параметров) и декодера VAE (276 операций и 49 миллионов параметров). Предполагая, что размер пакета равен 1, полная генерация изображения выполняется в 10 шагов, что даёт хорошие результаты (с использованием планировщика Euler Ancestral), требует 2 выполнений шифровщика текста, 20 (то есть 2*10) выполнений модели UNET и 1 выполнения декодера VAE.

Эта таблица показывает различное время вывода трёх моделей Stable Diffusion вместе с объёмом потребляемой памяти (то есть Peak Working Set Size в Windows или Maximum Resident Set Size в Linux).

В случае модели UNET (при выполнении с точностью FP16 с арифметикой FP16) OnnxStream может потреблять даже в 55 раз меньше памяти, чем OnnxRuntime, выполняясь всего в 0,5-2 раза медленнее.

Примечания:

• Первое выполнение OnnxRuntime — это прогревочный вывод, поскольку её InferenceSession создаётся до первого выполнения и повторно используется для всех последующих. Для OnnxStream такого понятия как «разогрев» не существует, поскольку этот инструмент по своей природе всегда наготове (тем не менее последующие выполнения могут получать преимущество благодаря кэшированию файлов весов в ОС).
• Пока что OnnxStream не поддерживает ввод с размером пакета != 1 при том, что OnnxRuntime при выполнении модели UNET может значительно ускорять весь процесс диффузии, используя размер пакета 2.
• В моих тестах изменение SessionOptions в OnnxRuntime (например, EnableCpuMemArena и ExecutionMode) не ведёт к значительным изменениям результата.
• Производительность OnnxRuntime очень близка к производительности NCNN (ещё один проанализированный мной фреймворк), как в плане потребления памяти, так и в плане времени вывода.
• Тесты я выполнял на своей машине для разработки: Windows Server 2019, 16GB RAM, 8750H CPU (AVX2), 970 EVO Plus SSD, 8 виртуальных ядер на VMWare.

▍ Разделение внимания и квантование

Использование «разделения внимания» (attention slicing) при выполнении модели UNET и использование квантования W8A8 для декодера VAE было необходимо для сокращения потребления памяти до уровня, который бы позволил выполнить программу на RPI Zero 2.

Несмотря на то, что в интернете есть много информации по теме квантования нейронных сетей, сложно найти хоть что-то про «разделение внимания». Идея этого механизма проста: задача избежать материализации всей матрицы Q @ K^T при вычислении скалярного произведения масштабированного внимания различных многопоточных вниманий в модели UNET.

При 8 лучах внимания в этой модели Q имеет форму (8,4096,40), в то время как K^T имеет форму (8,40,4096): поэтому результат первой MatMul имеет финальную форму (8,4096,4096), то есть является тензором размером 512 МБ (с точностью FP32):

Решением будет разделить Q вертикально и выполнить операции внимания стандартно для каждой полученной части. Q_sliced имеет форму (1,x,40), где x равен 4096 (в этом случае), поделённым на onnxstream::Model::m_attention_fused_ops_parts (с предустановленным значением 2, которое можно изменить). Этот простой приём позволяет уменьшить общий объём потребляемой моделью UNET памяти с 1,1 ГБ до 300 МБ (когда модель выполняется с точностью FP32). Возможной и явно более эффективной альтернативой будет использование FlashAttention. Однако FlashAttention потребует написания кастомного ядра для каждой поддерживаемой архитектуры (AVX, NEON и так далее), в нашем случае обходя XnnPack.

▍ Как работает OnnxStream

Этот код может выполнять модель, определённую в path_to_model_folder/model.txt: (все операции модели определены в файле model.txt. OnnxStream ожидает найти все веса в том же каталоге в виде серии файлов .bin)

#include "onnxstream.h"

using namespace onnxstream;

int main()
{
    Model model;

    //
    // опциональные параметры, которые можно установить для объекта Model:
    //
    // model.set_weights_provider( ... ); // устанавливает другого поставщика весов (по умолчанию это DiskPrefetchWeightsProvider)
    // model.read_range_data( ... ); // считывает файл диапазонов (который содержит диапазоны отрезания активаций квантуемой модели)
    // model.write_range_data( ... ); // записывает файл диапазонов (пригодится после калибровки)
    // model.m_range_data_calibrate = true; // калибрует модель
    // model.m_use_fp16_arithmetic = true; // использует при выводе арифметику FP16 (пригождается, если веса находятся в точности FP16)
    // model.m_use_uint8_arithmetic = true; // использует при выводе арифметику UINT8 
    // model.m_use_uint8_qdq = true; // использует динамическое квантование UINT8 (может сокращать потребление памяти некоторыми моделями)
    // model.m_fuse_ops_in_attention = true; // активирует разделение внимания
    // model.m_attention_fused_ops_parts = ... ; // читайте «Разделение внимания» выше
  
    model.read_file("path_to_model_folder/model.txt");

    tensor_vector<float> data;
    
    ... // заполняет tensor_vector данными. «tensor_vector» - это просто псевдоним для std::vector с кастомным аллокатором.

    Tensor t;
    t.m_name = "input";
    t.m_shape = { 1, 4, 64, 64 };
    t.set_vector(std::move(data));
    model.push_tensor(std::move(t));

    model.run();
    
    auto& result = model.m_data[0].get_vector<float>();
    
    ... // обработка результата: «result» - это ссылка на первый результат вывода (а также tensor_vector<float>).

    return 0;
}

Файл model.txt содержит все операции с моделями в формате ASCII в том виде, в каком они были экспортированы из исходного файла ONNX. Каждая строка соответствует отдельной операции: например, эта представляет свёртку в квантованной модели:

Conv_4:Conv*input:input_2E_1(1,4,64,64);post_5F_quant_5F_conv_2E_weight_nchw.bin(uint8[0.0035054587850383684,134]:4,4,1,1);post_5F_quant_5F_conv_2E_bias.bin(float32:4)*output:input(1,4,64,64)*dilations:1,1;group:1;kernel_shape:1,1;pads:0,0,0,0;strides:1,1

Для экспорта model.txt и его весов (в виде серии файлов .bin) из файла ONNX для использования в OnnxStream предоставляется блокнот (с одной ячейкой) (onnx2txt.ipynb).

При экспорте Pytorch nn.Module (в нашем случае) в ONNX для использования в OnnxStream нужно кое-что учесть:

1. Во время вызова torch.onnx.export, dynamic_axes нужно оставить пустым, поскольку OnnxStream не поддерживает ввод с динамической формой.
2. Настоятельно рекомендуется выполнять ONNX Simplifier ^[3] для экспортированного файла ONNX до его преобразования в файл model.txt.

▍ Как собрать пример Stable Diffusion в Linux/Mac/Windows

• Только Windows: запустите следующую командную строку: Visual Studio Tools > x64 Native Tools Command Prompt.
• Только Mac: установите cmake: brew install cmake.

Сначала нужно собрать XNNPACK ^[4].

Поскольку прототипы функций XnnPack могут измениться в любое время, я включил git checkout, чтобы обеспечить корректную компиляцию OnnxStream с совместимой на момент написания статьи версией XnnPack:

git clone https://github.com/google/XNNPACK.git
cd XNNPACK
git rev-list -n 1 --before="2023-06-27 00:00" master
git checkout <COMMIT_ID_FROM_THE_PREVIOUS_COMMAND>
mkdir build
cd build
cmake -DXNNPACK_BUILD_TESTS=OFF -DXNNPACK_BUILD_BENCHMARKS=OFF ..
cmake --build . --config Release

Затем можно собрать пример Stable Diffusion:

<КАТАЛОГ_КУДА_БЫЛ_КЛОНИРОВАН_XNNPACK>, например /home/vito/Desktop/XNNPACK или C:ProjectsSDXNNPACK (в Windows):

git clone https://github.com/vitoplantamura/OnnxStream.git
cd OnnxStream
cd src
mkdir build
cd build
cmake -DXNNPACK_DIR=<DIRECTORY_WHERE_XNNPACK_WAS_CLONED> ..
cmake --build . --config Release

Теперь можете выполнить полученный пример. Веса для него доступны в разделе Releases репозитория проекта ^[5]. Опции командной строки для этого примера Stable Diffusion:

--models-path       устанавливает каталог, содержащий модели Stable Diffusion.
--ops-printf        во время вывода записывает текущую операцию в stdout.
--output            устанавливает выходной файл PNG.
--decode-latents    пропускает диффузию и декодирует указанный файл латентов.
--prompt            устанавливает положительный (желаемый) запрос.
--neg-prompt        устанавливает отрицательный (нежелательный) запрос.
--steps             устанавливает количество шагов диффузии.
--save-latents      после диффузии сохраняет латенты в указанном файле.
--decoder-calibrate калибрует квантованную версию деокдера VAE.
--decoder-fp16      во время вывода использует версию FP16 декодера VAE.
--rpi               конфигурирует модели для выполнения на Raspberry Pi Zero 2.

▍ Примечание

Реализация Stable Diffusion в sd.cpp основана на этом проекте ^[6], который, в свою очередь, основан на этом проекте ^[7] @EdVince. Изначальный код был изменён для использования OnnxStream вместо NCNN.