Miglioramenti di WebAssembly e WebGPU per un'IA web più rapida, parte 2

Questo documento è una continuazione dei miglioramenti di WebAssembly e WebGPU per Web AI più veloce, parte 1. Prima di continuare, ti consigliamo di leggere questo post o di guardare la presentazione all'indirizzo IO 24.

Austin Eng
Austin Eng
Deepti Gandluri
Deepti Gandluri
François Beaufort
François Beaufort

WebGPU

WebGPU consente alle applicazioni web di accedere all'hardware GPU del client per eseguire un calcolo efficiente e altamente parallelo. Dal lancio di WebGPU in Chrome, abbiamo visto incredibili demo di intelligenza artificiale (IA) e machine learning (ML) sul web.

Ad esempio, Web Stable Diffusion ha dimostrato che è stato possibile utilizzare l'IA per generare immagini dal testo direttamente nel browser. All'inizio di quest'anno, il team Mediapipe di Google ha pubblicato il supporto sperimentale per l'inferenza dei modelli linguistici di grandi dimensioni (LLM).

La seguente animazione mostra Gemma, il modello linguistico di grandi dimensioni (LLM) open source di Google, in esecuzione interamente sul dispositivo in Chrome, in tempo reale.

La seguente demo di Hugging Face di Segment Anything Model di Meta produce maschere degli oggetti di alta qualità interamente sul client.

Questi sono solo un paio degli incredibili progetti che mostrano la potenza di WebGPU per AI e ML. WebGPU consente a questi e ad altri modelli di essere eseguiti molto più velocemente di quanto potrebbero fare con la CPU.

Il benchmark WebGPU per l'incorporamento del testo di Hugging Face ha registrato incredibili accelerazioni rispetto a un'implementazione CPU dello stesso modello. Su un laptop Apple M1 Max, WebGPU era oltre 30 volte più veloce. Altri hanno segnalato che WebGPU accelera il benchmark di oltre 120 volte.

Miglioramento delle funzionalità di WebGPU per AI e ML

WebGPU è un'ottima soluzione per i modelli di IA e ML, che possono avere miliardi di parametri, grazie al supporto degli shadowr computing. I Compute staging vengono eseguiti sulla GPU e consentono di eseguire operazioni parallele di array su grandi volumi di dati.

Tra i numerosi miglioramenti apportati a WebGPU nell'ultimo anno, abbiamo continuato ad aggiungere altre funzionalità per migliorare le prestazioni di ML e IA sul web. Di recente, abbiamo lanciato due nuove funzionalità: la rappresentazione in virgola mobile a 16 bit e i prodotti a punti interi compressi.

rappresentazione in virgola mobile a 16 bit

Ricorda che i carichi di lavoro ML non richiedono precisione. shader-f16 è una funzionalità che consente l'utilizzo del tipo f16 nel linguaggio di ombreggiatura di WebGPU. Questo tipo in virgola mobile occupa 16 bit, invece dei soliti 32 bit. f16 ha un intervallo più piccolo ed è meno preciso, ma per molti modelli ML questo è sufficiente.

Questa funzionalità aumenta l'efficienza in alcuni modi:

  • Memoria ridotta: i tensori con elementi f16 occupano metà dello spazio, dimezzando l'utilizzo della memoria. I calcoli della GPU sono spesso colli di bottiglia sulla larghezza di banda della memoria, quindi metà della memoria può spesso comportare un'esecuzione doppiata degli screener. Tecnicamente, non è necessario f16 per risparmiare larghezza di banda della memoria. È possibile archiviare i dati in un formato a bassa precisione e quindi espanderli a f32 completo nello shaker per il calcolo. Tuttavia, la GPU consuma più potenza di calcolo per pacchettizzare e decomprimere i dati.

  • Conversione dei dati ridotta: f16 utilizza meno risorse di calcolo riducendo al minimo la conversione dei dati. I dati a bassa precisione possono essere archiviati e quindi utilizzati direttamente senza conversione.

  • Aumento del parallelismo: le GPU moderne sono in grado di includere più valori contemporaneamente nelle unità di esecuzione della GPU, consentendo di eseguire un numero maggiore di calcoli paralleli. Ad esempio, una GPU che supporta fino a 5 bilioni di operazioni in virgola mobile con f32 al secondo potrebbe supportare 10 trilioni di operazioni in virgola mobile con f16 al secondo.

Screenshot del benchmark WebGPU per l'incorporamento del testo
Con shader-f16, il benchmark WebGPU per l'incorporamento del testo di Hugging Face esegue il benchmark 3 volte più velocemente di f32 su laptop Apple M1 Max.

WebLLM è un progetto che può eseguire più modelli linguistici di grandi dimensioni. Utilizza Apache TVM, un framework di compilazione di machine learning open source.

