Web AI を高速化するための WebAssembly と WebGPU の機能強化、パート 2

このドキュメントは、Web AI を高速化する WebAssembly と WebGPU の機能拡張（パート 1）の続きです。先に進む前に、この投稿を読むか、IO 24 の講演をご視聴ください。

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

WebGPU

WebGPU を使用すると、ウェブ アプリケーションがクライアントの GPU ハードウェアにアクセスして、効率的で並列計算を実行できるようになります。Chrome で WebGPU をリリースして以来、ウェブ上に AI と ML の素晴らしいデモが登場しました。

たとえば、Web Stable Diffusion は、AI を使用して、ブラウザ内で直接テキストから画像を生成できることを実証しました。今年、Google の Mediapipe チームは、大規模言語モデル推論の試験運用版サポートを発表しました。

次のアニメーションは、Google のオープンソースの大規模言語モデル（LLM）である Gemma が、すべてオンデバイスで Chrome でリアルタイムに実行されている様子を示しています。

Meta の Segment Anything Model の Hugging Face のデモでは、クライアント上で完全に高品質のオブジェクト マスクが生成されます。

これらは、AI と ML における WebGPU の力を示す素晴らしいプロジェクトのほんの数例です。WebGPU を使用することで、このようなモデルやその他のモデルは CPU よりもはるかに高速に実行できます。

Hugging Face のテキスト埋め込み用の WebGPU ベンチマークでは、同じモデルの CPU 実装と比較して、驚異的なスピードアップが実証されています。Apple M1 Max ノートパソコンでは、WebGPU が 30 倍以上高速になりました。WebGPU によってベンチマークが 120 倍以上高速化されたという報告もあります。

AI と ML 向けの WebGPU 機能の改善

WebGPU は、コンピューティング シェーダーをサポートしているため、何十億ものパラメータを持つことができる AI モデルや ML モデルに最適です。コンピューティング シェーダーは GPU 上で動作し、大量のデータに対して並列配列オペレーションを実行するのに役立ちます。

過去 1 年間に行われた WebGPU の多数の改善の中で、Google はウェブでの ML と AI のパフォーマンスを向上させる機能の追加を続けています。先日、16 ビット浮動小数点とパック整数のドット積という 2 つの新機能をリリースしました。

16 ビット浮動小数点

ML ワークロードに精度は必要ないことを覚えておいてください。shader-f16 は、WebGPU シェーディング言語で f16 タイプを使用できるようにする機能です。この浮動小数点型は、通常の 32 ビットではなく 16 ビットを使用します。f16 は範囲が小さく、精度は低くなりますが、多くの ML モデルではこれで十分です。

この機能により、次のような方法で効率が向上します。

  • メモリの削減: f16 要素を使用したテンソルはスペースの半分を占めるため、メモリ使用量が半分になります。GPU の計算はメモリ帯域幅のボトルネックになることが多いため、メモリが半分になると、シェーダーの実行速度が 2 倍になることがよくあります。技術的には、メモリ帯域幅を節約するために f16 は必要ありません。低精度形式でデータを保存し、それをシェーダーで f32 まで展開して計算することができます。しかし、GPU はデータのパッキングとアンパックに余分なコンピューティング能力を消費します。

  • データ変換の低減: f16 はデータ変換を最小限に抑えることで、使用するコンピューティングを減らします。精度の低いデータを保存して、変換なしで直接使用できる。

  • 並列処理の向上: 最新の GPU では、GPU の実行ユニットにより多くの値を同時に格納できるため、より多くの並列計算を実行できます。たとえば、1 秒あたり最大 5 兆の f32 浮動小数点演算をサポートする GPU は、1 秒あたり 10 兆個の f16 浮動小数点演算をサポートする可能性があります。

で確認できます。 <ph type="x-smartling-placeholder">
</ph> テキスト埋め込み用の WebGPU ベンチマークのスクリーンショット
Hugging Face のテキスト埋め込み用の WebGPU ベンチマークでは、shader-f16 を使用することで、Apple M1 Max ノートパソコンにおける f32 の 3 倍の速さで実行されています。

