Ottimizzazione della diffusione stabile per i processori Intel con NNCF e 🤗 Optimum

Ottimizzazione diffusione stabile processori Intel NNCF 🤗 Optimum.

I modelli di diffusione latenti sono dei cambiamenti di gioco quando si tratta di risolvere problemi di generazione di testo-immagine. La diffusione stabile è uno degli esempi più famosi che ha ottenuto un’ampia adozione nella comunità e nell’industria. L’idea alla base del modello di diffusione stabile è semplice e convincente: generi un’immagine da un vettore di rumore in più piccoli passaggi raffinando il rumore in una rappresentazione di immagine latente.

Tuttavia, tale approccio inevitabilmente aumenta il tempo di inferenza complessivo e causa una scarsa esperienza utente quando viene implementato su una macchina client. Si può notare che le potenti GPU possono aiutare come al solito, ed è vero, ma il costo di ciò aumenta drasticamente. A titolo di riferimento, nel H1’23, il prezzo di un’istanza di CPU potente r6i.2xlarge con 8 vCPU e 64 GB di RAM è di $0,504 all’ora, mentre una simile istanza g4dn.2xlarge con NVIDIA T4 con 16 GB di VRAM costa $0,75 all’ora, che è 1,5 volte di più ..

Ciò rende i servizi di generazione di immagini piuttosto costosi per i proprietari e gli utenti. Il problema è ancora più acuto nelle applicazioni client che vengono eseguite sul lato dell’utente. Potrebbe non esserci nessuna GPU affatto! Ciò rende la distribuzione del pipeline di diffusione stabile un problema impegnativo.

Negli ultimi cinque anni, OpenVINO Toolkit ha incorporato molte funzionalità per l’elaborazione ad alte prestazioni. Inizialmente progettato per modelli di visione artificiale, domina ancora in questo dominio mostrando prestazioni di inferenza di livello superiore per molti modelli contemporanei, inclusa la diffusione stabile. Tuttavia, l’ottimizzazione dei modelli di diffusione stabile per le applicazioni con risorse limitate richiede di andare molto oltre le semplici ottimizzazioni in fase di esecuzione. Ed è qui che entrano in gioco le capacità di ottimizzazione del modello dal framework di compressione delle reti neurali OpenVINO (NNCF).

• Nuovi modelli ViT e ALIGN da Kakao Brain
• Ottimizzazione avanzata di 20B LLM con RLHF su una GPU da 24GB per consumatori
• Jupyter X Hugging Face

In questo post del blog, descriveremo i problemi dell’ottimizzazione dei modelli di diffusione stabile e proporremo un flusso di lavoro che riduce sostanzialmente la latenza di tali modelli quando vengono eseguiti su una HW con risorse limitate come la CPU. In particolare, abbiamo ottenuto un’accelerazione dell’inferenza di 5,1 volte e una riduzione delle dimensioni del modello di 4 volte rispetto a PyTorch.

Ottimizzazione della diffusione stabile

Nel pipeline di diffusione stabile, il modello UNet è computazionalmente il più costoso da eseguire. Pertanto, ottimizzare un solo modello porta notevoli benefici in termini di velocità di inferenza.

Tuttavia, si scopre che i metodi tradizionali di ottimizzazione del modello, come la quantizzazione a 8 bit dopo l’allenamento, non funzionano per questo modello. Ci sono due ragioni principali per questo. In primo luogo, i modelli di previsione a livello di pixel, come la segmentazione semantica, la super risoluzione, ecc., sono tra i più complessi in termini di ottimizzazione del modello a causa della complessità del compito, quindi la modifica dei parametri del modello e della struttura compromette i risultati in numerosi modi. La seconda ragione è che il modello ha un livello inferiore di ridondanza perché ospita molte informazioni pur essendo addestrato su centinaia di milioni di campioni. Ecco perché i ricercatori devono utilizzare metodi di quantizzazione più sofisticati per preservare l’accuratezza dopo l’ottimizzazione. Ad esempio, Qualcomm ha utilizzato il metodo di distillazione della conoscenza a livello di layer (AdaRound) per quantizzare i modelli di diffusione stabile. Ciò significa che è comunque necessario sintonizzare il modello dopo la quantizzazione. In tal caso, perché non utilizzare semplicemente l’addestramento consapevole della quantizzazione (QAT) che può sintonizzare il modello e i parametri di quantizzazione contemporaneamente nello stesso modo in cui il modello di origine viene addestrato? Pertanto, abbiamo provato questo approccio nel nostro lavoro utilizzando NNCF, OpenVINO e Diffusers e l’abbiamo associato a Token Merging.