Ho chiesto a WebLLM di pianificare un viaggio a Parigi, utilizzando il modello Llama 3 da otto miliardi di parametri. I risultati mostrano che durante la fase di precompilazione del modello, f16 è 2,1 volte più veloce di f32. Durante la fase di decodifica, è oltre 1,3 volte più veloce.

Le applicazioni devono prima confermare che l'adattatore GPU supporti f16 e, se è disponibile, abilitarlo esplicitamente quando viene richiesto un dispositivo GPU. Se f16 non è supportato, non puoi richiederlo nell'array requiredFeatures.

// main.js

const adapter = await navigator.gpu.requestAdapter();
const supportsF16 = adapter.features.has('shader-f16');
if (supportsF16) {
  // Use f16.
  const device = await adapter.requestDevice({
    requiredFeatures: ['shader-f16'],
  });
  initApp(device);
}

Quindi, negli Shaper WebGPU, devi abilitare esplicitamente f16 nella parte superiore. Dopodiché, puoi utilizzarlo all'interno dello shaker come qualsiasi altro tipo di dati in virgola mobile.

// my-shader.wgsl

enable f16;

struct Data {
  values : array<vec4<f16>>
}
@group(0) @binding(0) var<storage, read> data : Data;
@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid : vec3u) {
  let value : vec4<f16> = data.values[gid.x];
  ...
}

Prodotti con valore intero pacchettizzato

Molti modelli funzionano ancora bene con soli 8 bit di precisione (metà di f16). Questa funzionalità è molto diffusa tra gli LLM e i modelli di immagine per la segmentazione e il riconoscimento degli oggetti. Detto questo, la qualità dell'output per i modelli si riduce con meno precisione, quindi la quantizzazione a 8 bit non è adatta a tutte le applicazioni.

Un numero relativamente ridotto di GPU supporta in modo nativo i valori a 8 bit. È qui che entrano in gioco i prodotti a punti interi compressi. Abbiamo spedito DP4a in Chrome 123.

Le GPU moderne hanno istruzioni speciali per prendere due numeri interi a 32 bit, interpretarli ciascuno come 4 numeri interi a 8 bit pacchettizzati consecutivamente e calcolare il prodotto scalare tra i componenti.

Questo è particolarmente utile per l'IA e il machine learning, perché i kernel di moltiplicazione matriciale sono composti da molti, molti prodotti di tipo "dot".

Ad esempio, proviamo a moltiplicare una matrice 4 x 8 per un vettore 8 x 1. Il calcolo di ciò comporta l'impiego di prodotti a 4 punti per calcolare ciascuno dei valori nel vettore di output. A, B, C e D.

Esempio di diagramma di moltiplicazione matrice-vettore

Il processo per calcolare ognuno di questi output è lo stesso; esamineremo i passaggi necessari per calcolarne uno. Prima di qualsiasi calcolo, dobbiamo prima convertire i dati interi a 8 bit in un tipo con cui possiamo eseguire operazioni aritmetiche, come f16. Quindi, eseguiamo una moltiplicazione per elemento e, infine, aggiungiamo tutti i prodotti. In totale, per l'intera moltiplicazione matriciale-vettore, eseguiamo 40 conversioni numeri interi per fluttuare in modo da decomprimere i dati, 32 moltiplicazioni in virgola mobile e 28 aggiunte in virgola mobile.

Per matrici più grandi con più operazioni, i prodotti a punti interi compressi possono aiutare a ridurre la quantità di lavoro.

Per ciascuno degli output nel vettore dei risultati, eseguiamo due operazioni di addizione del prodotto pacchettizzate utilizzando il linguaggio dot4U8Packed integrato di WebGPU Shading Language, quindi aggiungiamo i risultati. In totale, per l'intera moltiplicazione matrice-vettore, non viene eseguita alcuna conversione di dati. Eseguiamo 8 pacchetti di prodotti scalare e 4 aggiunte di numeri interi.

Diagramma dell&#39;esempio di moltiplicazione tra matrice e vettore di numeri interi compressi

Abbiamo testato prodotti a punti interi con dati a 8 bit su una serie di GPU consumer. Rispetto alla virgola mobile a 16 bit, possiamo vedere che la velocità a 8 bit è da 1,6 a 2,8 volte più veloce. Quando vengono utilizzati anche prodotti con numero intero pacchettizzato, le prestazioni sono persino migliori. È da 1,7 a 2,9 volte più veloce.

Screenshot dell&#39;accelerazione di moltiplicazione del vettore matrice: f16 contro u8
. Grafico 1: Velocità vettoriale a matrice, confrontando f16 con U8 e U8 con dot4U8Packed.