WebLLM は、複数の大規模言語モデルを実行できるプロジェクトです。オープンソースの ML コンパイラ フレームワークである Apache TVM を使用しています。

Llama 3 の 80 億パラメータ モデルを使用してパリへの旅行を計画するよう WebLLM に依頼しました。結果から、モデルのプレフィル段階では、f16 は f32 の 2.1 倍であることがわかりました。デコード フェーズでは、1.3 倍以上高速です。

アプリケーションは、まず GPU アダプターが f16 をサポートしていることを確認し、利用可能であれば、GPU デバイスをリクエストするときに明示的に有効にします。f16 がサポートされていない場合、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);
}

次に、WebGPU シェーダーで、上部で f16 を明示的に有効にする必要があります。その後は、他の 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];
  ...
}

パックされた整数のドット積

多くのモデルは、わずか 8 ビットの精度（f16 の半分）でうまく機能します。これは、セグメンテーションやオブジェクト認識のための LLM や画像モデルでよく使用されています。とはいえ、モデルの出力品質は精度が低くなるため、8 ビットの量子化はすべてのアプリケーションに適しているわけではありません。

8 ビット値をネイティブにサポートする GPU はそれほど多くありません。そこで必要になるのが、パックされた整数のドット積です。DP4a は Chrome 123 でリリースされました。

最新の GPU には、2 つの 32 ビット整数を取り、それらを 4 つの連続した 8 ビット整数として解釈し、コンポーネント間のドット積を計算する特別な命令があります。

行列乗算カーネルは多数のドット積で構成されているため、これは AI や ML で特に有用です。

たとえば、4 x 8 の行列と 8 x 1 のベクトルを乗算します。この計算では、4 つのドット積を使用して出力ベクトルの各値を計算します。A、B、C、D があります。

行列ベクトル乗算の例の図

それぞれの出力を計算するプロセスは同じです。そのうちの一つを計算するためのステップを見ていきます。計算を行う前に、まず 8 ビットの整数データを、演算が可能な型（f16 など）に変換する必要があります。次に、要素単位の乗算を実行し、最後にすべての積を加算します。行列ベクトルの乗算全体で、データを展開するために整数から浮動小数点数への変換を 40 回、浮動小数点数の乗算を 32 回、浮動小数点の加算を 28 回実行します。

演算回数が多い大規模な行列では、整数のドット積のパックを使用すると、処理量を削減できます。

結果ベクトルの出力ごとに、WebGPU シェーディング言語の組み込み dot4U8Packed を使用して 2 つのパックされたドット積演算を行い、結果を加算します。行列とベクトルの乗算全体では、全体としてデータ変換は行いません。8 つのパックされたドット積と 4 つの整数の加算を実行します。

パックされた整数行列ベクトル乗算の例の図

さまざまなコンシューマ GPU で、8 ビットデータによるパックされた整数ドット積をテストしました。16 ビットの浮動小数点数と比較すると、8 ビットは 1.6 ～ 2.8 倍高速であることがわかります。さらに、パックされた整数ドット積を使用すると、パフォーマンスはさらに向上します。1.7 ～ 2.9 倍高速です。

<ph type="x-smartling-placeholder">
</ph> 行列ベクトル乗算の高速化のスクリーンショット: f16 と u8 <ph type="x-smartling-placeholder">
</ph> グラフ 1: 行列ベクトルの高速化（dot4U8Packed の f16 を U8 および U8 と比較）。

wgslLanguageFeatures プロパティでブラウザのサポートを確認します。GPU がパックされたドット積をネイティブにサポートしていない場合、ブラウザは独自の実装をポリフィルします。

// main.js

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

次のコード スニペットの diff（差分）は、WebGPU シェーダーで、パックされた整数の製品を使用するために必要な変更を示しています。

