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

Dieses Dokument ist eine Fortsetzung von WebAssembly- und WebGPU-Verbesserungen für schnellere Web AI, Teil 1. Wir empfehlen Ihnen, diesen Beitrag zu lesen oder sich den Vortrag bei I/O 24 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 eine effiziente, hochparallele Berechnung durchzuführen. Seit der Einführung der WebGPU in Chrome gibt es unglaubliche Demonstrationen von künstlicher Intelligenz (KI) und maschinellem Lernen (ML) im Web.

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

Die folgende Animation zeigt Gemma, das Large Language Model (LLM) von Google, das vollständig direkt auf dem Gerät in Chrome ausgeführt wird.

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

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

Der WebGPU-Benchmark für Texteinbettungen von Hugging Face zeigt im Vergleich zur CPU-Implementierung desselben Modells enorme Beschleunigungen. Auf einem Apple M1 Max-Laptop war WebGPU mehr als 30-mal schneller. Andere berichteten, dass WebGPU die Benchmark um mehr als 120-mal beschleunigt.

Verbesserung der WebGPU-Funktionen für KI und ML

WebGPU eignet sich dank der Unterstützung von Computing-Shadern hervorragend für KI- und ML-Modelle, die Milliarden von Parametern haben können. Compute-Shader werden auf der GPU ausgeführt und tragen dazu bei, parallele Arrayvorgänge für große Datenmengen auszuführen.

Neben den zahlreichen Verbesserungen von WebGPU im letzten Jahr haben wir weitere Funktionen hinzugefügt, um die ML- und KI-Leistung im Web zu verbessern. Kürzlich haben wir zwei neue Funktionen eingeführt: 16-Bit-Gleitkomma- und gepackte Integer-Punkt-Produkte.

16-Bit-Gleitkomma

Denken Sie daran, dass ML-Arbeitslasten keine Präzision erfordern. shader-f16 ist ein Feature, das die Verwendung des Typs f16 in der WebGPU-Shading-Sprache ermöglicht. Dieser Gleitkommatyp benötigt 16 Bit anstelle der üblichen 32 Bit. f16 hat einen kleineren Bereich und ist weniger präzise, aber für viele ML-Modelle reicht dies aus.

Diese Funktion erhöht die Effizienz in mehreren Bereichen:

  • Geringerer Arbeitsspeicher: Tensoren mit f16-Elementen nehmen nur die Hälfte des Speicherplatzes ein, wodurch sich der Arbeitsspeicherverbrauch halbieren lässt. Bei GPU-Berechnungen kommt es häufig zu Engpässen bei der Speicherbandbreite. Daher kann die Hälfte des Arbeitsspeichers oft bedeuten, dass Shader doppelt so schnell ausgeführt werden. Technisch gesehen benötigen Sie f16 nicht, um Speicherbandbreite zu sparen. Es ist möglich, die Daten in einem Format mit niedriger Genauigkeit zu speichern und dann für die Berechnung auf das vollständige f32-Format im Shader auszudehnen. Allerdings benötigt die GPU zusätzliche Rechenleistung, um die Daten zu verpacken und zu entpacken.

  • Reduzierte Datenkonvertierung: f16 benötigt weniger Rechenressourcen, da die Datenkonvertierung minimiert wird. Daten mit geringer Genauigkeit können gespeichert und dann direkt und ohne Konvertierung verwendet werden.

  • Erhöhte Parallelität: Moderne GPUs können mehr Werte gleichzeitig in die Ausführungseinheiten der GPU aufnehmen, wodurch mehr parallele Berechnungen durchgeführt werden können. Beispielsweise unterstützt eine GPU, die bis zu 5 Billionen f32-Gleitkommavorgänge pro Sekunde unterstützt, möglicherweise 10 Billionen f16-Gleitkommavorgänge pro Sekunde.

Screenshot des WebGPU-Benchmarks für Texteinbettungen
Mit shader-f16 führt Hugging Faces WebGPU-Benchmark für die Texteinbettung die Benchmark dreimal schneller aus als f32 auf einem Apple M1 Max-Laptop.

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

Ich habe WebLLM gebeten, eine Reise nach Paris zu planen. Dafür habe ich das acht Milliarden Parameter-Modell von Llama 3 verwendet. Die Ergebnisse zeigen, dass f16 während der Vorabfüllphase des Modells 2,1-mal schneller ist als f32. Während der Decodierungsphase ist es mehr als 1, 3-mal schneller.

