🧨 Accelerazione della diffusione stabile XL dell’inferenza con JAX su Cloud TPU v5e

🚀 Accelerazione stabile dell'inferenza XL con JAX su TPU Cloud v5e

I modelli AI generativi, come Stable Diffusion XL (SDXL), consentono la creazione di contenuti di alta qualità e realistici con applicazioni molto ampie. Tuttavia, sfruttare il potere di tali modelli presenta sfide significative e costi computazionali elevati. SDXL è un grande modello di generazione di immagini il cui componente UNet è circa tre volte più grande rispetto alla versione precedente del modello. Implementare un modello come questo in produzione è difficile a causa delle maggiori esigenze di memoria e dei tempi di inferenza aumentati. Oggi, siamo entusiasti di annunciare che Hugging Face Diffusers supporta ora il servizio SDXL utilizzando JAX su Cloud TPUs, consentendo inferenze ad alte prestazioni ed efficienti dal punto di vista dei costi.

Le Google Cloud TPUs sono acceleratori AI dedicati, ottimizzati per la formazione e l’inferenza di modelli AI di grandi dimensioni, compresi gli LLM all’avanguardia e i modelli AI generativi come SDXL. La nuova Cloud TPU v5e è progettata appositamente per offrire l’efficienza e le prestazioni necessarie per la formazione e l’inferenza su larga scala di intelligenza artificiale. A meno della metà del costo del TPU v4, il TPU v5e consente a un maggior numero di organizzazioni di formare e implementare modelli AI.

🧨 L’integrazione di Diffusers JAX offre un modo conveniente per eseguire SDXL su TPU tramite XLA, e abbiamo creato una demo per mostrarlo. Puoi provarlo su questo Spazio o nel playground incorporato di seguito:

Alla base, questa demo viene eseguita su diverse istanze TPU v5e-4 (ognuna con 4 chip TPU) e sfrutta la parallelizzazione per servire quattro grandi immagini 1024×1024 in circa 4 secondi. Questo tempo include conversioni di formato, tempo di comunicazione e elaborazione front-end; il tempo effettivo di generazione è di circa 2,3 secondi, come vedremo di seguito!

• Implementare la AI Comic Factory utilizzando l’API di inferenza
• Accelerare oltre 130.000 modelli di Hugging Face con ONNX Runtime
• Probabilmente conosci le GAN, ma sai come addestrarle?

In questo post di blog,

1. Descriviamo perché JAX + TPU + Diffusers è un framework potente per eseguire SDXL
2. Spieghiamo come è possibile scrivere una semplice pipeline di generazione di immagini con Diffusers e JAX
3. Mostriamo i benchmark che confrontano diverse impostazioni TPU

Perché JAX + TPU v5e per SDXL?

È possibile eseguire SDXL con JAX su Cloud TPU v5e con elevate prestazioni e un costo efficiente grazie alla combinazione di hardware TPU progettato appositamente e una stack di software ottimizzata per le prestazioni. Di seguito evidenziamo due fattori chiave: la compilazione just-in-time (jit) di JAX e il parallelismo guidato dal compilatore XLA con pmap di JAX.

Compilazione JIT

Una caratteristica notevole di JAX è la sua compilazione just-in-time (jit). Il compilatore JIT traccia il codice durante il primo avvio e genera binari TPU altamente ottimizzati che vengono riutilizzati nelle chiamate successive. La particolarità di questo processo è che richiede tutte le forme di input, intermedi e di output che siano statiche, ossia devono essere conosciute in anticipo. Ogni volta che cambiamo le forme, verrà attivato nuovamente un processo di compilazione costoso. La compilazione JIT è ideale per i servizi che possono essere progettati intorno a forme statiche: la compilazione viene eseguita una volta, e poi sfruttiamo tempi di inferenza super veloci.

La generazione di immagini si presta bene alla compilazione JIT. Se generiamo sempre lo stesso numero di immagini e sono di dimensioni uguali, allora le forme di output sono costanti e conosciute in anticipo. Gli input di testo sono anche costanti: per progettazione, Stable Diffusion e SDXL utilizzano vettori di embedding di forma fissa (con imbottitura) per rappresentare le indicazioni inserite dall’utente. Pertanto, possiamo scrivere codice JAX che si basa su forme fisse, e che può essere notevolmente ottimizzato!

Prestazioni elevate per alte dimensioni di lotti

