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

Questo documento è la continuazione di Miglioramenti di WebAssembly e WebGPU per un'AI web più veloce, parte 1. Ti consigliamo di leggere questo post o di guardare il talk all'IO 24 prima di continuare.

Austin Eng

Deepti Gandluri

François Beaufort

WebGPU

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

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

L'animazione seguente 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 del modello Segment Anything di Meta produce maschere di oggetti di alta qualità interamente sul client.

Questi sono solo alcuni dei fantastici progetti che mostrano la potenza di WebGPU per l'IA e il ML. WebGPU consente a questi modelli e ad altri di funzionare molto più velocemente rispetto alla CPU.

Il benchmark WebGPU per l'embedding del testo di Hugging Face dimostra enormi accelerazioni rispetto a un'implementazione su CPU dello stesso modello. Su un laptop Apple M1 Max, WebGPU è stato più di 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 è ideale per i modelli di AI e ML, che possono avere miliardi di parametri, grazie al supporto degli shader di calcolo. Gli shader di calcolo vengono eseguiti sulla GPU e consentono di eseguire operazioni di array parallele 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 AI sul web. Di recente abbiamo lanciato due nuove funzionalità: i prodotti in virgola mobile a 16 bit e i prodotti in virgola mobile interi pacchettizzati.

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 shading WebGPU. Questo tipo di numeri in virgola mobile occupa 16 bit, anziché i consueti 32 bit. f16 ha un intervallo più ridotto ed è meno preciso, ma per molti modelli di ML è sufficiente.

Questa funzionalità aumenta l'efficienza in diversi modi:

• Memoria ridotta: i tensori con elementi f16 occupano metà dello spazio, il che dimezza l'utilizzo della memoria. I calcoli della GPU sono spesso limitati dalla larghezza di banda della memoria, quindi dimezzare la memoria può spesso significare che gli shader vengono eseguiti due volte più velocemente. Tecnicamente, non è necessario f16 per risparmiare sulla larghezza di banda della memoria. È possibile archiviare i dati in un formato a bassa precisione e poi espanderli in f32 completo nello shader per il calcolo. Tuttavia, la GPU utilizza una potenza di calcolo aggiuntiva per imballare e sballare i dati.

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

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

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

Ho chiesto a WebLLM di pianificare un viaggio a Parigi utilizzando il modello Llama 3 con 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, è più di 1,3 volte più veloce.

Le applicazioni devono prima verificare che l'adattatore GPU supporti f16 e, se è disponibile, attivarlo esplicitamente quando richiedono 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);
}

Poi, negli shader WebGPU, devi attivare esplicitamente f16 in alto. Dopodiché puoi utilizzarlo all'interno dello shader come qualsiasi altro tipo di dati float.

// 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 punti interi impacchettati

Molti modelli funzionano ancora bene con una precisione di soli 8 bit (la metà di f16). Questo approccio è molto utilizzato tra i modelli LLM e di immagini per la segmentazione e il riconoscimento degli oggetti. Detto questo, la qualità dell'output dei modelli peggiora con una precisione inferiore, pertanto la quantizzazione a 8 bit non è adatta a tutte le applicazioni.

Poche GPU supportano in modo nativo i valori a 8 bit. È qui che entrano in gioco i prodotti con punti interi impacchettati. Abbiamo rilasciato la DP4a in Chrome 123.

Le GPU moderne dispongono di istruzioni speciali per prendere due numeri interi a 32 bit, interpretarli ciascuno come quattro numeri interi a 8 bit con imballaggio consecutivo e calcolare il prodotto scalare tra i relativi componenti.

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

Ad esempio, moltiplichiamo una matrice 4 x 8 con un vettore 8 x 1. Il calcolo prevede l'esecuzione di 4 prodotti scalari per calcolare ciascuno dei valori nel vettore di output: A, B, C e D.

La procedura per calcolare ciascuno di questi output è la stessa. 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, ad esempio f16. Poi eseguiamo una moltiplicazione elemento per elemento e infine aggiungiamo tutti i prodotti. In totale, per l'intera moltiplicazione di matrici e vettori, vengono eseguite 40 conversioni da interi a valori float per scompattare i dati, 32 moltiplicazioni di valori float e 28 addizioni di valori float.

Per matrici più grandi con più operazioni, i prodotti in virgola mobile interi pacchettizzati possono contribuire a ridurre la quantità di lavoro.

Per ogni output del vettore del risultato, eseguiamo due operazioni di prodotto scalare imballato utilizzando il linguaggio di shading WebGPU integrato dot4U8Packed e poi sommiamo i risultati. In totale, per l'intera moltiplicazione di matrici e vettori, non viene eseguita alcuna conversione dei dati. Eseguiamo 8 prodotti in punti imballati e 4 addizioni di interi.

Abbiamo testato i prodotti in virgola mobile interi pacchettizzati con dati a 8 bit su una serie di GPU consumer. Rispetto al floating point a 16 bit, possiamo vedere che il floating point a 8 bit è da 1,6 a 2,8 volte più veloce. Se utilizziamo anche prodotti in virgola tra interi impacchettati, il rendimento è ancora migliore. È da 1,7 a 2,9 volte più veloce.

Verifica il supporto del browser con la proprietà wgslLanguageFeatures. Se la GPU non supporta in modo nativo i puntini raggruppati, 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 (diff) dello snippet di codice evidenzia le modifiche necessarie per utilizzare i prodotti interi pacchettizzati in uno shader WebGPU.

Prima: uno shader WebGPU che accumula prodotti scalari parziali nella variabile "somma". Al termine del ciclo, "somma" contiene il prodotto scalare 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);
  }
}