Anwendungen müssen zuerst bestätigen, dass der GPU-Adapter f16 unterstützt. Wenn f16 verfügbar ist, muss er beim Anfordern eines GPU-Geräts explizit aktiviert werden. Wenn f16 nicht unterstützt wird, können Sie es nicht im Array requiredFeatures 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);
}

Dann müssen Sie oben in Ihren WebGPU-Shadern f16 explizit aktivieren. 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 Integer-Punktprodukte

Viele Modelle funktionieren auch mit einer Genauigkeit von nur 8 Bit (Hälfte von f16) gut. Dies ist bei LLMs und Bildmodellen zur Segmentierung und Objekterkennung beliebt. Allerdings verschlechtert sich die Ausgabequalität von Modellen mit geringerer Genauigkeit, sodass die 8-Bit-Quantisierung nicht für jede Anwendung geeignet ist.

Nur relativ wenige GPUs unterstützen 8-Bit-Werte nativ. Hier kommen gepackte Integer-Punkt-Produkte ins Spiel. Wir haben DP4a in Chrome 123 ausgeliefert.

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

Dies ist besonders nützlich für KI und maschinelles Lernen, da die Kernel der Matrixmultiplikation aus vielen, vielen Punktprodukten bestehen.

Multiplizieren wir zum Beispiel eine 4 × 8-Matrix mit einem 8 × 1-Vektor. Bei der Berechnung werden die einzelnen Werte im Ausgabevektor mit 4-Punkt-Produkten berechnet. A, B, C und D.

Diagramm eines Beispiels mit einer Matrixvektor-Multiplikation

Der Prozess zur Berechnung jeder dieser Ausgaben ist gleich: sehen wir uns die Schritte zur Berechnung eines von ihnen an. Vor jeder Berechnung müssen die 8-Bit-Ganzzahldaten in einen Typ konvertiert werden, mit dem wir eine Arithmetik durchführen können, z. B. f16. Dann führen wir eine elementweise Multiplikation aus und addieren schließlich alle Produkte. Insgesamt führen wir für die gesamte Matrixvektormultiplikation 40 Ganzzahl-/Gleitkommazahlenkonvertierungen zum Entpacken der Daten, 32 Gleitkomma-Multiplikationen und 28 Gleitkomma-Additionen durch.

Bei größeren Matrizen mit mehr Operationen können gepackte Ganzzahlpunktprodukte den Arbeitsaufwand reduzieren.

Für jede der Ausgaben im Ergebnisvektor führen wir zwei gepackte Punktproduktvorgänge mit der in WebGPU Shading Language integrierten dot4U8Packed durch und addieren dann die Ergebnisse. Insgesamt wird für die gesamte Matrixvektormultiplikation keine Datenkonvertierung durchgeführt. Wir führen 8 gepackte Punktprodukte und 4 ganzzahlige Additionen aus.

Diagramm eines gepackten Beispiels mit einer ganzzahligen Matrix-Vektor-Multiplikation

Wir haben Integer-Punkt-Produkte mit 8-Bit-Daten auf einer Vielzahl von Verbraucher-GPUs getestet. Im Vergleich zu einem 16-Bit-Gleitkomma ist 8-Bit 1,6- bis 2,8-mal schneller. Wenn wir zusätzlich gepackte Integer-Punkt-Produkte verwenden, ist die Leistung noch besser. Das ist 1,7- bis 2,9-mal schneller.

<ph type="x-smartling-placeholder">
</ph> Screenshot der Matrix-Vektor-Multiplikationsbeschleunigung: f16 im Vergleich zu u8 <ph type="x-smartling-placeholder">
</ph> Diagramm 1: Beschleunigung des Matrixvektors, Vergleich von f16 mit U8 und U8 mit Punkt4U8Packed.

Prüfe mit dem Attribut wgslLanguageFeatures, ob Browser unterstützt werden. Wenn die GPU gepackte Punktprodukte nativ nicht unterstützt, führt der Browser eine eigene Implementierung aus.

// main.js

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

Der folgende Unterschied im Code-Snippet (Unterschied) hebt die Änderungen hervor, die für die Verwendung gepackter Ganzzahlprodukte in einem WebGPU-Shader erforderlich sind.

Vor – Ein WebGPU-Shader, der partielle Punktprodukte in die Variable „sum“ sammelt. Am Ende der Schleife enthält `sum` das vollständige Punktprodukt 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);
  }
}

