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. Bevor Sie fortfahren, empfehlen wir Ihnen, diesen Beitrag zu lesen oder sich den Vortrag bei der I/24 anzusehen.

Austin Eng

Deepti Gandluri

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 von WebGPU in Chrome gibt es unglaubliche Demos von künstlicher Intelligenz (KI) und maschinellem Lernen (ML) im Web.

Web Stable Diffusion zeigte beispielsweise, dass es möglich ist, 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, ML-Arbeitslasten erfordern keine Präzision. 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 belegen die Hälfte des Speicherplatzes, 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.

</ph> <ph type="x-smartling-placeholder">

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 erfassen, 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.

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 integrierten WebGPU Shading Language (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.

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-Dot-Produkte verwenden, ist die Leistung noch besser. Das ist 1,7- bis 2,9-mal schneller.

</ph> <ph type="x-smartling-placeholder">

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> <ph type="x-smartling-placeholder">

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> <ph type="x-smartling-placeholder">

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.

Im folgenden Diagramm 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> <ph type="x-smartling-placeholder">

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.