Flusso di lavoro di ottimizzazione

Solitamente iniziamo l’ottimizzazione di un modello dopo che è stato addestrato. Qui, partiamo da un modello ottimizzato a livello di testo-immagine sul dataset dei Pokémon che contiene immagini dei Pokémon e le loro descrizioni testuali.

Abbiamo utilizzato l’esempio di sintonizzazione fine del testo-immagine per la diffusione stabile da Diffusers e abbiamo integrato QAT da NNCF nello script di addestramento seguente. Abbiamo anche modificato la funzione di perdita per incorporare la distillazione della conoscenza dal modello di origine che funge da insegnante in questo processo mentre il modello effettivo in fase di addestramento funge da studente. Questo approccio è diverso dal metodo classico di distillazione della conoscenza, in cui il modello insegnante addestrato viene distillato in un modello studente più piccolo. Nel nostro caso, la distillazione della conoscenza viene utilizzata come metodo ausiliario che aiuta a migliorare l’accuratezza finale del modello di ottimizzazione. Utilizziamo anche il metodo di media mobile esponenziale (EMA) per i parametri del modello escludendo i quantizzatori che ci consente di rendere il processo di addestramento più stabile. Abbiamo sintonizzato il modello per sole 4096 iterazioni.

Con alcuni trucchi, come il checkpointing del gradiente e il mantenimento del modello EMA in RAM anziché nella VRAM, possiamo eseguire il processo di ottimizzazione utilizzando una sola GPU con 24 GB di VRAM. L’intera ottimizzazione richiede meno di un giorno utilizzando una sola GPU!

Andare oltre l’addestramento consapevole della quantizzazione

La quantizzazione da sola può portare miglioramenti significativi riducendo la dimensione del modello, il tempo di caricamento, il consumo di memoria e la latenza dell’inferenza. Ma la grande cosa della quantizzazione è che può essere applicata insieme ad altri metodi di ottimizzazione portando a un aumento cumulativo della velocità.

Recentemente, Facebook Research ha introdotto un metodo di unione dei token per i modelli Vision Transformer. La sostanza del metodo è che unisce i token ridondanti a quelli importanti utilizzando una delle strategie disponibili (media, valore massimo, ecc.). Questo viene fatto prima del blocco di auto-attenzione, che è la parte più computazionalmente impegnativa dei modelli Transformer. Pertanto, riducendo la dimensione del token si riduce il tempo complessivo di calcolo nei blocchi di auto-attenzione. Questo metodo è stato anche adattato per i modelli di diffusione stabile e ha mostrato risultati promettenti nell’ottimizzazione delle pipeline di diffusione stabile per la sintesi di immagini ad alta risoluzione eseguite su GPU.

Abbiamo modificato il metodo di unione dei token per renderlo conforme a OpenVINO e lo abbiamo impilato con la quantizzazione a 8 bit quando applicato al modello Attention UNet. Ciò coinvolge anche tutte le tecniche menzionate, compresa la distillazione della conoscenza, ecc. Per quanto riguarda la quantizzazione, richiede un’intonazione fine per essere applicata per ripristinare l’accuratezza. Iniziamo anche l’ottimizzazione e l’intonazione fine dal modello addestrato sul dataset di Pokemons. La figura qui sotto mostra un flusso di lavoro di ottimizzazione complessivo.

