Optimizando la Difusión Estable para CPUs de Intel con NNCF y 🤗 Optimum

Optimizing Stable Diffusion for Intel CPUs with NNCF and 🤗 Optimum.

Los modelos de difusión latente son revolucionarios cuando se trata de resolver problemas de generación de texto a imagen. Diffusion Estable es uno de los ejemplos más famosos que ha sido ampliamente adoptado en la comunidad y la industria. La idea detrás del modelo de difusión estable es simple y convincente: se genera una imagen a partir de un vector de ruido en múltiples pasos pequeños refinando el ruido a una representación de imagen latente.

Sin embargo, este enfoque inevitablemente aumenta el tiempo total de inferencia y provoca una mala experiencia de usuario cuando se implementa en una máquina cliente. Uno puede notar que las potentes GPUs pueden ayudar aquí como de costumbre, y esto es cierto, pero el costo de esto aumenta drásticamente. Como referencia, en H1’23, el precio de una instancia potente de CPU r6i.2xlarge con 8 vCPUs y 64 GB de RAM es de $0.504 por hora, mientras que una instancia similar g4dn.2xlarge con NVIDIA T4 con 16 GB de VRAM cuesta $0.75 por hora, lo que es 1.5 veces más ..

Esto hace que los servicios de generación de imágenes sean bastante caros tanto para sus propietarios como para los usuarios. El problema es aún más agudo en las aplicaciones cliente que se ejecutan en el lado del usuario. ¡Aquí puede que no haya GPU en absoluto! Esto hace que la implementación del pipeline Stable Diffusion sea un problema desafiante.

A lo largo de los últimos cinco años, la caja de herramientas OpenVINO ha encapsulado muchas características para una inferencia de alto rendimiento. Inicialmente diseñado para modelos de visión por computadora, todavía domina en este dominio mostrando un rendimiento de inferencia de vanguardia para muchos modelos contemporáneos, incluido Stable Diffusion. Sin embargo, optimizar modelos de difusión estable para aplicaciones con restricciones de recursos requiere ir mucho más allá de las simples optimizaciones en tiempo de ejecución. Y aquí es donde entran en juego las capacidades de optimización de modelos del Marco de Compresión de Redes Neuronales (NNCF) de OpenVINO.

• Anunciando la Jam de Juegos de Inteligencia Artificial de Código Abierto 🎮
• Reconocimiento de Voz de IA en Unity
• DuckDB analiza más de 50,000 conjuntos de datos almacenados en el Hugging Face Hub

En esta publicación de blog, describiremos los problemas de optimización de los modelos de difusión estable y propondremos un flujo de trabajo que reduce sustancialmente la latencia de dichos modelos cuando se ejecutan en un hardware con recursos limitados, como una CPU. En particular, logramos una aceleración de inferencia de 5.1 veces y una reducción del tamaño del modelo de 4 veces en comparación con PyTorch.

Optimización de Stable Diffusion

En el pipeline de Stable Diffusion, el modelo UNet es computacionalmente el más costoso de ejecutar. Por lo tanto, optimizar solo un modelo aporta beneficios sustanciales en términos de velocidad de inferencia.

Sin embargo, resulta que los métodos tradicionales de optimización de modelos, como la cuantificación de 8 bits después del entrenamiento, no funcionan para este modelo. Hay dos razones principales para esto. Primero, los modelos de predicción a nivel de píxel, como la segmentación semántica, la superresolución, etc., son uno de los más complicados en términos de optimización de modelos debido a la complejidad de la tarea, por lo que ajustar los parámetros y la estructura del modelo rompe los resultados de varias formas. La segunda razón es que el modelo tiene un nivel inferior de redundancia porque aloja mucha información mientras se entrena con cientos de millones de muestras. Es por eso que los investigadores deben emplear métodos de cuantificación más sofisticados para preservar la precisión después de la optimización. Por ejemplo, Qualcomm utilizó el método Knowledge Distillation por capas (AdaRound) para cuantificar los modelos de difusión estables. Esto significa que de todos modos se requiere ajustar el modelo después de la cuantificación. Si es así, ¿por qué no simplemente utilizar el Entrenamiento con Conciencia de Cuantización (QAT) que puede ajustar el modelo y los parámetros de cuantización simultáneamente de la misma manera en que se entrena el modelo fuente? Por lo tanto, probamos este enfoque en nuestro trabajo utilizando NNCF, OpenVINO y Diffusers y lo combinamos con la Fusión de Tokens.

Flujo de trabajo de optimización

Normalmente comenzamos la optimización de un modelo después de que se haya entrenado. Aquí, comenzamos con un modelo ajustado a los datos del conjunto de datos de Pokemons, que contiene imágenes de Pokemons y sus descripciones de texto.

Utilizamos el ejemplo de ajuste fino de texto a imagen para Stable Diffusion de Diffusers e integramos QAT de NNCF en el siguiente script de entrenamiento. También cambiamos la función de pérdida para incorporar destilación de conocimiento del modelo fuente que actúa como profesor en este proceso, mientras que el modelo actual que se entrena actúa como estudiante. Este enfoque es diferente al método clásico de destilación de conocimiento, donde el modelo profesor entrenado se destila en un modelo estudiante más pequeño. En nuestro caso, la destilación de conocimiento se utiliza como un método auxiliar que ayuda a mejorar la precisión final del modelo optimizado. También utilizamos el método de Promedio Móvil Exponencial (EMA) para los parámetros del modelo, excluyendo los cuantizadores, lo que nos permite hacer que el proceso de entrenamiento sea más estable. Ajustamos el modelo durante solo 4096 iteraciones.

