Penyempurnaan WebAssembly dan WebGPU untuk Web AI yang lebih cepat, bagian 2

Dokumen ini merupakan kelanjutan dari peningkatan WebAssembly dan WebGPU untuk mempercepat Web AI, bagian 1. Sebaiknya baca postingan ini atau tonton pembahasan di IO 24 sebelum melanjutkan.

WebGPU

WebGPU memberi aplikasi web akses ke hardware GPU klien untuk melakukan komputasi yang efisien dan sangat paralel. Sejak meluncurkan WebGPU di Chrome, kami telah melihat demo yang luar biasa terkait kecerdasan buatan (AI) dan machine learning (ML) di web.

Misalnya, Web Stable Diffusion menunjukkan bahwa AI dapat digunakan untuk membuat gambar dari teks, langsung di browser. Pada awal tahun ini, tim Mediapipe Google memublikasikan dukungan eksperimental untuk inferensi model bahasa besar.

Animasi berikut menunjukkan Gemma, model bahasa besar (LLM) open source dari Google, yang berjalan sepenuhnya di perangkat di Chrome, secara real time.

Demo Hugging Face berikut untuk Model Anything Segment Meta menghasilkan mask objek berkualitas tinggi sepenuhnya di klien.

Ini hanyalah beberapa project luar biasa yang menunjukkan kecanggihan WebGPU untuk AI dan ML. WebGPU memungkinkan model ini dan yang lainnya berjalan jauh lebih cepat daripada yang dapat dilakukan pada CPU.

Tolok ukur WebGPU untuk penyematan teks dari Hugging Face menunjukkan percepatan yang luar biasa dibandingkan dengan penerapan CPU dari model yang sama. Di laptop Apple M1 Max, WebGPU berjalan 30 kali lebih cepat. Sejumlah penyedia lain melaporkan bahwa WebGPU mempercepat tolok ukur lebih dari 120 kali.

Meningkatkan fitur WebGPU untuk AI dan ML

WebGPU sangat cocok untuk model AI dan ML, yang dapat memiliki miliaran parameter, berkat dukungan untuk shader komputasi. Shader komputasi berjalan di GPU dan membantu menjalankan operasi array paralel pada volume data yang besar.

Di antara berbagai peningkatan pada WebGPU dalam setahun terakhir, kami terus menambahkan lebih banyak kemampuan untuk meningkatkan performa ML dan AI di web. Baru-baru ini, kami meluncurkan dua fitur baru: produk floating point 16 bit dan produk dot bilangan bulat yang dipaketkan.

Floating point 16 bit

Ingat, beban kerja ML tidak memerlukan presisi. shader-f16 adalah fitur yang memungkinkan penggunaan jenis f16 dalam bahasa shading WebGPU. Jenis floating point ini memakan 16 bit, bukan 32 bit yang biasa. f16 memiliki rentang yang lebih kecil dan kurang presisi, tetapi untuk banyak model ML, ini sudah cukup.

Fitur ini meningkatkan efisiensi dalam beberapa cara:

• Memori lebih rendah: Tensor dengan elemen f16 menggunakan setengah ruang, yang mengurangi penggunaan memori setengahnya. Komputasi GPU sering kali mengalami bottleneck pada bandwidth memori, sehingga separuh memori tersebut sering kali berarti shader berjalan dua kali lebih cepat. Secara teknis, Anda tidak memerlukan f16 untuk menghemat bandwidth memori. Data dapat disimpan dalam format presisi rendah, lalu memperluasnya ke f32 penuh dalam shader untuk komputasi. Namun, GPU menghabiskan daya komputasi tambahan untuk mengemas dan membongkar data.

• Pengurangan konversi data: f16 menggunakan lebih sedikit komputasi dengan meminimalkan konversi data. Data presisi rendah dapat disimpan, lalu digunakan langsung tanpa konversi.

• Peningkatan paralelisme: GPU modern dapat memasukkan lebih banyak nilai secara bersamaan di unit eksekusi GPU, sehingga dapat melakukan lebih banyak komputasi paralel. Misalnya, GPU yang mendukung hingga 5 triliun operasi floating point f32 per detik dapat mendukung 10 triliun operasi floating point f16 per detik.

