Acelerando la Inferencia de Difusión Estable en CPUs Intel

'Accelerating Stable Diffusion Inference on Intel CPUs'.

Recientemente, presentamos la última generación de CPUs Intel Xeon (nombre en clave Sapphire Rapids), sus nuevas características de hardware para aceleración de aprendizaje profundo y cómo utilizarlas para acelerar el ajuste fino distribuido y la inferencia para los Transformers de procesamiento de lenguaje natural.

En esta publicación, vamos a mostrar diferentes técnicas para acelerar los modelos de Difusión Estable en CPUs Sapphire Rapids. Una publicación de seguimiento hará lo mismo para el ajuste fino distribuido.

En el momento de escribir esto, la forma más sencilla de obtener un servidor Sapphire Rapids es utilizar la familia de instancias Amazon EC2 R7iz. Como aún está en versión preliminar, tienes que registrarte para obtener acceso. Al igual que en publicaciones anteriores, estoy utilizando una instancia r7iz.metal-16xl (64 vCPU, 512GB de RAM) con una AMI de Ubuntu 20.04 (ami-07cd3e6c4915b2d18).

¡Empecemos! Los ejemplos de código están disponibles en Gitlab.

• Inferencia rápida en modelos de lenguaje grandes BLOOMZ en el acelerador Habana Gaudi2
• Boletín de Ética y Sociedad #3 Apertura Ética en Hugging Face
• StackLLaMA Una guía práctica para entrenar LLaMA con RLHF

La biblioteca Diffusers

La biblioteca Diffusers hace extremadamente sencillo generar imágenes con modelos de Difusión Estable. Si no estás familiarizado con estos modelos, aquí tienes una gran introducción ilustrada.

Primero, creemos un entorno virtual con las bibliotecas requeridas: Transformers, Diffusers, Accelerate y PyTorch.

virtualenv sd_inference
source sd_inference/bin/activate
pip install pip --upgrade
pip install transformers diffusers accelerate torch==1.13.1

Luego, escribamos una función simple de benchmarking que ejecuta repetidamente la inferencia y devuelve la latencia promedio para la generación de una sola imagen.

import time

def elapsed_time(pipeline, prompt, nb_pass=10, num_inference_steps=20):
    # calentamiento
    images = pipeline(prompt, num_inference_steps=10).images
    start = time.time()
    for _ in range(nb_pass):
        _ = pipeline(prompt, num_inference_steps=num_inference_steps, output_type="np")
    end = time.time()
    return (end - start) / nb_pass

Ahora, construyamos un StableDiffusionPipeline con el tipo de datos float32 por defecto y midamos su latencia de inferencia.

from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "barco de vela en tormenta por Rembrandt"
latencia = elapsed_time(pipe, prompt)
print(latencia)

La latencia promedio es de 32,3 segundos. Como se demuestra en este Espacio de Intel, el mismo código se ejecuta en una generación anterior de Intel Xeon (nombre en clave Ice Lake) en aproximadamente 45 segundos.

¡De entrada, podemos ver que las CPUs Sapphire Rapids son bastante más rápidas sin ningún cambio de código!

Ahora, ¡aceleremos!

Óptimo Intel y OpenVINO

Óptimo Intel acelera las canalizaciones de extremo a extremo en arquitecturas Intel. Su API es extremadamente similar a la API de Diffusers básica, lo que facilita la adaptación de código existente.

Óptimo Intel admite OpenVINO, un kit de herramientas de código abierto de Intel para inferencia de alto rendimiento.

Óptimo Intel y OpenVINO se pueden instalar de la siguiente manera:

pip install optimum[openvino]

A partir del código anterior, solo necesitamos reemplazar StableDiffusionPipeline por OVStableDiffusionPipeline. Para cargar un modelo PyTorch y convertirlo al formato OpenVINO sobre la marcha, puedes establecer export=True al cargar tu modelo.

from optimum.intel.openvino import OVStableDiffusionPipeline
...
ov_pipe = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
latencia = elapsed_time(ov_pipe, prompt)
print(latencia)

# No olvides guardar el modelo exportado
ov_pipe.save_pretrained("./openvino")

OpenVINO optimiza automáticamente el modelo para el formato bfloat16. Gracias a esto, la latencia promedio ahora es de 16,7 segundos, una aceleración de 2x.

La canalización anterior es compatible con formas de entrada dinámicas, sin restricciones en el número de imágenes o su resolución. Con Difusión Estable, tu aplicación suele estar restringida a una (o unas pocas) resoluciones de salida diferentes, como 512×512 o 256×256. Por lo tanto, tiene mucho sentido desbloquear una aceleración significativa mediante la remodelación de la canalización a una resolución fija. Si necesitas más de una resolución de salida, simplemente puedes mantener unas pocas instancias de canalización, una para cada resolución.

ov_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
latency = elapsed_time(ov_pipe, prompt)

Con una forma estática, la latencia promedio se reduce a 4.7 segundos, lo que representa una aceleración adicional de 3.5 veces.

Como se puede observar, OpenVINO es una forma simple y eficiente de acelerar la inferencia de Difusión Estable. Cuando se combina con una CPU Sapphire Rapids, ofrece una aceleración de casi 10 veces en comparación con la inferencia convencional en Ice Lake Xeons.

Si no puedes o no quieres usar OpenVINO, el resto de esta publicación te mostrará una serie de otras técnicas de optimización. ¡Asegúrate de abrocharte el cinturón!

Optimización a nivel de sistema