Il modello risultante è molto vantaggioso quando si esegue l’inferenza su dispositivi con risorse di calcolo limitate, come CPU client o edge. Come è stato menzionato, l’impilamento di Token Merging con la quantizzazione porta a una riduzione aggiuntiva della latenza dell’inferenza.

Risultati della demo di generazione di immagini utilizzando diversi modelli ottimizzati. L’input prompt è “cartoon bird”, il seed è 42. I modelli sono con OpenVINO 2022.3 in Hugging Face Spaces utilizzando un’istanza “CPU upgrade” che utilizza processori Intel® Xeon® Scalable di terza generazione con tecnologia Intel® Deep Learning Boost.

Risultati

Abbiamo utilizzato i flussi di lavoro di ottimizzazione divulgati per ottenere due tipi di modelli ottimizzati, quantizzati a 8 bit e quantizzati con Token Merging, e li abbiamo confrontati con il baselines di PyTorch. Abbiamo anche convertito il baselines in un modello OpenVINO floating-point (FP32) per il confronto esaustivo.

La figura sopra mostra i risultati della generazione di immagini e alcune caratteristiche del modello. Come si può vedere, la semplice conversione in OpenVINO porta a una significativa diminuzione della latenza dell’inferenza ( 1,9x ). L’applicazione della quantizzazione a 8 bit aumenta ulteriormente la velocità di inferenza portando a un’accelerazione del 3,9x rispetto a PyTorch. Un altro vantaggio della quantizzazione è una significativa riduzione della dimensione del modello, 0,25x del checkpoint di PyTorch, che migliora anche il tempo di caricamento del modello. L’applicazione di Token Merging (ToME) (con un rapporto di unione del 0,4 ) in cima alla quantizzazione porta a un’accelerazione delle prestazioni del 5,1x mantenendo la dimensione del modello allo stesso livello. Non abbiamo fornito un’analisi approfondita della qualità visiva dei modelli ottimizzati, ma, come si può vedere, i risultati sono abbastanza solidi.

Sotto mostriamo come eseguire l’inferenza con la pipeline finale ottimizzata per l’esecuzione su CPU Intel:

from optimum.intel import OVStableDiffusionPipeline

# Carica e compila la pipeline per le prestazioni.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Genera un'immagine.
prompt = "un disegno di un pokemon verde con occhi rossi"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("image.png")

Puoi trovare il codice di addestramento e quantizzazione nella libreria Hugging Face Optimum Intel. Il notebook che mostra la differenza tra modelli ottimizzati e originali è disponibile qui . Puoi trovare anche molti modelli sull’Hugging Face Hub nell’organizzazione OpenVINO . Inoltre, abbiamo creato una demo su Hugging Face Spaces che viene eseguita su un’istanza r6id.2xlarge con processore Intel Xeon Scalable di terza generazione.

Cosa dire del modello di diffusione stabile ad uso generale?

Come abbiamo mostrato con il compito di generazione di immagini di Pokemon, è possibile raggiungere un alto livello di ottimizzazione della pipeline di diffusione stabile utilizzando una quantità relativamente piccola di risorse di addestramento. Allo stesso tempo, è ben noto che addestrare un modello di diffusione stabile ad uso generale è un compito costoso . Tuttavia, con un budget sufficiente e risorse HW, è possibile ottimizzare il modello ad uso generale utilizzando l’approccio descritto e regolarlo per produrre immagini di alta qualità. L’unico avvertimento che abbiamo riguarda il metodo di unione dei token che riduce notevolmente la capacità del modello. La regola generale qui è che più complicato è il dataset che hai per l’addestramento, meno rapporto di unione dovresti utilizzare durante l’ottimizzazione.

Se hai apprezzato la lettura di questo post, potresti essere interessato/a a dare un’occhiata a questo post che discute altri approcci complementari per ottimizzare le prestazioni di Stable Diffusion sui processori Intel Xeon di 4a generazione.