WebLLM adalah project yang dapat menjalankan beberapa model bahasa besar. Game ini menggunakan Apache TVM, framework compiler machine learning open source.

Saya meminta WebLLM untuk merencanakan perjalanan ke Paris, menggunakan model parameter Llama 3 delapan miliar. Hasilnya menunjukkan bahwa selama fase {i>prefill<i} dari model, f16 adalah 2,1 kali lebih cepat daripada f32. Selama fase dekode, prosesnya 1,3 kali lebih cepat.

Aplikasi harus mengonfirmasi terlebih dahulu bahwa adaptor GPU mendukung f16, dan jika tersedia, aktifkan secara eksplisit saat meminta perangkat GPU. Jika f16 tidak didukung, Anda tidak dapat memintanya dalam 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);
}

Kemudian, dalam shader WebGPU, Anda harus mengaktifkan f16 secara eksplisit di bagian atas. Setelah itu, Anda bebas menggunakannya dalam shader seperti jenis data float lainnya.

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

Produk bilangan bulat titik yang dikemas

Banyak model masih bekerja dengan baik hanya dengan presisi 8 bit (setengah dari f16). Cara ini populer di kalangan LLM dan model gambar untuk segmentasi dan pengenalan objek. Meskipun demikian, kualitas output untuk model menurun dengan presisi yang lebih rendah, sehingga kuantisasi 8-bit tidak cocok untuk setiap aplikasi.

Hanya sedikit GPU yang secara native mendukung nilai 8 bit. Di sinilah produk bilangan bulat yang dikemas dapat berperan. Kami meluncurkan DP4a di Chrome 123.

GPU modern memiliki instruksi khusus untuk mengambil dua bilangan bulat 32-bit, menafsirkannya masing-masing sebagai 4 bilangan bulat 8-bit yang dikemas secara berurutan, dan menghitung produk dot di antara komponennya.

Fungsi ini sangat berguna untuk AI dan machine learning karena kernel perkalian matriks terdiri dari banyak produk titik.

Sebagai contoh, mari kita kalikan matriks 4 x 8 dengan vektor 8 x 1. Komputasi ini melibatkan pengambilan 4 titik perkalian untuk menghitung masing-masing nilai dalam vektor {i>output<i}; A, B, C, dan D.

Proses untuk menghitung setiap {i>output<i} ini sama; kita akan melihat langkah-langkah yang terlibat dalam menghitung salah satunya. Sebelum komputasi apa pun, pertama-tama kita harus mengonversi data integer 8-bit ke tipe yang dapat kita gunakan untuk melakukan aritmatika, seperti f16. Kemudian, kita menjalankan perkalian berbasis elemen dan terakhir, menjumlahkan semua produk. Secara total, untuk seluruh perkalian matriks-vektor, kita melakukan 40 bilangan bulat untuk konversi mengambang untuk membongkar data, 32 perkalian {i>float<i}, dan 28 penambahan {i>float<i}.

Untuk matriks yang lebih besar dengan lebih banyak operasi, produk titik bilangan bulat yang dikemas dapat membantu mengurangi jumlah pekerjaan.

Untuk setiap output dalam vektor hasil, kami melakukan dua operasi produk titik yang dikemas menggunakan dot4U8Packed WebGPU Shading Language bawaan, lalu menambahkan hasilnya bersama-sama. Secara total, untuk seluruh perkalian matriks-vektor, kita tidak melakukan konversi data apa pun. Kita menjalankan 8 produk packed dot dan 4 penambahan bilangan bulat.

Kami menguji produk dot integer yang dikemas dengan data 8-bit pada berbagai GPU konsumen. Dibandingkan dengan floating point 16-bit, kita dapat melihat bahwa 8-bit 1,6 hingga 2,8 kali lebih cepat. Jika kami juga menggunakan produk dot bilangan bulat yang dipaketkan, performanya akan lebih baik. 1,7 hingga 2,9 kali lebih cepat.

Periksa dukungan browser dengan properti wgslLanguageFeatures. Jika GPU tidak mendukung produk dot paket secara native, browser akan mem-polyfill implementasinya sendiri.

// main.js

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

Perbedaan cuplikan kode berikut menunjukkan perubahan yang diperlukan untuk menggunakan produk bilangan bulat yang dipaketkan dalam shader WebGPU.

