WebAssembly- und WebGPU-Verbesserungen für schnellere Web AI, Teil 2

Dieses Dokument ist eine Fortsetzung von WebAssembly- und WebGPU-Erweiterungen für schnellere Web-KI, Teil 1. Wir empfehlen Ihnen, diesen Beitrag zu lesen oder sich den Vortrag auf der IO 24 noch einmal anzusehen, bevor Sie fortfahren.

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

WebGPU

WebGPU bietet Webanwendungen Zugriff auf die GPU-Hardware des Clients, um effiziente, hochparallele Berechnungen durchzuführen. Seit der Einführung von WebGPU in Chrome haben wir unglaubliche Demos zu künstlicher Intelligenz (KI) und maschinellem Lernen (ML) im Web gesehen.

Mit Web Stable Diffusion konnte beispielsweise gezeigt werden, dass es möglich ist, mithilfe von KI Bilder direkt im Browser aus Text zu generieren. Anfang des Jahres hat das Mediapipe-Team von Google experimentelle Unterstützung für die Inferenz von Large Language Models veröffentlicht.

In der folgenden Animation sehen Sie Gemma, das Open-Source-Large Language Model (LLM) von Google, das vollständig auf dem Gerät in Chrome ausgeführt wird.

In der folgenden Demo von Hugging Face des Segment Anything-Modells von Meta werden hochwertige Objektmasken vollständig auf dem Client erstellt.

Dies sind nur einige der beeindruckenden Projekte, die die Leistungsfähigkeit von WebGPU für KI und ML zeigen. Mit WebGPU können diese und andere Modelle deutlich schneller ausgeführt werden als auf der CPU.

Der WebGPU-Benchmark für Text-Embeddings von Hugging Face zeigt enorme Geschwindigkeitssteigerungen im Vergleich zu einer CPU-Implementierung desselben Modells. Auf einem Apple M1 Max-Laptop war WebGPU über 30-mal schneller. Andere haben berichtet, dass WebGPU den Benchmark um mehr als das 120-fache beschleunigt.

Verbesserung der WebGPU-Funktionen für KI und ML

WebGPU eignet sich hervorragend für KI- und ML-Modelle, die dank der Unterstützung von Compute Shadern Milliarden von Parametern haben können. Compute Shader werden auf der GPU ausgeführt und unterstützen parallele Array-Vorgänge für große Datenmengen.

Neben den zahlreichen Verbesserungen an WebGPU im letzten Jahr haben wir weitere Funktionen hinzugefügt, um die ML- und KI-Leistung im Web zu verbessern. Vor Kurzem haben wir zwei neue Funktionen eingeführt: 16‑Bit-Gleitkomma- und Packte-Ganzzahl-Punktprodukte.

16-Bit-Gleitkomma

Denken Sie daran, dass ML-Arbeitslasten keine Präzision erfordern. shader-f16 ist eine Funktion, die die Verwendung des Typs „f16“ in der WebGPU-Shading-Sprache ermöglicht. Dieser Gleitkommatyp belegt 16 Bit anstelle der üblichen 32 Bit. f16 hat einen kleineren Bereich und ist weniger genau, aber für viele ML-Modelle ist das ausreichend.

Diese Funktion erhöht die Effizienz auf verschiedene Weise:

  • Reduzierter Arbeitsspeicher: Tensoren mit f16-Elementen belegen nur noch die Hälfte des Speicherplatzes, wodurch die Speichernutzung halbiert wird. GPU-Berechnungen sind oft durch die Arbeitsspeicherbandbreite begrenzt. Daher kann es sein, dass Shader mit der Hälfte des Arbeitsspeichers doppelt so schnell laufen. Technisch gesehen benötigen Sie f16 nicht, um die Arbeitsspeicherbandbreite zu sparen. Es ist möglich, die Daten in einem Format mit niedriger Präzision zu speichern und sie dann im Shader für die Berechnung auf volle f32-Werte auszuweiten. Die GPU benötigt jedoch zusätzliche Rechenleistung, um die Daten zu verpacken und zu entpacken.

  • Verringerte Datenkonvertierung: f16 benötigt weniger Rechenleistung, da die Datenkonvertierung minimiert wird. Daten mit niedriger Genauigkeit können gespeichert und dann direkt ohne Umwandlung verwendet werden.

  • Erhöhter Parallelismus: Moderne GPUs können mehr Werte gleichzeitig in die Ausführungseinheiten der GPU einfügen, wodurch mehr parallele Berechnungen ausgeführt werden können. Eine GPU, die beispielsweise bis zu 5 Billionen f32-Gleitkommaoperationen pro Sekunde unterstützt, kann 10 Billionen f16-Gleitkommaoperationen pro Sekunde unterstützen.