Nachher: Ein WebGPU-Shader, der für gepackte Integer-Punktprodukte geschrieben ist. Der Hauptunterschied besteht darin, dass dieser Shader statt 4 Gleitkommawerte aus dem Vektor und der Matrix eine einzelne 32-Bit-Ganzzahl lädt. Diese 32-Bit-Ganzzahl enthält die Daten von vier ganzzahligen 8-Bit-Werten. Dann 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 gepackte Integer-Punkt-Produkte sind in Chrome enthalten, die KI und ML beschleunigen. 16-Bit-Gleitkommazahlen sind verfügbar, wenn die Hardware dies unterstützt. Chrome implementiert gepackte Integer-Punkt-Produkte auf allen Geräten.

Sie können diese Funktionen ab sofort in der stabilen Chrome-Version verwenden, um eine bessere Leistung zu erzielen.

Vorgeschlagene Funktionen

In Zukunft untersuchen wir zwei weitere Merkmale: Untergruppen und kooperative Matrixmultiplikation.

Die Funktion für Untergruppen ermöglicht die Parallelität der SIMD-Ebene, um zu kommunizieren oder kollektive mathematische Operationen durchzuführen, z. B. eine Summe für mehr als 16 Zahlen. Dies ermöglicht eine effiziente domainübergreifende Datenfreigabe. Untergruppen werden in modernen GPUs APIs mit unterschiedlichen Namen und in leicht unterschiedlichen Formen unterstützt.

Wir haben das gemeinsame Set in einem Vorschlag zusammengefasst, den wir der WebGPU-Standardisierungsgruppe vorgelegt haben. Außerdem haben wir Prototyp-Untergruppen in Chrome hinter einem experimentellen Flag und haben unsere ersten Ergebnisse in die Diskussion eingebracht. Das Hauptproblem besteht darin, wie die Portabilität sichergestellt werden kann.

Die kooperative Matrixmultiplikation ist eine neuere Ergänzung zu GPUs. Eine große Matrixmultiplikation kann in mehrere kleinere Matrixmultiplikationen aufgeteilt werden. Bei der kooperativen Matrixmultiplikation werden diese kleineren Blöcke fester Größe in einem einzigen logischen Schritt multipliziert. Innerhalb dieses Schritts arbeiten eine Gruppe von Threads effizient zusammen, um das Ergebnis zu berechnen.

Wir haben die Unterstützung für zugrunde liegende GPU-APIs befragt und planen, der WebGPU-Standardisierungsgruppe einen Vorschlag vorzuschlagen. Wie bei Untergruppen gehen wir davon aus, dass sich ein Großteil der Diskussion um die Übertragbarkeit dreht.

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

Wir haben Untergruppen in GPU-Kerneln der Prefill-Phase des Large Language Model verwendet, daher melde ich nur die Beschleunigung für die Prefill-Phase. Auf einer Intel-GPU sehen wir, dass Untergruppen zweieinhalb Mal schneller arbeiten als die Baseline. Diese Verbesserungen sind jedoch nicht über verschiedene GPUs hinweg konsistent.

<ph type="x-smartling-placeholder">
</ph> Screenshot der Beschleunigung von Untergruppen in der MediaPipe-LLM-Inferenz <ph type="x-smartling-placeholder">
</ph> Diagramm 2. Mit Untergruppen wird die Vorabausführung auf der Intel Tiger Lake GT2-GPU um das 2, 5-Fache beschleunigt und bietet experimentelle Unterstützung in Chrome und Mediapipe.

Das nächste Diagramm zeigt die Ergebnisse der Anwendung von Untergruppen, um eine Mikro-Benchmark mit Matrixmultiplikation über mehrere Verbraucher-GPUs hinweg zu optimieren. Die Matrixmultiplikation ist einer der schwereren Operationen in Large Language Models. Die Daten zeigen, dass Untergruppen auf vielen GPUs die Geschwindigkeit um das Zwei-, Fünf- und sogar Dreizehnfache der Basisgeschwindigkeit erhöhen. Beachten Sie jedoch, dass Untergruppen auf der ersten GPU überhaupt nicht viel besser sind.

<ph type="x-smartling-placeholder">
</ph> Screenshot der Untergruppenbeschleunigung für die Matrixmultiplikation <ph type="x-smartling-placeholder">
</ph> Diagramm 3: Die Anwendung von Untergruppen für die Matrixmultiplikation kann die Leistung weiter erhöhen.

GPU-Optimierung ist schwierig