Sebelum — Shader WebGPU yang mengakumulasi produk titik parsial ke dalam variabel `sum`. Di akhir loop, 'sum' menyimpan perkalian titik penuh antara vektor dan satu baris matriks 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);
  }
}

Setelah — Shader WebGPU yang ditulis untuk menggunakan produk titik bilangan bulat yang dipaketkan. Perbedaan utamanya adalah, bukannya memuat 4 nilai float dari vektor dan matriks, shader ini memuat satu integer 32-bit. Bilangan bulat 32-bit ini menyimpan data empat nilai bilangan bulat 8-bit. Lalu, kita memanggil dot4U8Packed untuk menghitung produk titik dari dua nilai.

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

Produk floating point 16-bit dan produk dot bilangan bulat yang dipaketkan adalah fitur yang dikirimkan di Chrome yang mempercepat AI dan ML. Floating point 16-bit tersedia jika hardware mendukungnya, dan Chrome menerapkan produk titik bilangan bulat yang dipaketkan di semua perangkat.

Anda dapat menggunakan fitur ini di Chrome Stabil sekarang untuk mencapai performa yang lebih baik.

Fitur yang diusulkan

Ke depannya, kami akan menyelidiki dua fitur lainnya: subgrup dan perkalian matriks kooperatif.

Fitur subgrup memungkinkan paralelisme tingkat SIMD untuk berkomunikasi atau melakukan operasi matematika kolektif, misalnya menjumlahkan lebih dari 16 angka. Hal ini memungkinkan pembagian data lintas thread yang efisien. Subgrup didukung di API GPU modern, dengan nama yang bervariasi dan dalam bentuk yang sedikit berbeda.

Kami telah menyaring set umum menjadi proposal yang telah kami masukkan ke grup standardisasi WebGPU. Selain itu, kami telah membuat prototipe subgrup di Chrome di balik tanda eksperimental, dan telah membawa hasil awal kami ke dalam diskusi. Masalah utamanya adalah bagaimana memastikan perilaku portabel.

Perkalian matriks kooperatif adalah penambahan yang lebih baru untuk GPU. Perkalian matriks besar dapat dipecah menjadi beberapa perkalian matriks yang lebih kecil. Perkalian matriks kooperatif melakukan perkalian pada blok berukuran tetap yang lebih kecil ini dalam satu langkah logis. Dalam langkah tersebut, sekelompok thread bekerja sama secara efisien untuk menghitung hasilnya.

Kami melakukan survei dukungan pada API GPU dasar, dan berencana untuk mengajukan proposal kepada grup standardisasi WebGPU. Seperti pada subkelompok, kami berharap sebagian besar diskusi akan berpusat pada portabilitas.

Untuk mengevaluasi performa operasi subgrup, dalam aplikasi nyata, kami mengintegrasikan dukungan eksperimental untuk subgrup ke dalam MediaPipe dan mengujinya dengan prototipe Chrome untuk operasi subgrup.

Kami menggunakan subgrup dalam kernel GPU pada fase pengisian otomatis dari model bahasa besar, jadi saya hanya melaporkan percepatan untuk fase pengisian otomatis. Di GPU Intel, kita melihat bahwa subgrup berjalan dua setengah kali lebih cepat dari dasar pengukuran. Namun, peningkatan ini tidak konsisten di berbagai GPU.

Diagram berikutnya menunjukkan hasil penerapan subgrup untuk mengoptimalkan matriks perkalian microbenchmark di beberapa GPU konsumen. Perkalian matriks adalah salah satu operasi yang lebih berat dalam model bahasa besar. Data menunjukkan bahwa di banyak GPU, sub-kelompok meningkatkan kecepatan dua, lima, dan bahkan tiga belas kali lipat dari dasar pengukuran. Namun, perhatikan bahwa di GPU pertama, subgrup sama sekali tidak lebih baik.

Pengoptimalan GPU sulit

Pada akhirnya, cara terbaik untuk mengoptimalkan GPU Anda bergantung pada GPU yang ditawarkan klien. Menggunakan fitur GPU baru yang canggih tidak selalu memberikan hasil seperti yang Anda harapkan, karena mungkin ada banyak faktor kompleks yang terlibat. Strategi pengoptimalan terbaik di satu GPU mungkin bukan strategi terbaik di GPU lain.