Screenshot des WebGPU-Benchmarks für Text-Embedding
Mit shader-f16 läuft der WebGPU-Benchmark für das Text-Embedding von Hugging Face dreimal schneller als f32 auf einem Apple M1 Max-Laptop.

WebLLM ist ein Projekt, mit dem mehrere Large Language Models ausgeführt werden können. Dabei wird Apache TVM verwendet, ein Open-Source-Compiler-Framework für maschinelles Lernen.

Ich habe WebLLM gebeten, eine Reise nach Paris mit dem Llama 3-Modell mit acht Milliarden Parametern zu planen. Die Ergebnisse zeigen, dass f16 während der Prefill-Phase des Modells 2,1-mal schneller ist als f32. Während der Dekodierungsphase ist es über 1, 3-mal schneller.

Anwendungen müssen zuerst prüfen, ob der GPU-Adapter f16 unterstützt, und ihn bei der Anforderung eines GPU-Geräts explizit aktivieren, falls er verfügbar ist. Wenn „f16“ nicht unterstützt wird, kannst du es nicht im requiredFeatures-Array anfordern.

// 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);
}

Aktivieren Sie dann in Ihren WebGPU-Shadern oben explizit f16. Danach können Sie ihn wie jeden anderen Float-Datentyp im Shader verwenden.

// 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];
  ...
}

Gepackte Ganzzahl-Punktprodukte

Viele Modelle funktionieren auch mit nur 8 Bit Genauigkeit (die Hälfte von f16) gut. Diese Methode ist bei LLMs und Bildmodellen für Segmentierung und Objekterkennung beliebt. Die Ausgabequalität von Modellen verschlechtert sich jedoch mit weniger Genauigkeit. Daher ist die 8‑Bit-Quantisierung nicht für jede Anwendung geeignet.

Nur relativ wenige GPUs unterstützen nativ 8‑Bit-Werte. Hier kommen komprimierte Ganzzahl-Punktprodukte ins Spiel. Wir haben DP4a in Chrome 123 eingeführt.

Moderne GPUs haben spezielle Anweisungen, um zwei 32‑Bit-Ganzzahlen zu nehmen, sie jeweils als vier aufeinanderfolgend verpackte 8‑Bit-Ganzzahlen zu interpretieren und die Punktprodukt zwischen ihren Komponenten zu berechnen.

Das ist besonders nützlich für KI und maschinelles Lernen, da Matrixmultiplikationskerne aus vielen, vielen Skalarprodukten bestehen.

Multiplizieren wir beispielsweise eine 4 × 8-Matrix mit einem 8 × 1-Vektor. Dazu sind vier Skalarprodukte erforderlich, um die einzelnen Werte im Ausgabevektor A, B, C und D zu berechnen.

Diagramm für ein Beispiel für die Matrix-Vektor-Multiplikation

Die Berechnung dieser Ausgaben erfolgt auf die gleiche Weise. Wir sehen uns die Schritte zur Berechnung einer davon an. Vor jeder Berechnung müssen wir die 8-Bit-Ganzzahldaten in einen Typ konvertieren, mit dem wir arithmetische Operationen ausführen können, z. B. f16. Anschließend führen wir eine elementweise Multiplikation durch und addieren schließlich alle Produkte. Insgesamt führen wir für die gesamte Matrix-Vektor-Multiplikation 40 Ganzzahl-zu-Float-Konvertierungen aus, um die Daten zu entpacken, 32 Float-Multiplikationen und 28 Float-Additionen.