After: uno shader WebGPU scritto per utilizzare prodotti scalari interi pacchettizzati. La differenza principale è che, anziché caricare 4 valori float dal vettore e dalla matrice, questo shader 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 sia i prodotti in virgola mobile con numeri interi pacchettizzati sono le funzionalità fornite in Chrome che accelerano l'AI e l'ML. I numeri in virgola mobile a 16 bit sono disponibili se l'hardware li supporta e Chrome implementa i prodotti in virgola mobile interi pacchettizzati su tutti i dispositivi.

Puoi utilizzare queste funzionalità in Chrome Stable oggi stesso per ottenere un rendimento migliore.

Funzionalità proposte

In futuro, esamineremo altre due funzionalità: i sottogruppi e la moltiplicazione di matrici cooperative.

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

Abbiamo distillato l'insieme comune in una proposta che abbiamo presentato al gruppo di standardizzazione WebGPU. Inoltre, abbiamo implementato un prototipo di sottogruppi in Chrome dietro un flag sperimentale e abbiamo presentato i nostri risultati iniziali nella discussione. Il problema principale è come garantire il comportamento portatile.

La moltiplicazione di matrici cooperativa è una funzionalità più recente delle GPU. Una moltiplicazione di matrici di grandi dimensioni può essere suddivisa in più moltiplicazioni di matrici più piccole. La moltiplicazione di matrici cooperativa esegue le moltiplicazioni su questi blocchi di dimensioni fisse più piccoli in un unico 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 riguarderà la portabilità.

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

Abbiamo utilizzato sottogruppi nei kernel GPU della fase di precompilazione del grande modello linguistico, quindi riporto solo l'aumento di velocità per la fase di precompilazione. Su una GPU Intel, abbiamo notato che i sottogruppi hanno un rendimento due volte e mezzo superiore rispetto al valore di riferimento. Tuttavia, questi miglioramenti non sono coerenti tra GPU diverse.

Il grafico seguente mostra i risultati dell'applicazione di sottogruppi per ottimizzare un microbenchmark di moltiplicazione di matrici su più GPU consumer. La moltiplicazione di matrici è una delle operazioni più pesanti nei modelli linguistici di grandi dimensioni. I dati mostrano che su molte GPU, i sottogruppi aumentano la velocità di due, cinque e persino tredici volte rispetto al valore di riferimento. Tuttavia, tieni presente che sulla prima GPU i sottogruppi non sono molto migliori.

L'ottimizzazione della GPU è difficile

In definitiva, il modo migliore per ottimizzare la GPU dipende dalla GPU offerta dal cliente. L'utilizzo di nuove funzionalità avanzate della GPU non sempre ripaga come ci si aspetterebbe, perché possono essere coinvolti molti fattori complessi. La strategia di ottimizzazione migliore su una GPU potrebbe non essere la migliore su un'altra.

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

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

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

Puoi notare che la modifica della memoria ha un impatto significativo; a volte è persino più efficace dei sottogruppi. Sulla GPU 6, la permutazione offre un aumento di velocità di 12 volte, mentre i sottogruppi offrono un aumento di velocità di 13 volte. Insieme, offrono un incredibile aumento di velocità pari a 26 volte. Per altre GPU, a volte la combinazione di swizzle e sottogruppi ha un rendimento migliore rispetto a ciascuno da solo. Su altre GPU, il rendimento migliore si ottiene utilizzando esclusivamente swizzle.

La messa a punto e l'ottimizzazione degli algoritmi GPU per il corretto funzionamento su ogni componente hardware può richiedere molta esperienza. Fortunatamente, però, c'è un'enorme quantità di lavoro di talento in corso per i framework di librerie di livello superiore, come Mediapipe, Transformers.js, Apache TVM, ONNX Runtime Web e altri ancora.

Le librerie e i framework sono ben posizionati per gestire la complessità della gestione di diverse architetture GPU e generare codice specifico per la piattaforma che funzioni 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 ci assicuriamo che i modelli, sia grandi che piccoli, possano funzionare in modo efficiente su tutti i dispositivi.

Il nostro obiettivo è massimizzare le funzionalità della piattaforma mantenendo al contempo il meglio del web: la sua copertura, usabilità e portabilità. E non lo stiamo facendo da soli. Stiamo collaborando con gli altri fornitori di browser del W3C e con molti partner di sviluppo.

Quando lavori con WebAssembly e WebGPU, ti consigliamo di ricordare quanto segue:

• L'inferenza AI è ora disponibile sul web e su tutti i dispositivi. Ciò offre il vantaggio di essere eseguito sui dispositivi client, ad esempio costi ridotti del server, bassa latenza e maggiore privacy.
• Sebbene molte delle funzionalità discusse siano pertinenti principalmente per gli autori del framework, le tue applicazioni possono trarne vantaggio senza un eccessivo sovraccarico.
• Gli standard web sono fluidi e in continua evoluzione e siamo sempre alla ricerca di feedback. Condividi le tue per WebAssembly e WebGPU.

Ringraziamenti

Vogliamo ringraziare il team di grafica web di Intel, che ha contribuito in modo determinante allo sviluppo delle funzionalità di prodotto in virgola mobile f16 e di prodotto intero con imballaggio di WebGPU. Ringraziamo gli altri membri dei gruppi di lavoro WebAssembly e WebGPU del W3C, inclusi gli altri fornitori di browser.

Grazie ai team di IA e ML di Google e della community open source per essere partner incredibili. E, naturalmente, a tutti i nostri colleghi che rendono possibile tutto questo.