Anda ingin meminimalkan bandwidth memori, sembari menggunakan sepenuhnya thread komputasi GPU.

Pola akses memori juga bisa menjadi sangat penting. GPU cenderung berperforma jauh lebih baik saat thread komputasi mengakses memori dengan pola yang optimal untuk hardware. Penting: Anda akan melihat karakteristik performa yang berbeda pada hardware GPU yang berbeda. Anda mungkin perlu menjalankan pengoptimalan yang berbeda, bergantung pada GPU.

Dalam diagram berikut, kami telah mengambil algoritma perkalian matriks yang sama, tetapi menambahkan dimensi lain untuk menunjukkan lebih lanjut dampak dari berbagai strategi pengoptimalan, serta kompleksitas dan varians di berbagai GPU. Kami telah memperkenalkan teknik baru di sini, yang kami sebut "Swizzle." Swizzle mengoptimalkan pola akses memori agar lebih optimal untuk hardware.

Anda bisa melihat bahwa {i>memory swizzle<i} memiliki dampak yang signifikan; kadang-kadang bahkan lebih berdampak daripada sub-kelompok. Pada GPU 6, swizzle menyediakan percepatan 12x, sedangkan subkelompok memberikan percepatan 13x. Jika digabungkan, keduanya memiliki kecepatan 26x yang luar biasa. Untuk GPU lain, terkadang swizzle dan subgrup yang digabungkan akan berperforma lebih baik daripada hanya salah satu saja. Dan di GPU lain, penggunaan swizzle secara eksklusif memberikan performa terbaik.

Menyesuaikan dan mengoptimalkan algoritma GPU agar berfungsi dengan baik di setiap hardware, memerlukan banyak keahlian. Namun, untungnya ada banyak karya berbakat yang masuk ke dalam framework library dengan level yang lebih tinggi, seperti Mediapipe, Transformers.js, Apache TVM, ONNX Runtime Web, dan banyak lagi.

Library dan framework diposisikan dengan baik untuk menangani kompleksitas pengelolaan beragam arsitektur GPU, dan menghasilkan kode khusus platform yang akan berjalan dengan baik pada klien.

Poin-poin penting

Tim Chrome terus membantu mengembangkan standar WebAssembly dan WebGPU guna meningkatkan platform web untuk workload machine learning. Kami berinvestasi pada primitif komputasi yang lebih cepat, interop yang lebih baik di seluruh standar web, dan memastikan bahwa model, baik besar maupun kecil, dapat berjalan secara efisien di berbagai perangkat.

Tujuan kami adalah memaksimalkan kemampuan platform ini sekaligus mempertahankan yang terbaik dari web: jangkauan, kegunaan, dan portabilitas. Dan kita tidak melakukan ini sendirian. Kami bekerja sama dengan vendor browser lainnya di W3C, dan banyak mitra pengembangan.

Kami harap Anda mengingat hal-hal berikut, saat Anda bekerja dengan WebAssembly dan WebGPU:

• Inferensi AI kini tersedia di web, di berbagai perangkat. Hal ini memberikan keuntungan jika berjalan di perangkat klien, seperti pengurangan biaya server, latensi rendah, dan peningkatan privasi.
• Meskipun banyak fitur yang dibahas relevan terutama bagi penulis framework, aplikasi Anda dapat memanfaatkan banyak hal tanpa biaya tambahan.
• Standar web berubah-ubah, dan berkembang, dan kami selalu mencari masukan. Bagikan kartu Anda untuk WebAssembly dan WebGPU.

Ucapan terima kasih

Kami ingin berterima kasih kepada tim grafis web Intel, yang berperan penting dalam mendorong WebGPU f16 dan mengemas fitur produk integer dot. Kami ingin berterima kasih kepada anggota grup kerja WebAssembly dan WebGPU lainnya di W3C, termasuk vendor browser lainnya.

Terima kasih kepada tim AI dan ML di Google dan di komunitas open source karena telah menjadi partner yang luar biasa. Dan tentu saja, semua rekan tim kami yang membuat semua ini mungkin terjadi.