Bei größeren Matrizen mit mehr Vorgängen können gepackte Ganzzahl-Punktprodukte den Arbeitsaufwand reduzieren.

Für jede der Ausgaben im Ergebnisvektor führen wir zwei gepackte Punktproduktoperationen mit der integrierten dot4U8Packed der WebGPU-Shading-Language aus und addieren dann die Ergebnisse. Insgesamt führen wir für die gesamte Matrix-Vektor-Multiplikation keine Datenkonvertierung durch. Wir führen 8 gepackte Punktprodukte und 4 Ganzzahladditionen aus.

Diagramm für ein Beispiel für die Multiplikation von Ganzzahlmatrizen und -vektoren

Wir haben gepackte Ganzzahlpunktprodukte mit 8‑Bit-Daten auf einer Vielzahl von Verbraucher-GPUs getestet. Im Vergleich zu 16-Bit-Gleitkommazahlen ist 8-Bit 1,6- bis 2,8-mal schneller. Wenn wir zusätzlich verpackte Ganzzahl-Punktprodukte verwenden, ist die Leistung noch besser. Sie ist 1,7- bis 2,9-mal schneller.

Screenshot der Beschleunigung der Matrix-Vektor-Multiplikation: f16 im Vergleich zu u8
Abbildung 1: Beschleunigung von Matrixvektoren, Vergleich von f16 mit U8 und U8 mit dot4U8Packed.

Prüfen Sie mit dem wgslLanguageFeatures-Attribut, ob der Browser unterstützt wird. Wenn die GPU keine nativ unterstützten gepackten Punktprodukte unterstützt, führt der Browser eine eigene polyfill-Implementierung aus.

// main.js

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

Im folgenden Code-Snippet wird der Unterschied zwischen den beiden Code-Versionen verdeutlicht, die für die Verwendung von komprimierten Ganzzahlprodukten in einem WebGPU-Shader erforderlich sind.

Vorher: Ein WebGPU-Shader, der partielle Skalarprodukte in der Variablen „sum“ sammelt. Am Ende der Schleife enthält „sum“ das vollständige Skalarprodukt zwischen einem Vektor und einer Zeile der Eingabematrix.

// 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: Ein WebGPU-Shader, der für die Verwendung von verpackten Ganzzahl-Punktprodukten geschrieben wurde. Der Hauptunterschied besteht darin, dass dieser Shader anstelle von vier Gleitkommawerten aus dem Vektor und der Matrix eine einzelne 32‑Bit-Ganzzahl lädt. Diese 32‑Bit-Ganzzahl enthält die Daten von vier 8‑Bit-Ganzzahlwerten. Anschließend rufen wir dot4U8Packed auf, um das Skalarprodukt der beiden Werte zu berechnen.

// 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);
  }
}

Sowohl 16-Bit-Gleitkomma- als auch Packte-Ganzzahl-Punktprodukte sind in Chrome verfügbare Funktionen, die KI und ML beschleunigen. 16-Bit-Gleitkommazahlen sind verfügbar, wenn die Hardware sie unterstützt. Chrome implementiert auf allen Geräten gepackte Ganzzahl-Punktprodukte.

Sie können diese Funktionen bereits jetzt in der stabilen Chrome-Version verwenden, um die Leistung zu verbessern.

Vorgeschlagene Funktionen

Wir arbeiten derzeit an zwei weiteren Funktionen: Subgruppen und kooperative Matrixmultiplikation.

Die Funktion „Untergruppen“ ermöglicht die Kommunikation oder die Ausführung kollektiver mathematischer Vorgänge auf SIMD-Ebene, z. B. die Summe von mehr als 16 Zahlen. Das ermöglicht eine effiziente threadübergreifende Datenfreigabe. Untergruppen werden in modernen GPUs APIs mit unterschiedlichen Namen und in leicht unterschiedlichen Formen unterstützt.

