Оптимизация Stable Diffusion для CPU: 8-битная квантизация и Token Merging

Введение: проблема инференса Stable Diffusion на CPU

Stable Diffusion генерирует изображения высокого качества, но обычно требует GPU с большим объёмом видеопамяти. Серверы без GPU, edge-устройства и сценарии с жёстким бюджетом на инфраструктуру часто остаются за бортом. Инференс на CPU возможен, но базовые версии моделей работают медленно и занимают много памяти.

Практический пайплайн, который разбирается в этой статье, даёт ускорение в 5.1 раза и уменьшение размера модели в 4 раза по сравнению с базовым PyTorch FP32. Достигается это комбинацией трёх методов: 8-битная квантизация с Quantization-Aware Training и Knowledge Distillation, Token Merging для сокращения вычислений в self-attention блоках UNet и оптимизации OpenVINO. Примеры кода основаны на Optimum Intel, готовые модели размещены на Hugging Face Hub.

Материал ориентирован на разработчиков, которые хотят запускать диффузионные модели без GPU и при этом сохранять приемлемое качество генерации.

Почему традиционная квантизация не работает для Stable Diffusion

Пост-обучающая квантизация, или PTQ, применяется к уже обученной модели без дополнительного обучения. Для Stable Diffusion такой подход приводит к заметной деградации изображений: появляются артефакты, теряются детали, нарушается цветопередача. Причина в высокой чувствительности диффузионных моделей к ошибкам округления весов и активаций.

UNet, ключевой компонент Stable Diffusion, содержит множество слоёв с разным распределением значений. При переводе в INT8 часть информации теряется, и ошибки накапливаются по мере прохождения через десятки шагов денойзинга. Результат: модель работает быстрее, но генерирует визуально неприемлемые изображения.

Quantization-Aware Training (QAT) с Knowledge Distillation

QAT имитирует квантизацию во время обучения. Модель видит, как будут округляться веса и активации, и адаптируется к этому. Вместо резкого перехода в INT8 после обучения модель постепенно учится работать с ограниченной точностью.

Knowledge Distillation усиливает этот процесс. Полноточная модель выступает учителем, квантизованная - учеником. Ученик обучается не только на исходных данных, но и на выходных распределениях учителя. Это помогает восстановить качество, потерянное при сжатии.

Для реализации используется NNCF, фреймворк квантизации от OpenVINO. Он интегрируется с PyTorch и поддерживает QAT из коробки. Дистилляция настраивается через функцию потерь, которая сравнивает логиты учителя и ученика.

Token Merging (ToMe): сокращение вычислений в self-attention

Token Merging уменьшает количество токенов в self-attention блоках UNet. Похожие токены объединяются в один, что снижает вычислительную сложность квадратично зависящую от числа токенов. Для изображений высокого разрешения это даёт существенный выигрыш.

Метод работает без переобучения модели. Он анализирует сходство токенов и объединяет те, которые несут избыточную информацию. В UNet токены соответствуют пространственным областям изображения. Соседние области с похожими признаками объединяются, сокращая общее число операций в attention.

Вклад ToMe в общее ускорение зависит от настроек агрессивности. При умеренных параметрах качество изображений остаётся визуально близким к оригиналу, а скорость инференса растёт на десятки процентов. Для CPU это критично, так как self-attention на процессоре выполняется медленнее, чем на GPU.

Практическая реализация: OpenVINO, NNCF и Optimum Intel

Optimum Intel связывает Hugging Face Transformers и Diffusers с инструментами оптимизации OpenVINO. Через него можно применять квантизацию, экспортировать модели в формат OpenVINO IR и запускать инференс на CPU.

Пример кода: 8-битная квантизация с QAT

from optimum.intel import OVQuantizer
from optimum.intel import OVConfig

# Загрузка модели Stable Diffusion
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")

# Настройка квантизации с QAT
quantization_config = OVConfig(
    quantization=True,
    quantization_approach="QAT",
    bits=8
)

quantizer = OVQuantizer.from_pretrained(pipeline)
quantizer.quantize(
    quantization_config=quantization_config,
    calibration_dataset=calibration_dataset,
    save_directory="./sd_int8_qat"
)