Verifica il supporto del browser con la proprietà wgslLanguageFeatures. Se la GPU non supporta in modo nativo i prodotti di base compressi, il browser esegue il polyfill della propria implementazione.

// main.js

if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
  // Use dot4U8Packed, dot4I8Packed builtin
  // functions in the shaders.
}

La seguente differenza (differenza) degli snippet di codice che evidenzia le modifiche necessarie per utilizzare prodotti con numeri interi pacchettizzati in uno streamr WebGPU.

Prima: uno Shar WebGPU che accumula prodotti di tipo parziale nella variabile "sum". Alla fine del ciclo, "somma" contiene il prodotto di punti completo tra un vettore e una riga della matrice di input.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f16;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : vec4<f16> = vector.values[i];
    let v2 : vec4<f16> = matrix.values[start + i];
    sum += dot(v1, v2);
  }
}

Dopo: uno shaker WebGPU scritto per utilizzare i prodotti basati su valori interi compressi. La differenza principale è che, invece di caricare 4 valori in virgola mobile fuori dal vettore e dalla matrice, questo shaker carica un singolo numero intero a 32 bit. Questo numero intero a 32 bit contiene i dati di quattro valori interi a 8 bit. Quindi, chiamiamo dot4U8Packed per calcolare il prodotto scalare dei due valori.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f32;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : u32 = vector.values[i];
    let v2 : u32 = matrix.values[start + i];
    sum += dot4U8Packed(v1, v2);
  }
}

Sia i prodotti in virgola mobile a 16 bit che i prodotti a virgola mobile compressi sono le funzionalità incluse in Chrome che accelerano l'IA e il machine learning. La rappresentazione in virgola mobile a 16 bit è disponibile quando l'hardware la supporta e Chrome implementa prodotti a virgola mobile compressi su tutti i dispositivi.

Puoi utilizzare queste funzionalità nella versione stabile di Chrome oggi stesso per ottenere prestazioni migliori.

Funzionalità proposte

Guardando al futuro, stiamo esaminando altre due funzionalità: i sottogruppi e la moltiplicazione della matrice cooperativa.

La funzionalità dei sottogruppi consente al parallelismo a livello di SIMD di comunicare o eseguire operazioni matematiche collettive, come la somma di più di 16 numeri. Ciò consente una condivisione efficiente dei dati tra thread. I sottogruppi sono supportati sulle API GPU moderne, con nomi variabili e in forme leggermente diverse.

Abbiamo sintetizzato l'insieme comune in una proposta che abbiamo portato al gruppo di standardizzazione WebGPU. Inoltre, abbiamo prototipato sottogruppi in Chrome dietro un flag sperimentale e abbiamo introdotto i nostri risultati iniziali nella discussione. Il problema principale è come garantire un comportamento portabile.

La moltiplicazione matriciale cooperativa è un'aggiunta più recente alle GPU. Una moltiplicazione matriciale grande può essere suddivisa in più moltiplicazioni matriciali più piccole. La moltiplicazione della matrice cooperativa esegue moltiplicazioni su questi blocchi di dimensioni fisse più piccoli in un singolo passaggio logico. In questo passaggio, un gruppo di thread collabora in modo efficiente per calcolare il risultato.

Abbiamo esaminato il supporto nelle API GPU sottostanti e prevediamo di presentare una proposta al gruppo di standardizzazione WebGPU. Come per i sottogruppi, prevediamo che gran parte della discussione sarà incentrata sulla portabilità.

Per valutare le prestazioni delle operazioni sui sottogruppi, in un'applicazione reale abbiamo integrato il supporto sperimentale per i sottogruppi in MediaPipe e l'abbiamo testato con il prototipo di Chrome per le operazioni sui sottogruppi.

Abbiamo usato sottogruppi nei kernel GPU della fase di precompilazione del modello LLM, quindi riporto solo l'accelerazione per la fase di precompilazione. Su una GPU Intel, notiamo che i sottogruppi sono due volte e mezzo più veloci rispetto alla base di riferimento. Tuttavia, questi miglioramenti non sono coerenti nelle diverse GPU.

Screenshot dell&#39;accelerazione di sottogruppi nell&#39;inferenza LLM di MediaPipe
. Grafico 2. I sottogruppi rendono la precompilazione 2,5 volte più veloce sulla GPU Intel Tiger Lake GT2, con supporto sperimentale in Chrome e Mediapipe.

Il grafico successivo mostra i risultati dell'applicazione di sottogruppi per ottimizzare un microbenchmark di moltiplicazione matriciale su più GPU consumer. La moltiplicazione matriciale è una delle operazioni più complesse nei modelli linguistici di grandi dimensioni (LLM). I dati mostrano che, in molte delle GPU, i sottogruppi aumentano la velocità di due, cinque e addirittura tredici volte la base. Tuttavia, noterai che sulla prima GPU, i sottogruppi non sono molto migliori.