Wir haben die gemeinsamen Funktionen in einem Vorschlag zusammengefasst, den wir der WebGPU-Standardisierungsgruppe vorgelegt haben. Außerdem haben wir Prototypen für Untergruppen in Chrome mithilfe einer experimentellen Flag erstellt und unsere ersten Ergebnisse in die Diskussion eingebracht. Das Hauptproblem besteht darin, wie das portable Verhalten sichergestellt werden kann.

Die kooperative Matrixmultiplikation ist eine neuere Funktion von GPUs. Eine große Matrixmultiplikation kann in mehrere kleinere Matrixmultiplikationen zerlegt werden. Bei der kooperativen Matrixmultiplikation werden diese kleineren Blöcke mit fester Größe in einem einzigen logischen Schritt multipliziert. In diesem Schritt arbeiten mehrere Threads effizient zusammen, um das Ergebnis zu berechnen.

Wir haben die Unterstützung in den zugrunde liegenden GPU-APIs untersucht und planen, der WebGPU-Standardisierungsgruppe einen Vorschlag zu unterbreiten. Wie bei den Untergruppen gehen wir davon aus, dass sich ein Großteil der Diskussion um die Portabilität drehen wird.

Um die Leistung von Subgruppenoperationen zu bewerten, haben wir in einer echten Anwendung experimentelle Unterstützung für Subgruppen in MediaPipe integriert und mit dem Chrome-Prototyp für Subgruppenoperationen getestet.

Wir haben Subgruppen in GPU-Kerneln der Prefill-Phase des Large Language Models verwendet. Daher wird nur die Beschleunigung für die Prefill-Phase angegeben. Auf einer Intel-GPU sind Subgruppen zweieinhalbmal schneller als die Baseline. Diese Verbesserungen sind jedoch nicht für alle GPUs gleich.

Screenshot der Beschleunigung von Untergruppen bei der LLM-Inferenz in MediaPipe
Diagramm 2. Mit Untergruppen wird das Prefill auf der Intel Tiger Lake GT2 GPU 2, 5-mal schneller ausgeführt.Chrome und Mediapipe unterstützen diese Funktion derzeit nur experimentell.

Das nächste Diagramm zeigt die Ergebnisse der Anwendung von Untergruppen zur Optimierung eines Matrixmultiplikations-Mikrobenchmarks auf mehreren Consumer-GPUs. Die Matrixmultiplikation ist einer der aufwendigsten Vorgänge in Large Language Models. Die Daten zeigen, dass bei vielen GPUs die Geschwindigkeit durch Untergruppen um das Zwei-, Fünf- und sogar Dreizehnfache gesteigert wird. Beachten Sie jedoch, dass die Untergruppen bei der ersten GPU nicht viel besser abschneiden.

Screenshot der Subgruppenbeschleunigung für die Matrixmultiplikation
Diagramm 3. Durch die Anwendung von Untergruppen für die Matrixmultiplikation lässt sich die Leistung weiter steigern.

GPU-Optimierung ist schwierig

Die beste Optimierungsmethode für die GPU hängt letztendlich davon ab, welche GPU der Kunde anbietet. Die Verwendung neuer GPU-Funktionen zahlt sich nicht immer so aus, wie Sie es erwarten würden, da viele komplexe Faktoren eine Rolle spielen können. Die beste Optimierungsstrategie für eine GPU ist möglicherweise nicht die beste Strategie für eine andere GPU.

Sie möchten die Arbeitsspeicherbandbreite minimieren und gleichzeitig die Rechenthreads der GPU voll ausnutzen.

Auch die Speicherzugriffsmuster können sehr wichtig sein. GPUs erzielen in der Regel eine deutlich bessere Leistung, wenn die Compute-Threads auf den Arbeitsspeicher in einem für die Hardware optimalen Muster zugreifen. Wichtig: Die Leistung kann je nach GPU-Hardware variieren. Je nach GPU müssen Sie möglicherweise verschiedene Optimierungen ausführen.

