Accelerazione dell’Inferenza di Diffusione Stabile su CPU Intel

Accelerazione dell'Inferenza di Diffusione Stabile su CPU Intel

Recentemente, abbiamo introdotto l’ultima generazione di processori Intel Xeon (nome in codice Sapphire Rapids), le sue nuove caratteristiche hardware per l’accelerazione dell’apprendimento approfondito e come utilizzarle per accelerare il raffinamento distribuito e l’elaborazione per l’elaborazione del linguaggio naturale con i Transformer.

In questo post, ti mostreremo diverse tecniche per accelerare i modelli Stable Diffusion su processori Sapphire Rapids. Un post successivo farà lo stesso per il raffinamento distribuito.

Al momento della scrittura, il modo più semplice per mettere le mani su un server Sapphire Rapids è utilizzare la famiglia di istanze Amazon EC2 R7iz. Poiché è ancora in anteprima, devi iscriverti per ottenere l’accesso. Come nei post precedenti, sto utilizzando un’istanza r7iz.metal-16xl (64 vCPU, 512 GB di RAM) con un’AMI Ubuntu 20.04 (ami-07cd3e6c4915b2d18 ).

Iniziamo! I campioni di codice sono disponibili su Gitlab .

La libreria Diffusers

La libreria Diffusers rende estremamente semplice generare immagini con modelli di Diffusione Stabile. Se non sei familiare con questi modelli, ecco una grande introduzione illustrata .

Per prima cosa, creiamo un ambiente virtuale con le librerie richieste: Transformers, Diffusers, Accelerate e PyTorch.

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

Quindi, scriviamo una semplice funzione di benchmarking che esegue ripetutamente l’elaborazione e restituisce la latenza media per la generazione di un’immagine singola.

import time

def elapsed_time(pipeline, prompt, nb_pass=10, num_inference_steps=20):
    # riscaldamento
    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

Ora, creiamo una StableDiffusionPipeline con il tipo di dati predefinito float32 e misuriamo la latenza dell’elaborazione.

from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "nave a vela in tempesta di Rembrandt"
latenza = elapsed_time(pipe, prompt)
print(latenza)

La latenza media è di 32,3 secondi . Come dimostrato da questo Intel Space , lo stesso codice viene eseguito su una generazione precedente di processori Intel Xeon (nome in codice Ice Lake) in circa 45 secondi.

Di default, possiamo vedere che i processori Sapphire Rapids sono abbastanza più veloci senza alcuna modifica del codice!

Ora, acceleriamo!

Optimum Intel e OpenVINO

Optimum Intel accelera le pipeline end-to-end su architetture Intel. La sua API è estremamente simile all’API Diffusers di base, rendendo facile adattare il codice esistente.

Optimum Intel supporta OpenVINO , un toolkit open-source Intel per l’elaborazione ad alte prestazioni.

Optimum Intel e OpenVINO possono essere installati come segue:

pip install optimum[openvino]

Partendo dal codice sopra, dobbiamo solo sostituire StableDiffusionPipeline con OVStableDiffusionPipeline . Per caricare un modello PyTorch e convertirlo nel formato OpenVINO al volo, è possibile impostare export=True durante il caricamento del modello.

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

# Non dimenticare di salvare il modello esportato
ov_pipe.save_pretrained("./openvino")

OpenVINO ottimizza automaticamente il modello per il formato bfloat16. Grazie a questo, la latenza media è ora di 16,7 secondi , un bel vantaggio del 2x.

La pipeline sopra supporta forme di input dinamiche, senza restrizioni sul numero di immagini o sulla loro risoluzione. Con Stable Diffusion, di solito la tua applicazione è limitata a una (o poche) diverse risoluzioni di output, come ad esempio 512×512 o 256×256. Pertanto, ha molto senso sbloccare un’accelerazione significativa ridimensionando la pipeline a una risoluzione fissa. Se hai bisogno di più di una risoluzione di output, puoi semplicemente mantenere diverse istanze della pipeline, una per ogni risoluzione.

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

Con una forma statica, la latenza media si riduce a 4.7 secondi, un ulteriore aumento di velocità del 3.5x.

Come puoi vedere, OpenVINO è un modo semplice ed efficiente per accelerare l’infertza Stable Diffusion. Quando combinato con una CPU Sapphire Rapids, offre un aumento di velocità di quasi 10x rispetto all’infertza standard su Ice Lake Xeons.