Ключевые параметры: quantization_approach="QAT" включает обучение с имитацией квантизации, bits=8 задаёт целевую разрядность. Для дистилляции используется учитель - полноточная модель, которая передаёт знания через функцию потерь.

Пример кода: применение Token Merging

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float32
)

# Применение ToMe к UNet
pipeline.unet = apply_tome(pipeline.unet, ratio=0.5)

Функция apply_tome модифицирует self-attention блоки UNet, добавляя слой объединения токенов. Параметр ratio определяет долю объединяемых токенов. Значение 0.5 означает сокращение числа токенов вдвое.

Экспорт и инференс в OpenVINO

from optimum.intel import OVStableDiffusionPipeline

# Загрузка оптимизированной модели
ov_pipeline = OVStableDiffusionPipeline.from_pretrained(
    "./sd_int8_qat",
    export=True,
    compile=True
)

# Инференс на CPU
image = ov_pipeline("a cat sitting on a table", num_inference_steps=20).images[0]
image.save("output.png")

Экспорт в формат OpenVINO IR выполняется через export=True. Параметр compile=True запускает компиляцию модели под конкретное железо. После этого инференс выполняется через OpenVINO Runtime, который использует все доступные ядра CPU.

Готовые модели с уже применённой квантизацией и ToMe доступны на Hugging Face Hub. Это позволяет пропустить этап обучения и сразу перейти к тестированию.

Результаты: ускорение в 5.1 раза и уменьшение размера в 4 раза

Сравнение проводилось на серверном CPU Intel Sapphire Rapids. Базовая модель - Stable Diffusion v1.5 в PyTorch FP32. Оптимизированная версия - INT8 с QAT и дистилляцией, ToMe с ratio 0.5, экспорт в OpenVINO IR.

Ускорение складывается из трёх составляющих. Квантизация снижает объём вычислений и памяти, ToMe сокращает число операций в self-attention, OpenVINO оптимизирует исполнение под конкретный CPU с использованием AVX-512 и AMX.

Для контекста: отдельные методы дают меньший прирост. Квантизация без ToMe ускоряет инференс примерно в 2 раза, ToMe без квантизации - на 30-40%. Комбинация даёт синергетический эффект.

Похожие подходы к ускорению CPU-инференса разбираются в статье про оптимизацию Stable Diffusion на Intel Sapphire Rapids. Там показано, как системные настройки jemalloc и numactl добавляют ещё почти 3-кратный прирост без изменения кода.

Ограничения и возможные проблемы

Результаты варьируются в зависимости от оборудования. На процессорах без поддержки AMX или AVX-512 ускорение будет меньше. Конкретная модель также влияет: fine-tuned версии могут быть более чувствительны к квантизации, чем базовая.

QAT требует дополнительного обучения и данных. Это не бесплатный шаг: нужен калибровочный датасет, время на обучение и вычислительные ресурсы. Для разовой задачи затраты могут не окупиться. Для серийного использования оптимизированной модели - окупаются.

ToMe при агрессивных настройках снижает качество. Слишком высокий ratio объединения токенов приводит к потере мелких деталей и размытию текстур. Подбирать параметр нужно под конкретную задачу и тип генерируемых изображений.

Дистилляция не гарантирует полное восстановление качества. В некоторых случаях артефакты квантизации остаются заметными, особенно на сложных сценах с мелкими объектами. Тестирование на своих данных обязательно.

Заключение

8-битная квантизация с QAT и Knowledge Distillation, Token Merging и оптимизации OpenVINO дают практический путь к запуску Stable Diffusion на CPU. Итоговый пайплайн ускоряет инференс в 5.1 раза и уменьшает размер модели в 4 раза без значимых потерь качества.

Для разработчиков, которые уже используют CPU-инференс, этот подход открывает возможность генерации изображений на существующей инфраструктуре. Для тех, кто только оценивает варианты, цифры показывают: CPU-инференс диффузионных моделей стал реальной альтернативой GPU в сценариях с ограниченным бюджетом.

Готовые модели на Hugging Face Hub позволяют быстро проверить результат на своём железе. Если вы работаете с трансформерами на CPU, обратите внимание на практический тест ускорения PyTorch-трансформеров на Intel Sapphire Rapids. Для визуальных языковых моделей есть отдельный гайд по запуску VLM на Intel CPU.

Подписаться на канал