Con algunos trucos, como la checkpointing de gradientes y mantener el modelo EMA en RAM en lugar de VRAM, podemos ejecutar el proceso de optimización utilizando una sola GPU con 24 GB de VRAM. ¡Toda la optimización lleva menos de un día usando una GPU!

Trascendiendo el entrenamiento consciente de la cuantización

La cuantización por sí sola puede aportar mejoras significativas al reducir el tamaño del modelo, el tiempo de carga, el consumo de memoria y la latencia de inferencia. Pero lo maravilloso de la cuantización es que se puede aplicar junto con otros métodos de optimización, lo que lleva a una aceleración acumulativa.

Recientemente, Facebook Research introdujo un método de fusión de tokens para modelos de Vision Transformer. La esencia del método es que fusiona tokens redundantes con tokens importantes utilizando una de las estrategias disponibles (promedio, toma de valores máximos, etc.). Esto se hace antes del bloque de autoatención, que es la parte más computacionalmente exigente de los modelos Transformer. Por lo tanto, reducir la dimensión del token reduce el tiempo de cálculo general en los bloques de autoatención. Este método también se ha adaptado para modelos de Difusión Estable y ha mostrado resultados prometedores al optimizar pipelines de Difusión Estable para la síntesis de imágenes de alta resolución en ejecución en GPUs.

Modificamos el método de fusión de tokens para que sea compatible con OpenVINO y lo combinamos con la cuantización de 8 bits cuando se aplica al modelo Attention UNet. Esto también incluye todas las técnicas mencionadas, incluida la destilación de conocimiento, etc. En cuanto a la cuantización, se requiere un ajuste fino para restaurar la precisión. También comenzamos la optimización y el ajuste fino a partir del modelo entrenado en el conjunto de datos de Pokemons. La figura a continuación muestra un flujo de trabajo de optimización general.

El modelo resultante es altamente beneficioso cuando se ejecuta la inferencia en dispositivos con recursos computacionales limitados, como CPUs de cliente o de borde. Como se mencionó, combinar la fusión de tokens con la cuantización conduce a una reducción adicional en la latencia de inferencia.

Resultados de la demostración de generación de imágenes utilizando diferentes modelos optimizados. La entrada es “pájaro de dibujos animados”, el número aleatorio es 42. Los modelos utilizan OpenVINO 2022.3 en Hugging Face Spaces utilizando una instancia de “actualización de CPU” que utiliza procesadores escalables Intel® Xeon® de tercera generación con tecnología Intel® Deep Learning Boost.

Resultados

Utilizamos los flujos de trabajo de optimización divulgados para obtener dos tipos de modelos optimizados, cuantizados de 8 bits y cuantizados con fusión de tokens, y los comparamos con la línea de base de PyTorch. También convertimos la línea de base a un modelo de punto flotante (FP32) de OpenVINO para una comparación exhaustiva.

La imagen de arriba muestra los resultados de la generación de imágenes y algunas características del modelo. Como se puede ver, simplemente la conversión a OpenVINO produce una disminución significativa en la latencia de inferencia ( 1.9x ). La aplicación de la cuantización de 8 bits aumenta aún más la velocidad de inferencia, lo que lleva a una aceleración de 3.9x en comparación con PyTorch. Otro beneficio de la cuantización es una reducción significativa del tamaño del modelo, 0.25x del punto de control de PyTorch, lo que también mejora el tiempo de carga del modelo. La aplicación de la fusión de tokens (ToME) (con una tasa de fusión de 0.4 ) sobre la cuantización aporta una aceleración de rendimiento de 5.1x manteniendo el tamaño del modelo en el mismo nivel. No proporcionamos un análisis exhaustivo de la calidad visual de los modelos optimizados, pero, como se puede ver, los resultados son bastante sólidos.

A continuación, mostramos cómo realizar inferencias con el pipeline final optimizado para ejecutarse en CPUs de Intel:

from optimum.intel import OVStableDiffusionPipeline

# Cargar y compilar el pipeline para obtener un rendimiento óptimo.
nombre = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(nombre, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generar una imagen.
prompt = "un dibujo de un pokemon verde con ojos rojos"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("imagen.png")

Puede encontrar el código de entrenamiento y cuantización en la biblioteca Hugging Face Optimum Intel. El cuaderno que muestra la diferencia entre los modelos optimizados y los originales está disponible aquí . También puede encontrar muchos modelos en el Hugging Face Hub en la organización de OpenVINO . Además, hemos creado una demostración en Hugging Face Spaces que se ejecuta en una instancia r6id.2xlarge con procesador escalable Intel Xeon de tercera generación.

¿Qué pasa con el modelo de Difusión Estable de propósito general?

Como mostramos con la tarea de generación de imágenes de Pokemon, es posible lograr un alto nivel de optimización del pipeline de Difusión Estable utilizando una cantidad relativamente pequeña de recursos de entrenamiento. Al mismo tiempo, es bien sabido que entrenar un modelo de Difusión Estable de propósito general es una tarea costosa . Sin embargo, con suficiente presupuesto y recursos de hardware, es posible optimizar el modelo de propósito general utilizando el enfoque descrito y ajustarlo para producir imágenes de alta calidad. La única precaución que tenemos está relacionada con el método de fusión de tokens que reduce sustancialmente la capacidad del modelo. La regla general aquí es que cuanto más complicado sea el conjunto de datos que tenga para el entrenamiento, menor debería ser la tasa de fusión que debe utilizar durante la optimización.

Si disfrutaste leyendo esta publicación, es posible que también estés interesado en revisar esta publicación que discute otros enfoques complementarios para optimizar el rendimiento de Stable Diffusion en CPUs Intel Xeon de 4ta generación.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!