Le carichi di lavoro possono essere scalati su dispositivi multipli utilizzando il pmap di JAX, che esprime programmi singoli dati multipli (SPMD). L’applicazione di pmap ad una funzione compilerà una funzione con XLA, quindi l’eseguirà in parallelo su diversi dispositivi XLA. Per i carichi di lavoro di generazione di testo in immagine, ciò significa che aumentare il numero di immagini renderizzate contemporaneamente è semplice da implementare e non compromette le prestazioni. Ad esempio, l’esecuzione di SDXL su un TPU con 8 chip genererà 8 immagini nello stesso tempo impiegato da 1 chip per creare un’immagine singola.

Le istanze TPU v5e sono disponibili in diverse forme, tra cui forme a 1, 4 e 8 chip, fino a 256 chip (un pod completo TPU v5e), con collegamenti ICI ultraveloci tra i chip. Ciò consente di scegliere la forma TPU che meglio si adatta al proprio caso d’uso e di sfruttare facilmente il parallelismo offerto da JAX e TPUs.

Come scrivere una pipeline di generazione di immagini in JAX

Andremo passo dopo passo sul codice che devi scrivere per eseguire l’infrazione super veloce utilizzando JAX! Prima di tutto, importiamo le dipendenze.

  # Mostra le migliori pratiche per SDXL JAX importa jax importa jax.numpy come jnp importa numpy come np from flax.jax_utils import replicate from diffusers import FlaxStableDiffusionXLPipeline import time  

Ora caricheremo il modello base SDXL e il resto dei componenti necessari per l’infrazione. La pipeline di diffusers si occupa di scaricare e memorizzare nella cache tutto per noi. Rispettando l’approccio funzionale di JAX, i parametri del modello vengono restituiti separatamente e dovranno essere passati alla pipeline durante l’infrazione:

  pipeline, params = FlaxStableDiffusionXLPipeline.from_pretrained ("stabilityai / stable-diffusion-xl-base-1.0", split_head_dim = True)  

I parametri del modello vengono scaricati con una precisione di 32 bit di default. Per risparmiare memoria ed eseguire il calcolo più velocemente, li convertiremo in bfloat16 , una rappresentazione efficiente a 16 bit. Tuttavia, c’è un’avvertenza: per ottenere i migliori risultati, dobbiamo mantenere lo stato dello scheduler in float32 , altrimenti gli errori di precisione si accumulano e si traducono in immagini di bassa qualità o addirittura nere.

  scheduler_state = params.pop ("scheduler") params = jax.tree_util.tree_map (lambda x: x.astype (jnp.bfloat16), params) params ["scheduler"] = scheduler_state  

Ora siamo pronti per impostare il nostro prompt e il resto degli input alla pipeline.

  default_prompt = "foto di alta qualità di un delfino bebè che gioca in piscina e indossa un cappello da festa" default_neg_prompt = "illustrazione, bassa qualità" default_seed = 33 default_guidance_scale = 5.0 default_num_steps = 25  

I prompt devono essere forniti come tensori alla pipeline, e devono sempre avere le stesse dimensioni in ogni invocazione. Ciò consente di compilare la chiamata di infrazione. Il metodo prepare_inputs della pipeline esegue tutti i passaggi necessari per noi, quindi creeremo una funzione di aiuto per preparare sia il prompt che il prompt negativo come tensori. Lo utilizzeremo in seguito dalla nostra funzione generate :

  def tokenize_prompt (prompt, neg_prompt): prompt_ids = pipeline.prepare_inputs (prompt) neg_prompt_ids = pipeline.prepare_inputs (neg_prompt) return prompt_ids, neg_prompt_ids  

Per sfruttare il parallelismo, replicheremo gli input su più dispositivi. Un Cloud TPU v5e-4 ha 4 chip, quindi replicando gli input otteniamo ogni chip per generare un’immagine diversa, in parallelo. Dobbiamo fare attenzione a fornire un seed casuale diverso a ogni chip in modo che le 4 immagini siano diverse:

  NUM_DEVICES = jax.device_count () # I parametri del modello non cambiano durante l'infrazione, # quindi dobbiamo replicarli solo una volta. p_params = replicate (params) def replicate_all (prompt_ids, neg_prompt_ids, seed): p_prompt_ids = replicate (prompt_ids) p_neg_prompt_ids = replicate (neg_prompt_ids) rng = jax.random.PRNGKey (seed) rng = jax.random.split (rng, NUM_DEVICES) return p_prompt_ids, p_neg_prompt_ids, rng  