Screenshot dell&#39;accelerazione per sottogruppo per la moltiplicazione matriciale
. Grafico 3. L'applicazione di sottogruppi per la moltiplicazione matriciale può aumentare ulteriormente le prestazioni.

L'ottimizzazione della GPU è difficile

In definitiva, il modo migliore per ottimizzare la GPU dipende dalla GPU offerta dal client. L'utilizzo di nuove e sofisticate funzionalità GPU non sempre dà i risultati sperati, perché possono essere coinvolti molti fattori complessi. La migliore strategia di ottimizzazione su una GPU potrebbe non esserlo su un'altra GPU.

Vuoi ridurre al minimo la larghezza di banda della memoria, utilizzando completamente i thread di calcolo della GPU.

Anche i pattern di accesso alla memoria possono essere molto importanti. Le GPU tendono a funzionare di gran lunga meglio quando i thread di calcolo accedono alla memoria in un pattern ottimale per l'hardware. Importante: dovresti aspettarti caratteristiche di prestazioni diverse su hardware GPU diversi. Potresti dover eseguire ottimizzazioni diverse a seconda della GPU.

Nel grafico seguente abbiamo preso lo stesso algoritmo di moltiplicazione delle matrici, ma abbiamo aggiunto un'altra dimensione per dimostrare ulteriormente l'impatto di varie strategie di ottimizzazione, nonché la complessità e la varianza tra GPU diverse. Abbiamo introdotto una nuova tecnica qui, che chiameremo "Girandola". Swizzle ottimizza i pattern di accesso alla memoria per renderli più ottimali per l'hardware.

Puoi vedere che la variazione della memoria ha un impatto significativo: a volte ha un impatto anche maggiore rispetto ai sottogruppi. Sulla GPU 6, swizzle fornisce una velocità 12x, mentre i sottogruppi forniscono una velocità 13x. Tutti insieme, hanno un'incredibile velocità di 26 volte. Per altre GPU, a volte swizzle e sottogruppi combinati hanno prestazioni migliori di uno solo. Su altre GPU, invece, l'utilizzo esclusivo di swizzle garantisce le migliori prestazioni.

Screenshot della velocità per le strategie di moltiplicazione matriciale
. Grafico 4.

L'ottimizzazione e l'ottimizzazione degli algoritmi della GPU in modo che funzionino bene su qualsiasi hardware può richiedere molta esperienza. Per fortuna c'è un'enorme quantità di lavoro di talento dedicato ai framework di librerie di livello superiore, come Mediapipe, Transformers.js, Apache TVM, ONNX Runtime Web e altri ancora.

Librerie e framework sono ben posizionati per gestire la complessità della gestione di diverse architetture GPU e della generazione di codice specifico della piattaforma che funzionerà bene sul client.

Concetti principali

Il team di Chrome continua a contribuire all'evoluzione degli standard WebAssembly e WebGPU per migliorare la piattaforma web per i carichi di lavoro di machine learning. Stiamo investendo in primitive di calcolo più veloci, in una migliore interoperabilità tra gli standard web e garantendo che modelli, grandi e piccoli, siano in grado di funzionare in modo efficiente su tutti i dispositivi.

Il nostro obiettivo è massimizzare le funzionalità della piattaforma senza rinunciare al meglio del web: copertura, usabilità e portabilità. E non lo facciamo da soli. Stiamo collaborando con gli altri fornitori di browser di W3C e con molti partner di sviluppo.

Al momento di lavorare con WebAssembly e WebGPU, ci auguriamo che tu possa tenere a mente quanto segue:

  • L'inferenza IA è ora disponibile sul web, su tutti i dispositivi. Ciò offre il vantaggio dell'esecuzione sui dispositivi client, come la riduzione dei costi dei server, la bassa latenza e una maggiore privacy.
  • Anche se molte funzionalità discusse sono rilevanti principalmente per gli autori del framework, le tue applicazioni possono trarne vantaggio senza un overhead eccessivo.
  • Gli standard del web sono fluidi e in evoluzione, perciò siamo sempre alla ricerca di feedback. Condividi il tuo per WebAssembly e WebGPU.

Ringraziamenti

Desideriamo ringraziare il team di grafica web Intel, che è stato determinante nella realizzazione di WebGPU f16 e di funzionalità del prodotto "intero". Vorremmo ringraziare gli altri membri dei gruppi di lavoro WebAssembly e WebGPU presso W3C, inclusi gli altri fornitori di browser.

Grazie ai team di IA e ML sia di Google che della community open source per essere partner incredibili. Ovviamente, anche di tutti i membri del nostro team che rendono possibile tutto questo.