Die beste Möglichkeit zur Optimierung Ihrer GPU hängt letztendlich davon ab, welche GPU der Client bietet. Die Verwendung innovativer neuer GPU-Features zahlt sich nicht immer so aus, wie Sie es erwarten, da viele komplexe Faktoren beteiligt sein 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 Rechen-Threads der GPU vollständig nutzen.

Auch Muster für den Arbeitsspeicherzugriff können wichtig sein. GPUs funktionieren in der Regel weitaus besser, wenn die Compute-Threads in einem für die Hardware optimalen Muster auf den Speicher zugreifen. Wichtig: Je nach GPU-Hardware sollten Sie unterschiedliche Leistungsmerkmale erwarten. Je nach GPU müssen Sie möglicherweise verschiedene Optimierungen ausführen.

In der folgenden Tabelle verwenden wir denselben Matrixmultiplikationsalgorithmus, aber wir haben eine weitere Dimension hinzugefügt, um die Auswirkungen verschiedener Optimierungsstrategien sowie die Komplexität und Varianz der verschiedenen GPUs weiter zu veranschaulichen. Wir haben hier eine neue Technik eingeführt, die wir „Swizzle“ nennen. Swizzle optimiert die Zugriffsmuster des Arbeitsspeichers für die Hardware.

Wie Sie sehen, hat der Erinnerungseffekt eine große Wirkung: manchmal sogar noch wirkungsvoller als Untergruppen. Bei GPU 6 bietet Swizzle eine 12-fache Beschleunigung, während Untergruppen eine 13-fache Beschleunigung ermöglichen. Insgesamt ist die Geschwindigkeit um das 26-Fache gestiegen. Bei anderen GPUs erzielen Swizzle- und Untergruppen mitunter eine bessere Leistung als jede einzelne Gruppe. Bei anderen GPUs erzielen Sie mit der ausschließlichen Verwendung von Swizzle die beste Leistung.

<ph type="x-smartling-placeholder">
</ph> Screenshot der Beschleunigung für Matrixmultiplikationsstrategien <ph type="x-smartling-placeholder">
</ph> Diagramm 4:

Die Feinabstimmung und Optimierung von GPU-Algorithmen, damit sie auf jeder Hardware gut funktionieren, kann viel Fachwissen erfordern. Zum Glück gibt es jedoch eine enorme Menge an talentierten Arbeiten, die in höherstufige Bibliotheken-Frameworks wie Mediapipe, Transformers.js, Apache TVM und ONNX Runtime Web einfließen.

Bibliotheken und Frameworks sind gut auf die Verwaltung verschiedener GPU-Architekturen und das Generieren von plattformspezifischem Code ausgerichtet, der auf dem Client einwandfrei ausgeführt werden kann.

Fazit

Das Chrome-Team arbeitet weiterhin an der Weiterentwicklung der WebAssembly- und WebGPU-Standards, um die Webplattform für ML-Arbeitslasten zu verbessern. Wir investieren in schnellere Computing-Primitive, bessere Interoperabilität über Webstandards hinweg und sorgen dafür, dass große und kleine Modelle effizient auf verschiedenen Geräten 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 Übertragbarkeit. Und das machen wir nicht allein. Wir arbeiten mit den anderen Browseranbietern bei W3C und vielen Entwicklungspartnern zusammen.

Wir hoffen, dass Sie sich bei der Arbeit mit WebAssembly und WebGPU an Folgendes erinnern:

  • KI-Inferenzen sind jetzt geräteübergreifend im Web verfügbar. Dies bietet den Vorteil der Ausführung auf Clientgeräten, z. B. geringere Serverkosten, niedrige Latenz und erhöhter Datenschutz.
  • Während viele der besprochenen Funktionen in erster Linie für die Framework-Autoren relevant sind, können Ihre Anwendungen ohne großen Aufwand davon profitieren.
  • Die Webstandards sind fließend und entwickeln sich weiter und wir freuen uns immer über Feedback. Teilen Sie Ihre für WebAssembly und WebGPU.

Danksagungen

Wir möchten uns bei dem Intel-Webgrafikteam bedanken, das maßgeblich an der Entwicklung der WebGPU f16 mitgewirkt hat und Integer-Punkt-Produktfunktionen packt. Wir möchten uns bei allen Mitgliedern der Arbeitsgruppen WebAssembly und WebGPU bei W3C, einschließlich der anderen Browseranbieter, bedanken.

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