Se non puoi o non vuoi utilizzare OpenVINO, il resto di questo post ti mostrerà una serie di altre tecniche di ottimizzazione. Allaccia le cinture!

Ottimizzazione a livello di sistema

I modelli Diffuser sono modelli di grandi dimensioni di multi-gigabyte e la generazione di immagini è un’operazione intensiva in termini di memoria. Installando una libreria di allocazione della memoria ad alte prestazioni, dovremmo essere in grado di velocizzare le operazioni di memoria e parallelizzarle sui core Xeon. Si prega di notare che ciò cambierà la libreria di allocazione della memoria predefinita nel tuo sistema. Naturalmente, puoi tornare alla libreria predefinita disinstallando quella nuova.

jemalloc e tcmalloc sono entrambi interessanti. Qui, sto installando jemalloc poiché i miei test gli conferiscono un leggero vantaggio in termini di prestazioni. Può anche essere regolato per un carico di lavoro particolare, ad esempio per massimizzare l’utilizzo della CPU. Puoi fare riferimento alla guida di ottimizzazione per ulteriori dettagli.

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"

Successivamente, installiamo la libreria libiomp per ottimizzare l’elaborazione parallela. Fa parte di 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

Infine, installiamo lo strumento da riga di comando numactl. Ciò ci consente di fissare il nostro processo Python a core specifici ed evitare parte del sovraccarico legato allo switch di contesto.

numactl -C 0-31 python sd_blog_1.py

Grazie a queste ottimizzazioni, il nostro codice Diffusers originale ora predice in 11.8 secondi. È quasi 3 volte più veloce, senza alcuna modifica del codice. Questi strumenti funzionano sicuramente molto bene sul nostro Xeon a 32 core.

Non siamo ancora finiti. Aggiungiamo l’Estensione Intel per PyTorch alla miscela.

IPEX e BF16

L’Estensione Intel per Pytorch (IPEX) estende PyTorch e sfrutta le funzionalità di accelerazione hardware presenti sui processori Intel, come le Istruzioni di Rete Neurale Vettoriale AVX-512 (AVX512 VNNI) ed Estensioni Matriciali Avanzate (AMX).

Installiamolo.

pip install intel_extension_for_pytorch==1.13.100

Aggiorniamo quindi il nostro codice per ottimizzare ogni elemento del pipeline con IPEX (puoi elencarli stampando l’oggetto pipe). Ciò richiede di convertirli nel formato channels-last.

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

# to channels last
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)

# Create random input to enable JIT compilation
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)

# optimize with 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)

Abilitiamo anche il formato dati bloat16 per sfruttare l’acceleratore AMX tile matrix multiply unit (TMMU) presente sui processori Sapphire Rapids.

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

Con questa versione aggiornata, la latenza dell’inferenza si riduce ulteriormente da 11,9 secondi a 5,4 secondi. Questo significa un’accelerazione superiore al 2x grazie a IPEX e AMX.

Possiamo ottenere ancora migliori prestazioni? Sì, con i scheduler!

Scheduler

La libreria Diffusers ci consente di allegare uno scheduler a una pipeline Stable Diffusion. Gli scheduler cercano il miglior compromesso tra velocità di denoising e qualità del denoising.

Secondo la documentazione: “Al momento della stesura di questo documento, DPMSolverMultistepScheduler offre probabilmente il miglior compromesso tra velocità e qualità e può essere eseguito con soli 20 passaggi.”

Proviamolo.

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

Con questa versione finale, la latenza dell’inferenza è ora di 5,05 secondi. Rispetto al nostro punto di riferimento iniziale Sapphire Rapids (32,3 secondi), questo è quasi 6,5 volte più veloce!

*Ambiente: 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*

Conclusione

La capacità di generare immagini di alta qualità in pochi secondi può essere molto utile per molti casi d’uso, come le app per i clienti, la generazione di contenuti per il marketing e i media, o la generazione di dati sintetici per l’aumento del dataset.

Ecco alcune risorse per aiutarti a iniziare:

  • Documentazione di Diffusers
  • Documentazione di Optimum Intel
  • Intel IPEX su GitHub
  • Risorse per sviluppatori da Intel e Hugging Face.

Se hai domande o feedback, saremmo felici di leggerli sul forum di Hugging Face.

Grazie per la lettura!