変更前 - 部分的なドット積を変数「sum」に蓄積する WebGPU シェーダー。ループの最後で `sum` は、ベクトルと入力行列の 1 行の完全ドット積を保持します。

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

変更後 - パックされた整数のドット積を使用して記述される WebGPU シェーダー。主な違いは、ベクトルと行列から 4 つの浮動小数点値を読み込むのではなく、1 つの 32 ビット整数を読み込むことです。この 32 ビット整数は、4 つの 8 ビット整数値のデータを保持します。次に、dot4U8Packed を呼び出して 2 つの値のドット積を計算します。

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

Chrome には、16 ビットの浮動小数点数とパックされた整数ドット積の両方が、AI と ML を高速化する機能として標準装備されています。16 ビット浮動小数点はハードウェアでサポートされている場合に使用可能で、Chrome ではすべてのデバイスでパックされた整数ドット積が実装されます。

Chrome Stable で今すぐこれらの機能を使用して、パフォーマンスを向上させることができます。

提案された機能

今後、さらに 2 つの機能、サブグループと協調行列乗算について調査する予定です。

サブグループ機能を使用すると、SIMD レベルの並列処理で通信を行ったり、16 を超える数値の合計など、集団的な数学演算を実行したりできます。これにより、効率的なクロススレッド データ共有が可能になります。サブグループは最新の GPU API でサポートされていますが、名前や形式は若干異なります。

そこで、共通セットを精製して 1 つの提案書にまとめて、WebGPU 標準化グループと共有しました。また、実験的なフラグを使って Chrome でサブグループのプロトタイプを作成し、最初の結果を議論に盛り込みました。主な問題は、どのようにポータブルな動作を保証するかです。

協調行列乗算は、最近 GPU に追加された機能です。大きな行列の乗算は、複数の小さな行列の乗算に分割できます。協調行列乗算は、これらの小さな固定サイズのブロックに対して、単一の論理ステップで乗算を実行します。そのステップ内では、スレッドのグループが効率的に連携して結果を計算します。

基盤となる GPU API についてサポートに関するアンケートを実施し、WebGPU の標準化グループに提案を行う予定です。サブグループと同様、ポータビリティを中心に議論を進めることが予想されます。

サブグループ操作のパフォーマンスを評価するため、実際のアプリケーションでサブグループの試験運用版のサポートを MediaPipe に統合し、サブグループ操作の Chrome プロトタイプでテストしました。

大規模言語モデルの事前入力フェーズで GPU カーネルのサブグループを使用したため、プレフィル フェーズの高速化についてのみ報告します。Intel GPU では、サブグループのパフォーマンスがベースラインの 2.5 倍になることがわかりました。ただし、これらの改善は異なる GPU 間で一貫していません。

<ph type="x-smartling-placeholder">
</ph> MediaPipe LLM 推論でのサブグループの高速化のスクリーンショット <ph type="x-smartling-placeholder">
</ph> グラフ 2.Intel Tiger Lake GT2 GPU でサブグループを使用すると、事前入力処理が 2.5 倍速くなり、Chrome と Mediapipe で試験運用版がサポートされます。

次のグラフは、サブグループを適用して、複数のコンシューマ GPU にわたって行列乗算マイクロベンチマークを最適化した結果を示しています。行列の乗算は、大規模言語モデルにおける負荷の高いオペレーションの一つです。データによると、多くの GPU で、サブグループによって速度がベースラインの 2 倍、5 倍、さらには 13 倍になっていることがわかります。ただし、最初の GPU では、サブグループはそれほど改善されていません。

<ph type="x-smartling-placeholder">
</ph> 行列乗算のサブグループの高速化のスクリーンショット <ph type="x-smartling-placeholder">
</ph> グラフ 3.行列の乗算にサブグループを適用すると、パフォーマンスがさらに向上する可能性があります。

GPU の最適化が難しい

