Diffusione più veloce e stabile con Core ML su iPhone, iPad e Mac

Diffusione più veloce e stabile con Core ML su dispositivi Apple.

La WWDC’23 (Apple Worldwide Developers Conference) si è tenuta la scorsa settimana. Molte delle novità si sono concentrate sull’annuncio di Vision Pro durante il keynote, ma c’è molto di più. Come ogni anno, la settimana della WWDC è ricca di oltre 200 sessioni tecniche che approfondiscono le nuove funzionalità in arrivo sui sistemi operativi e i framework Apple. Quest’anno siamo particolarmente entusiasti dei cambiamenti in Core ML dedicati alle tecniche di compressione e ottimizzazione. Questi cambiamenti rendono l’esecuzione di modelli come Stable Diffusion più veloce e con un uso di memoria inferiore! A titolo di esempio, considera il seguente test eseguito sul mio iPhone 13 lo scorso dicembre, confrontato con la velocità attuale utilizzando la palettizzazione a 6 bit:

Stable Diffusion su iPhone, dicembre scorso e ora con palettizzazione a 6 bit

Contenuti

• Nuove ottimizzazioni di Core ML
• Utilizzo di modelli di Stable Diffusion quantizzati e ottimizzati
• Conversione e ottimizzazione di modelli personalizzati
• Utilizzo di meno di 6 bit
• Conclusioni

Nuove ottimizzazioni di Core ML

Core ML è un framework maturo che consente l’esecuzione efficiente di modelli di machine learning sui dispositivi, sfruttando tutta l’hardware di calcolo nei dispositivi Apple: CPU, GPU e Neural Engine specializzato in compiti di ML. L’esecuzione on-device sta attraversando un periodo di straordinario interesse, stimolato dalla popolarità di modelli come Stable Diffusion e Large Language Models con interfacce di chat. Molte persone vogliono eseguire questi modelli sul proprio hardware per una serie di motivi, tra cui comodità, privacy e risparmi sui costi delle API. Naturalmente, molti sviluppatori stanno esplorando modi per eseguire efficientemente questi modelli on-device e creando nuove app e casi d’uso. Le migliorie di Core ML che contribuiscono a raggiungere questo obiettivo sono una grande notizia per la comunità!

Le modifiche di ottimizzazione di Core ML comprendono due pacchetti software diversi (ma complementari):

• Sì, i Transformers sono efficaci per la previsione delle serie temporali (+ Autoformer)
• Politica di AI @🤗 Risposta alla richiesta di commenti sulla responsabilità dell’AI dell’NTIA degli Stati Uniti.
• Ottimizza i modelli di adattatore MMS per ASR a bassa risorsa

• Il framework Core ML stesso. Questo è il motore che esegue i modelli di ML sull’hardware Apple ed è parte del sistema operativo. I modelli devono essere esportati in un formato speciale supportato dal framework, e questo formato è anche chiamato “Core ML”.
• Il pacchetto di conversione coremltools. Questo è un modulo Python open-source il cui scopo è convertire modelli PyTorch o Tensorflow nel formato Core ML.

coremltools include ora un nuovo sottomodulo chiamato coremltools.optimize con tutti gli strumenti di compressione e ottimizzazione. Per tutti i dettagli su questo pacchetto, dai un’occhiata a questa sessione WWDC. Nel caso di Stable Diffusion, utilizzeremo la palettizzazione a 6 bit, un tipo di quantizzazione che comprime i pesi del modello da una rappresentazione in virgola mobile a 16 bit a soli 6 bit per parametro. Il nome “palettizzazione” si riferisce a una tecnica simile a quella utilizzata nei grafici computerizzati per lavorare con un insieme limitato di colori: la tavolozza dei colori (o “palette”) contiene un numero fisso di colori, e i colori nell’immagine vengono sostituiti dagli indici dei colori più vicini disponibili nella tavolozza. Ciò fornisce immediatamente il vantaggio di ridurre drasticamente la dimensione di archiviazione, riducendo così il tempo di download e l’utilizzo del disco on-device.

Illustrazione della palettizzazione a 2 bit. Credito immagine: Apple WWDC’23 Session Use Core ML Tools per la compressione dei modelli di machine learning.

I pesi compressi a 6 bit non possono essere utilizzati per i calcoli, perché sono solo indici in una tabella e non rappresentano più l’entità dei pesi originali. Pertanto, Core ML deve decomprimere i pesi palettizzati prima dell’uso. Nelle versioni precedenti di Core ML, la decompressione avveniva quando il modello veniva caricato per la prima volta dal disco, quindi la quantità di memoria utilizzata era uguale alla dimensione del modello non compresso. Con le nuove migliorie, i pesi vengono mantenuti come numeri a 6 bit e convertiti al volo durante l’inferenza da uno strato all’altro. Questo potrebbe sembrare lento, poiché l’esecuzione di un’inferenza richiede molte operazioni di decompressione, ma è tipicamente più efficiente rispetto alla preparazione di tutti i pesi in modalità a 16 bit! La ragione è che i trasferimenti di memoria sono nel percorso critico dell’esecuzione e trasferire meno memoria è più veloce che trasferire dati non compressi.

Utilizzo di modelli di Stable Diffusion quantizzati e ottimizzati

Lo scorso dicembre, Apple ha introdotto ml-stable-diffusion, un repository open-source basato su diffusori per convertire facilmente modelli di Stable Diffusion in Core ML. Applica anche ottimizzazioni agli strati di attenzione dei trasformatori che rendono l’inferenza più veloce sul Neural Engine (sui dispositivi su cui è disponibile). ml-stable-diffusion è stato appena aggiornato dopo la WWDC con quanto segue:

• La quantizzazione è supportata utilizzando --quantize-nbits durante la conversione. Puoi quantizzare a 8, 6, 4 o addirittura 2 bit! Per ottenere i migliori risultati, consigliamo di utilizzare la quantizzazione a 6 bit, poiché la perdita di precisione è ridotta mentre si ottiene un’elaborazione rapida e notevoli risparmi di memoria. Se desideri andare oltre, consulta questa sezione per tecniche avanzate.
• Ulteriori ottimizzazioni dei livelli di attenzione che offrono prestazioni ancora migliori sull’Engine Neurale! Il trucco consiste nel suddividere le sequenze di query in blocchi da 512 per evitare la creazione di grandi tensori intermedi. Questo metodo è chiamato SPLIT_EINSUM_V2 nel codice e può migliorare le prestazioni tra il 10% e il 30%.

Per rendere facile per tutti trarre vantaggio da queste migliorie, abbiamo convertito i quattro modelli Stable Diffusion ufficiali e li abbiamo caricati su Hub . Queste sono tutte le varianti:

Nota che ogni variante è disponibile nel formato Core ML e anche come archivio zip. I file zip sono ideali per le app native, come la nostra app demo open-source e altri strumenti di terze parti . Se desideri solo eseguire i modelli sul tuo hardware, il modo più semplice è utilizzare la nostra app demo e selezionare il modello quantizzato che desideri testare. Devi compilare l’app utilizzando Xcode, ma un aggiornamento sarà disponibile per il download sull’App Store a breve. Per maggiori dettagli, consulta il nostro post precedente .

Esecuzione del modello stable-diffusion-2-1-base a 6 bit nell’app demo

Se desideri scaricare un particolare pacchetto Core ML per integrarlo nel tuo progetto Xcode, puoi clonare i repository o semplicemente scaricare la versione di interesse utilizzando il codice seguente.

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-2-1-base-palettized"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Modello scaricato in {model_path}")

Conversione e ottimizzazione di modelli personalizzati

Se desideri utilizzare un modello Stable Diffusion personalizzato (ad esempio, se hai affinato o creato i tuoi modelli), puoi utilizzare il repository ml-stable-diffusion di Apple per effettuare la conversione tu stesso. Di seguito è riportato un breve riassunto di come procedere, ma ti consigliamo di leggere i dettagli della documentazione .

1. Seleziona il modello che desideri convertire. Puoi addestrarne uno tuo o sceglierne uno dalla Galleria dei modelli Diffusers di Hugging Face . Ad esempio, convertiamo prompthero/openjourney-v4 .
2. Installa apple/ml-stable-diffusion ed esegui una prima conversione utilizzando l’implementazione di attenzione ORIGINAL come segue:

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-vae-encoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation ORIGINAL \
    --compute-unit CPU_AND_GPU \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/original/openjourney-6-bit

1. Ripeti la conversione per l’implementazione di attenzione SPLIT_EINSUM_V2:

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation SPLIT_EINSUM_V2 \
    --compute-unit ALL \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/split_einsum_v2/openjourney-6-bit

1. Testa i modelli convertiti sull’hardware desiderato. Come regola generale, la versione ORIGINAL di solito funziona meglio su macOS, mentre SPLIT_EINSUM_V2 è di solito più veloce su iOS. Per ulteriori dettagli e punti di dati aggiuntivi, consulta questi test contribuiti dalla community sulla versione precedente di Stable Diffusion per Core ML.

2. Per integrare il modello desiderato nella tua app:

   • Se stai distribuendo il modello all’interno dell’app, utilizza i file .mlpackage. Tieni presente che questo aumenterà la dimensione del tuo pacchetto dell’app.
   • In caso contrario, puoi utilizzare le Risorse compilati per scaricarle in modo dinamico quando l’app viene avviata.

Utilizzo di meno di 6 bit

La quantizzazione a 6 bit è un punto di equilibrio tra qualità del modello, dimensione del modello e comodità: devi solo fornire un’opzione di conversione per poter quantizzare qualsiasi modello pre-addestrato. Questo è un esempio di compressione post-addestramento.

La versione beta di coremltools rilasciata la scorsa settimana include anche metodi di compressione durante l’addestramento. L’idea qui è che puoi affinare un modello Stable Diffusion pre-addestrato e eseguire la compressione dei pesi durante l’operazione di affinamento. Ciò ti consente di utilizzare una compressione a 4 o addirittura a 2 bit riducendo al minimo la perdita di qualità. La ragione per cui ciò funziona è perché il clustering dei pesi viene eseguito utilizzando un algoritmo differenziabile e quindi possiamo applicare gli usuali ottimizzatori di addestramento per trovare la tabella di quantizzazione riducendo al minimo la perdita del modello.

Abbiamo in programma di valutare presto questo metodo e non vediamo l’ora di vedere come funzionano i modelli ottimizzati a 4 bit e quanto velocemente eseguono. Se ci batti sul tempo, per favore inviaci una nota e saremo felici di controllare 🙂

Conclusione

I metodi di quantizzazione possono essere utilizzati per ridurre la dimensione dei modelli Stable Diffusion, farli eseguire più velocemente sul dispositivo e consumare meno risorse. Le ultime versioni di Core ML e coremltools supportano tecniche come la palettizzazione a 6 bit che sono facili da applicare e hanno un impatto minimo sulla qualità. Abbiamo aggiunto modelli palettizzati a 6 bit all’Hub, che sono sufficientemente piccoli da funzionare su iOS e macOS. Abbiamo anche mostrato come è possibile convertire i modelli affinati autonomamente e non vediamo l’ora di vedere cosa farai con questi strumenti e tecniche!