Im folgenden Diagramm haben wir denselben Algorithmus für die Matrixmultiplikation verwendet, aber eine weitere Dimension hinzugefügt, um die Auswirkungen verschiedener Optimierungsstrategien sowie die Komplexität und Varianz bei verschiedenen GPUs zu veranschaulichen. Hier haben wir ein neues Verfahren eingeführt, das wir „Swizzle“ nennen. Durch Swizzle werden die Speicherzugriffsmuster für die Hardware optimiert.

Sie sehen, dass die Speicher-Swizzle eine erhebliche Auswirkung hat, manchmal sogar eine größere als Untergruppen. Bei GPU 6 führt die Swizzle-Technologie zu einer 12-fachen Beschleunigung, während Subgruppen eine 13-fache Beschleunigung bieten. Zusammengenommen ergibt sich eine unglaubliche Beschleunigung von 26-fach. Bei anderen GPUs erzielen ggf. eine Kombination aus Swizzle und Untergruppen eine bessere Leistung als die einzelnen Optionen. Bei anderen GPUs ist die ausschließliche Verwendung von swizzle am effektivsten.

Screenshot der Beschleunigung bei Strategien für die Matrixmultiplikation
Diagramm 4.

Das Optimieren von GPU-Algorithmen, damit sie auf jeder Hardware gut funktionieren, erfordert viel Fachwissen. Glücklicherweise gibt es aber eine enorme Menge an talentierter Arbeit in Bibliotheken und Frameworks höherer Ebene, wie Mediapipe, Transformers.js, Apache TVM und ONNX Runtime Web.

Bibliotheken und Frameworks sind gut gerüstet, um die Komplexität der Verwaltung verschiedener GPU-Architekturen zu bewältigen und plattformspezifischen Code zu generieren, der auf dem Client gut funktioniert.

Fazit

Das Chrome-Team trägt weiterhin dazu bei, die WebAssembly- und WebGPU-Standards weiterzuentwickeln, um die Webplattform für Arbeitslasten im Bereich Machine Learning zu verbessern. Wir investieren in schnellere Computing-Primitive, eine bessere Interoperabilität zwischen Webstandards und sorgen dafür, dass sowohl große als auch kleine Modelle geräteübergreifend effizient ausgeführt werden können.

Unser Ziel ist es, die Funktionen der Plattform zu maximieren und gleichzeitig das Beste aus dem Web zu erhalten: Reichweite, Nutzerfreundlichkeit und Portabilität. Und wir sind nicht allein. Wir arbeiten mit den anderen Browseranbietern im W3C und vielen Entwicklungspartnern zusammen.

Wir hoffen, dass Sie bei der Arbeit mit WebAssembly und WebGPU Folgendes berücksichtigen:

  • KI-Inferenzen sind jetzt im Web und geräteübergreifend verfügbar. Dies bietet Vorteile wie geringere Serverkosten, eine geringe Latenz und mehr Datenschutz.
  • Viele der beschriebenen Funktionen sind zwar in erster Linie für die Framework-Entwickler relevant, aber auch Ihre Anwendungen können davon profitieren, ohne dass viel Mehraufwand entsteht.
  • Webstandards sind fließend und entwickeln sich weiter. Wir freuen uns immer über Feedback. Teilen Sie uns Ihre für WebAssembly und WebGPU mit.

Danksagungen

Wir möchten uns beim Intel-Webgrafikteam bedanken, das maßgeblich an der Entwicklung der WebGPU-f16- und der Funktionen für gepackte Ganzzahl-Punktprodukte beteiligt war. Wir möchten uns bei den anderen Mitgliedern der WebAssembly- und WebGPU-Arbeitsgruppen im W3C bedanken, einschließlich der anderen Browseranbieter.

Vielen Dank an die KI- und ML-Teams bei Google und in der Open-Source-Community für ihre großartige Unterstützung. Und natürlich an alle unsere Teammitglieder, die das alles möglich machen.