結局のところ、GPU を最適化する最善の方法は、クライアントが提供している GPU によって異なります。多くの複雑な要素が関係する可能性があるため、高度な新しい GPU 機能を使用しても、期待どおりの結果が得られるとは限りません。ある GPU で最適な最適化戦略が、別の GPU では最適ではない場合があります。

GPU のコンピューティング スレッドを最大限に活用しながら、メモリ帯域幅を最小限に抑えるには、

メモリアクセス パターンも非常に重要です。コンピューティング スレッドがハードウェアに最適なパターンでメモリにアクセスすると、GPU のパフォーマンスが大幅に向上する傾向があります。 重要: パフォーマンス特性は GPU ハードウェアによって異なることが想定されます。GPU に応じて異なる最適化の実行が必要になる場合があります。

次の図では、同じ行列乗算アルゴリズムを使用していますが、さまざまな最適化戦略の影響、異なる GPU 間の複雑さとばらつきをさらに示すため、別のディメンションを追加しました。ここでは、「Swizzle」という新しい手法を導入しました。Swizzle は、ハードウェアにより最適化されるようメモリアクセス パターンを最適化します。

メモリの実装入れ替えが重大な影響を及ぼしていることがわかります。サブグループよりもインパクトが強くなることもあります。GPU 6 では、swizzle で 12 倍の高速化、サブグループによる 13 倍の高速化を実現しています。組み合わせると、26 倍もの速度が上がります。他の GPU では、入れ替えとサブグループを組み合わせた方が、どちらかを単独で使用するよりもパフォーマンスが向上する場合があります。他の GPU では、入れ替えのみが最も効果的です。

<ph type="x-smartling-placeholder">
</ph> 行列乗算戦略の高速化のスクリーンショット <ph type="x-smartling-placeholder">
</ph> グラフ 4.

あらゆるハードウェアで適切に動作するように GPU アルゴリズムを調整および最適化するには、多くの専門知識が必要です。ありがたいことに、Mediapipe、Transformers.js、Apache TVM、ONNX Runtime Web などの上位レベルのライブラリ フレームワークには、膨大な量の才能が注ぎ込まれています。

ライブラリとフレームワークは、多様な GPU アーキテクチャの管理や、クライアント上で適切に実行されるプラットフォーム固有のコードの生成に伴う複雑さに対処できる体制が整っています。

要点

Chrome チームは、WebAssembly と WebGPU の標準を進化させ、ML ワークロードに適したウェブ プラットフォームを改善し続けるための取り組みを続けています。Google は、コンピューティング プリミティブを高速化し、ウェブ標準間の相互運用性を改善し、大小のモデルがデバイス間で効率的に実行されるようにすることに投資しています。

Google の目標は、ウェブの強みであるリーチ、ユーザビリティ、ポータビリティを維持しながら、プラットフォームの機能を最大化することです。これは Google だけの作業ではありません。Google は、W3C の他のブラウザ ベンダーや多くの開発パートナーと協力しています。

WebAssembly と WebGPU を使用する際は、次の点に留意してください。

  • AI 推論は、ウェブでも複数のデバイスで利用できるようになりました。これにより、サーバー費用の削減、低レイテンシ、プライバシー保護の強化など、クライアント デバイスで稼働するメリットが得られます。
  • ここで説明する機能の多くは主にフレームワークの作成者に関連していますが、アプリケーションは多くのオーバーヘッドなしにメリットを得ることができます。
  • ウェブの標準は流動的で進化し続けているため、Google は常にフィードバックを求めています。WebAssembly と WebGPU のグラフを共有します。

謝辞

WebGPU f16 と整数ドット製品機能の開発に貢献してくれた Intel ウェブ グラフィック チームに感謝の意を表します。他のブラウザ ベンダーを含め、W3C の WebAssembly および WebGPU ワーキング グループのメンバーの皆様に感謝いたします。

Google とオープンソース コミュニティの AI チームと ML チームの皆さまのご協力に感謝いたします。そしてもちろん、これらすべてを可能にしているすべてのチームメイトです。