Ora siamo pronti per mettere tutto insieme in una funzione di generazione:

def generate(prompt, negative_prompt, seed=default_seed, guidance_scale=default_guidance_scale, num_inference_steps=default_num_steps,):
    prompt_ids, neg_prompt_ids = tokenize_prompt(prompt, negative_prompt)
    prompt_ids, neg_prompt_ids, rng = replicate_all(prompt_ids, neg_prompt_ids, seed)
    images = pipeline(
        prompt_ids,
        p_params,
        rng,
        num_inference_steps=num_inference_steps,
        neg_prompt_ids=neg_prompt_ids,
        guidance_scale=guidance_scale,
        jit=True,
    ).images
    # convert the images to PIL
    images = images.reshape((images.shape[0] * images.shape[1], ) + images.shape[-3:])
    return pipeline.numpy_to_pil(np.array(images))

jit=True indica che vogliamo che la chiamata alla pipeline sia compilata. Questo accadrà la prima volta che chiamiamo generate, e sarà molto lenta: JAX deve tracciare le operazioni, ottimizzarle e convertirle in primitive di basso livello. Eseguiamo una prima generazione per completare questo processo e preparare il terreno:

start = time.time()
print(f"Compilazione in corso...")
generate(default_prompt, default_neg_prompt)
print(f"Compilato in {time.time() - start}")

La compilazione ha impiegato circa tre minuti la prima volta che l’abbiamo eseguita. Ma una volta compilato il codice, l’inferenza sarà estremamente veloce. Proviamo di nuovo!

start = time.time()
prompt = "lama nell'antica Grecia, olio su tela"
neg_prompt = "cartone animato, illustrazione, animazione"
images = generate(prompt, neg_prompt)
print(f"Inferenza in {time.time() - start}")

Ora ci sono voluti circa 2 secondi per generare le 4 immagini!

Benchmark

Le seguenti misurazioni sono state ottenute eseguendo SDXL 1.0 base per 20 passaggi, con il programmatore discreto di Eulero. Confrontiamo Cloud TPU v5e con TPUv4 per le stesse dimensioni di batch. Si noti che, a causa della parallelizzazione, un TPU v5e-4 come quelli che utilizziamo nella nostra demo genererà 4 immagini quando si utilizza una dimensione di batch di 1 (o 8 immagini con una dimensione di batch di 2). Allo stesso modo, un TPU v5e-8 genererà 8 immagini quando si utilizza una dimensione di batch di 1.

Le prove su Cloud TPU sono state eseguite utilizzando Python 3.10 e la versione 0.4.16 di Jax. Queste sono le stesse specifiche utilizzate nella nostra demo Spazio.

TPU v5e raggiunge una perf/$ fino a 2,4 volte superiore su SDXL rispetto a TPU v4, dimostrando l’efficienza in termini di costi della generazione TPU più recente.

Per misurare le prestazioni di inferenza, utilizziamo la metrica standard del settore della throughput. Prima misuriamo la latenza per immagine quando il modello è stato compilato e caricato. Quindi calcoliamo il throughput dividendo la dimensione del batch per la latenza per chip. Di conseguenza, il throughput misura le prestazioni del modello in ambienti di produzione indipendentemente da quanti chip vengono utilizzati. Quindi dividiamo il throughput per il prezzo di listino per ottenere le prestazioni per dollaro.

Come funziona la demo?

La demo che abbiamo mostrato prima è stata creata utilizzando uno script che segue essenzialmente il codice che abbiamo pubblicato in questo post del blog. Viene eseguita su alcuni dispositivi Cloud TPU v5e con 4 chip ciascuno, e c’è un semplice server di bilanciamento del carico che indirizza le richieste degli utenti ai server backend in modo casuale. Quando inserisci un prompt nella demo, la tua richiesta verrà assegnata a uno dei server backend e riceverai le 4 immagini che genera.

Si tratta di una soluzione semplice basata su diverse istanze TPU preallocati. In un post futuro, affronteremo la creazione di soluzioni dinamiche che si adattano al carico utilizzando GKE.

Tutto il codice per la demo è open source e disponibile in Hugging Face Diffusers oggi. Siamo entusiasti di vedere cosa costruirete con Diffusers + JAX + Cloud TPUs!