Los modelos de Diffuser son modelos grandes de varios gigabytes y la generación de imágenes es una operación intensiva en memoria. Al instalar una biblioteca de asignación de memoria de alto rendimiento, deberíamos poder acelerar las operaciones de memoria y paralelizarlas en los núcleos de Xeon. Ten en cuenta que esto cambiará la biblioteca de asignación de memoria predeterminada en tu sistema. Por supuesto, puedes volver a la biblioteca predeterminada desinstalando la nueva.

jemalloc y tcmalloc son igualmente interesantes. Aquí, estoy instalando jemalloc ya que mis pruebas indican que tiene una ligera ventaja en rendimiento. También se puede ajustar para una carga de trabajo específica, por ejemplo, para maximizar la utilización de la CPU. Puedes consultar la guía de ajuste para obtener más detalles.

sudo apt-get install -y libjemalloc-dev
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libjemalloc.so
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms: 60000,muzzy_decay_ms:60000"

A continuación, instalamos la biblioteca libiomp para optimizar el procesamiento paralelo. Es parte de Intel OpenMP* Runtime.

sudo apt-get install intel-mkl
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libiomp5.so
export OMP_NUM_THREADS=32

Finalmente, instalamos la herramienta de línea de comandos numactl. Esto nos permite fijar nuestro proceso de Python en núcleos específicos y evitar parte de los gastos generados por los cambios de contexto.

numactl -C 0-31 python sd_blog_1.py

Gracias a estas optimizaciones, nuestro código original de Diffusers ahora predice en 11.8 segundos. Eso es casi 3 veces más rápido, sin ningún cambio de código. Estas herramientas funcionan muy bien en nuestro Xeon de 32 núcleos.

Aún no hemos terminado. Agreguemos la Extensión de Intel para PyTorch a la mezcla.

IPEX y BF16

La Extensión de Intel para PyTorch (IPEX) extiende PyTorch y aprovecha las características de aceleración de hardware presentes en las CPUs de Intel, como las Instrucciones de Redes Neuronales Vectoriales AVX-512 (AVX512 VNNI) y las Extensiones de Matriz Avanzadas (AMX).

Instalémoslo.

pip install intel_extension_for_pytorch==1.13.100

Luego actualizamos nuestro código para optimizar cada elemento del pipeline con IPEX (puedes listarlos imprimiendo el objeto pipe). Esto requiere convertirlos al formato de canales al final.

import torch
import intel_extension_for_pytorch as ipex
...
pipe = StableDiffusionPipeline.from_pretrained(model_id)

# al formato de canales al final
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.vae = pipe.vae.to(memory_format=torch.channels_last)
pipe.text_encoder = pipe.text_encoder.to(memory_format=torch.channels_last)
pipe.safety_checker = pipe.safety_checker.to(memory_format=torch.channels_last)

# Crear una entrada aleatoria para habilitar la compilación JIT
sample = torch.randn(2,4,64,64)
timestep = torch.rand(1)*999
encoder_hidden_status = torch.randn(2,77,768)
input_example = (sample, timestep, encoder_hidden_status)

# optimizar con IPEX
pipe.unet = ipex.optimize(pipe.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
pipe.vae = ipex.optimize(pipe.vae.eval(), dtype=torch.bfloat16, inplace=True)
pipe.text_encoder = ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
pipe.safety_checker = ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)

También habilitamos el formato de datos bloat16 para aprovechar la unidad de multiplicación de matrices de azulejos (TMMU) AMX presente en las CPUs Sapphire Rapids.

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    latencia = elapsed_time(pipe, prompt)
    print(latencia)

Con esta versión actualizada, la latencia de inferencia se reduce aún más de 11.9 segundos a 5.4 segundos. Eso es más de 2 veces la aceleración gracias a IPEX y AMX.

¿Podemos obtener un poco más de rendimiento? ¡Sí, con los planificadores!

Planificadores

La biblioteca Diffusers nos permite adjuntar un planificador a una tubería de difusión estable. Los planificadores intentan encontrar el mejor equilibrio entre la velocidad de eliminación de ruido y la calidad de eliminación de ruido.

Según la documentación: “En el momento de escribir este documento, DPMSolverMultistepScheduler ofrece posiblemente el mejor equilibrio entre velocidad y calidad, y se puede ejecutar con tan solo 20 pasos”.

Vamos a probarlo.

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
...
dpm = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=dpm)

Con esta versión final, la latencia de inferencia ahora es de 5.05 segundos. ¡Comparado con nuestro punto de referencia inicial de Sapphire Rapids (32.3 segundos), esto es casi 6.5 veces más rápido!

*Entorno: Amazon EC2 r7iz.metal-16xl, Ubuntu 20.04, Linux 5.15.0-1031-aws, libjemalloc-dev 5.2.1-1, intel-mkl 2020.0.166-1, PyTorch 1.13.1, Intel Extension for PyTorch 1.13.1, transformers 4.27.2, diffusers 0.14, accelerate 0.17.1, openvino 2023.0.0.dev20230217, optimum 1.7.1, optimum-intel 1.7*

Conclusión

La capacidad de generar imágenes de alta calidad en segundos debería funcionar bien para muchos casos de uso, como aplicaciones de clientes, generación de contenido para marketing y medios, o datos sintéticos para la ampliación de conjuntos de datos.

Aquí hay algunos recursos para ayudarte a empezar:

• Documentación de Diffusers
• Documentación de Optimum Intel
• Intel IPEX en GitHub
• Recursos para desarrolladores de Intel y Hugging Face.

Si tienes preguntas o comentarios, nos encantaría leerlos en el foro de Hugging Face.

¡Gracias por leer!

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