Ulepszenia WebAssembly i WebGPU przyspieszające korzystanie z AI w internecie, część 2

Ten dokument jest kontynuacją ulepszeń WebAssembly i WebGPU przyspieszających działanie Web AI, część 1. Zanim przejdziesz dalej, przeczytaj ten post lub obejrzyj wykład na stronie IO 24.

WebGPU

WebGPU zapewnia aplikacjom internetowym dostęp do GPU klienta w celu wykonywania wydajnych, bardzo równoległych obliczeń. Od czasu wprowadzenia WebGPU w Chrome obserwujemy w internecie niesamowite wersje demonstracyjne sztucznej inteligencji (AI) i systemów uczących się.

Na przykład zespół Web Stable Diffusion pokazał, że można wykorzystać AI do generowania obrazów na podstawie tekstu bezpośrednio w przeglądarce. Na początku roku zespół Mediapipe Google opublikował eksperymentalną obsługę wnioskowania na podstawie dużego modelu językowego.

Animacja poniżej przedstawia Gemma, duży model językowy Google (LLM) typu open source, który działa w całości na urządzeniu w Chrome w czasie rzeczywistym.

Poniższa prezentacja Hugging Face modelu Segment dowolny model Meta tworzy wysokiej jakości maski obiektów całkowicie na kliencie.

To tylko kilka niezwykłych projektów, które pokazują możliwości WebGPU w zakresie AI i ML. WebGPU sprawia, że te i inne modele działają znacznie szybciej, niż mogłyby być w przypadku procesora.

Opracowana przez Hugging Face test porównawcza WebGPU w zakresie umieszczania tekstu potwierdza ogromne przyspieszenia w porównaniu z implementacją tego samego modelu przy użyciu procesora. Na laptopie Apple M1 Max procesor WebGPU był ponad 30 razy szybszy. Z innych źródeł wynika, że interfejs WebGPU przyspiesza test porównawczy ponad 120 razy.

Ulepszanie funkcji WebGPU w zakresie AI i ML

WebGPU świetnie sprawdza się w przypadku modeli AI i ML, które mogą mieć miliardy parametrów dzięki obsłudze cieniowania obliczeniowego. Moduły do cieniowania Compute działają z procesorem graficznym i ułatwiają wykonywanie operacji na tablicach równoległych na dużych ilościach danych.

Wśród wielu ulepszeń w WebGPU, które wprowadziliśmy w zeszłym roku, wciąż dodajemy kolejne funkcje zwiększające wydajność systemów uczących się i AI w internecie. Niedawno wprowadziliśmy 2 nowe funkcje: 16-bitową liczbę zmiennoprzecinkową i spakowaną liczbę całkowitą.

16-bitowa liczba zmiennoprzecinkowa

Pamiętaj, że zadania ML nie wymagają precyzji. shader-f16 to funkcja, która umożliwia korzystanie z typu f16 w języku cieniowania WebGPU. Ten typ zmiennoprzecinkowy zajmuje 16 bitów (a nie 32 bity). f16 ma mniejszy zakres i jest mniej precyzyjne, ale w przypadku wielu modeli ML jest to wystarczające.

Ta funkcja zwiększa wydajność na kilka sposobów:

• Mniejsza ilość pamięci: procesory Tensor z elementami f16 zajmują połowę miejsca, co zmniejsza zużycie pamięci o połowę. Obliczenia przy użyciu GPU często zmniejszają przepustowość pamięci, dlatego połowa pamięci może często powodować 2 razy szybsze działanie programów do cieniowania. Z technicznego punktu widzenia nie musisz mieć klawisza F16, aby oszczędzać przepustowość pamięci. Możesz przechowywać dane w formacie o niskiej dokładności, a potem rozszerzyć je do pełnego formatu f32 w cieniowaniu na potrzeby obliczeń. Proces ten wymaga jednak dodatkowej mocy obliczeniowej, by pakować i rozpakowywać dane.

• Mniejsza konwersja danych: f16 używa mniej mocy obliczeniowej dzięki minimalizacji konwersji danych. Dane o niskiej dokładności mogą być przechowywane, a potem używane bezpośrednio bez konwersji.

• Większa równoległość: nowoczesne układy GPU mogą pasować jednocześnie do większej liczby wartości w jednostkach wykonania GPU, co pozwala na wykonywanie większej liczby równoległych obliczeń. Na przykład GPU, które obsługuje do 5 bilionów operacji zmiennoprzecinkowych f32 na sekundę, może obsługiwać 10 bilionów operacji zmiennoprzecinkowych f16 na sekundę.

WebLLM to projekt, który może obsługiwać wiele dużych modeli językowych. Wykorzystuje Apache TVM – platformę open source do kompilowania systemów uczących się.

Poprosiłem WebLLM o zaplanowanie wycieczki do Paryża przy użyciu modelu Llama 3 zawierającego 8 miliardów parametrów. Wyniki pokazują, że podczas fazy wstępnego wypełniania modelu klawisz f16 jest 2,1 raza szybszy niż przyspieszenie f32. W fazie dekodowania jest to ponad 1, 3 raza szybsze.

Aplikacje muszą najpierw potwierdzić, że adapter GPU obsługuje kodowanie f16, a jeśli jest dostępny, musisz go włączyć bezpośrednio w żądaniu urządzenia GPU. Jeśli funkcja f16 nie jest obsługiwana, nie możesz zażądać jej w tablicy 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);
}

Następnie w cieniowaniu WebGPU musisz jawnie włączyć u góry klawisz f16. Potem możesz go używać w cieniowaniu tak samo jak każdego innego typu danych zmiennoprzecinkowych.

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

Spakowane iloczyny całkowite

Wiele modeli działa dobrze już przy 8-bitowej dokładności (połowa f16). Jest to popularne wśród LLM i modeli graficznych do segmentacji i rozpoznawania obiektów. Jakość danych wyjściowych modeli pogarsza się z mniejszą precyzją, więc kwantyzacja 8-bitowa nie jest odpowiednia dla każdego zastosowania.

Względnie niewiele GPU natywnie obsługuje wartości 8-bitowe. Tutaj wkraczają do świata produkty z spakowanymi liczbami całkowitymi. Wdrożyliśmy DP4a w Chrome 123.

Nowoczesne procesory graficzne mają specjalne instrukcje dotyczące pobierania dwóch 32-bitowych liczb całkowitych, interpretowania każdej z nich jako 4 składanych kolejno po sobie 8-bitowych liczb całkowitych i obliczania iloczynu skalarnego między ich komponentami.

Jest to szczególnie przydatne w przypadku sztucznej inteligencji i systemów uczących się, ponieważ jądra mnożenia macierzy składają się z wielu produktów punktowych.

Na przykład pomnóżmy macierz 4 x 8 przez wektor 8 x 1. Aby obliczyć tę wartość, trzeba użyć 4 iloczynów skalarnych do obliczenia każdej wartości w wektorze wyjściowym. A, B, C i D.

Proces obliczania każdego z tych wyników jest taki sam: Przyjrzyjmy się czynnościom niezbędnym do ich obliczenia. Przed obliczeniami musimy najpierw przekonwertować 8-bitowe liczby całkowite na typ, za pomocą którego można przeprowadzać obliczenia arytmetyczne, np. f16. Następnie przeprowadzamy mnożenie przez pierwiastki, a na koniec dodajesz wszystkie produkty do siebie. Łącznie dla całego mnożenia wektorów macierzy wykonujemy 40 konwersji liczb całkowitych na zmiennoprzecinkowe w celu rozpakowania danych, 32 mnożenia zmiennoprzecinkowe i 28 dodanych liczb zmiennoprzecinkowych.

W przypadku większych macierzy z większą liczbą operacji podane w formie iloczynu liczb całkowitych mogą pomóc zmniejszyć ilość pracy.

W przypadku każdego z wyjściowych wektorów wynikowych wykonujemy 2 operacje w formie iloczynu skalarnego, używając wbudowanego języka dot4U8Packed cieniowania WebGPU, a następnie sumujemy wyniki. Łącznie dla całego mnożenia wektorów macierzy nie przeprowadzamy żadnej konwersji danych. Realizujemy 8 produktów z kropkami i 4 dodane liczby całkowite.

Przetestowaliśmy 8-bitowe dane w postaci spakowanych liczb całkowitych przy użyciu różnych GPU. W porównaniu do 16-bitowej liczby zmiennoprzecinkowej widzimy, że 8-bitowa szybkość przetwarzania jest od 1,6 do 2,8 raza szybsza. Gdy dodatkowo korzystamy z produktów w formie spakowanych liczb całkowitych, wydajność jest jeszcze lepsza. Jest to od 1,7 do 2,9 raza szybsze.

Sprawdź, czy przeglądarki są obsługiwane przez właściwość wgslLanguageFeatures. Jeśli GPU nie obsługuje natywnie produktów spakowanych z kropkami, przeglądarka stosuje własną implementację.

// main.js

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

Poniższy fragment kodu różni się (różnica) i wyróżnia zmiany niezbędne do użycia produktów w postaci spakowanych liczb całkowitych w cieniu WebGPU.

Przed – cieniowanie WebGPU, które kumuluje częściowe produkty skalarne do zmiennej „sum”. Na końcu pętli „suma” określa pełny iloczyn skalarny między wektorem a jednym wierszem macierzy wejściowej.

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

Po – cieniowanie WebGPU napisane w celu użycia spakowanych liczb całkowitych z kropką. Główna różnica polega na tym, że zamiast 4 wartości zmiennoprzecinkowych z wektora i matrycy, ten cieniowanie wczytuje jedną 32-bitową liczbę całkowitą. Ta 32-bitowa liczba całkowita zawiera dane z 4 8-bitowych wartości całkowitych. Następnie wywołujemy funkcję dot4U8Packed, aby obliczyć iloczyn skalarny tych dwóch wartości.

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

Zarówno 16-bitowa liczba zmiennoprzecinkowa, jak i spakowane liczby całkowite są oferowane w Chrome, które przyspieszają działanie sztucznej inteligencji i systemów uczących się. 16-bitowa liczba zmiennoprzecinkowa jest dostępna, jeśli sprzęt ją obsługuje, a Chrome implementuje spakowane liczby całkowite z kropkami na wszystkich urządzeniach.

Obecnie możesz używać tych funkcji w wersji stabilnej Chrome, aby uzyskać lepszą wydajność.

Proponowane funkcje

Badamy jeszcze dwie funkcje: podgrupy i mnożenie macierzy kooperacyjnej.

Funkcja podgrup umożliwia komunikowanie się i wykonywanie zbiorczych operacji matematycznych na poziomie SIMD, na przykład przy użyciu sumy ponad 16 cyfr. Pozwala to na skuteczne udostępnianie danych w różnych wątkach. Podgrupy są obsługiwane przez nowoczesne interfejsy API GPU, różniące się nazwami i nieco inną formą.

Opracowaliśmy wspólny zestaw, który przekazaliśmy grupie standaryzacyjnej WebGPU. Opracowaliśmy pierwotne podgrupy w Chrome, opierając się na eksperymentalnej flagie, i właśnie do nich podeszliśmy w ramach dyskusji. Głównym problemem jest zapewnienie użytkownikom łatwości w działaniu aplikacji.

Mnożenie macierzy kooperacyjnej to nowszy dodatek do GPU. Duży mnożenie macierzy można podzielić na kilka mniejszych mnożeń macierzy. Funkcja mnożenia macierzy współpracy wykonuje mnożenie na tych mniejszych blokach o stałym rozmiarze w jednym kroku logicznym. Podczas wykonywania tego kroku grupa wątków skutecznie współpracuje przy obliczaniu wyniku.

Przeanalizowaliśmy obsługę podstawowych interfejsów API GPU i planujemy przedstawić ofertę grupie standaryzacyjnej WebGPU. Podobnie jak w przypadku podgrup, oczekujemy, że spora część dyskusji będzie się koncentrować na przenośności.

Aby ocenić wydajność operacji na podgrupach, w rzeczywistej aplikacji zintegrowaliśmy z MediaPipe eksperymentalną obsługę podgrup i przetestowaliśmy ją z prototypem Chrome do obsługi operacji w podgrupach.

Użyliśmy podgrup w jądkach GPU w fazie wstępnego wypełniania dużego modelu językowego, więc raportuję przyspieszenie tylko dla fazy wstępnego uzupełniania. W przypadku GPU firmy Intel zaobserwowaliśmy, że podgrupy działają 2,5 raza szybciej niż w przypadku punktu odniesienia. Te ulepszenia nie są jednak spójne w przypadku różnych GPU.

Następny wykres pokazuje wyniki zastosowania podgrup w celu optymalizacji wyniku mnożenia macierzy w mikroporównaniu z wieloma układami GPU konsumentów. Mnożenie macierzy to jedne z cięższych operacji w dużych modelach językowych. Dane pokazują, że w przypadku wielu układów GPU podgrupy przyspieszają szybkość o 2, 5, a nawet 13 razy względem wartości bazowej. Warto jednak zauważyć, że przy pierwszym procesorze graficznym podgrupy nie są dużo lepsze.

Optymalizacja GPU jest trudna

Ostatecznie najlepszy sposób optymalizacji GPU zależy od tego, jaki GPU oferuje klient. Korzystanie z nowych, wymyślnych funkcji GPU nie zawsze sprawdza się w oczekiwany sposób, ponieważ składa się to z wielu złożonych czynników. Najlepsza strategia optymalizacji w przypadku jednego GPU może nie być najlepsza w przypadku innego.

Chcesz zminimalizować przepustowość pamięci, jednocześnie w pełni wykorzystując wątki obliczeniowe procesora graficznego.

Wzorce dostępu do pamięci również mogą być bardzo ważne. Procesory graficzne działają znacznie lepiej, gdy wątki obliczeniowe korzystają z pamięci zgodnie ze wzorcem optymalnym dla sprzętu. Ważne: w zależności od sprzętu GPU możesz spodziewać się różnych parametrów wydajności. W zależności od GPU może być konieczne zastosowanie różnych optymalizacji.

Na wykresie poniżej zastosowaliśmy ten sam algorytm mnożenia macierzy, ale dodaliśmy kolejny wymiar, aby lepiej zademonstrować wpływ różnych strategii optymalizacji oraz złożoność i różnice między różnymi GPU. Wprowadziliśmy nową technikę o nazwie „Swizzle”. Swizzle optymalizuje wzorce dostępu do pamięci, by były bardziej optymalne pod kątem sprzętu.

Jak widać, wirowanie pamięci ma duży wpływ na jej funkcjonowanie. a czasem nawet bardziej niż podgrupy. Na GPU 6 gra Swizzle działa 12-krotnie, a podgrupy – 13-krotnie. Łącznie dają one niesamowite 26-krotne przyspieszenie. W przypadku innych układów GPU czasami połączenie pomysłu z podgrupami jest skuteczniejsze niż jedno z tych układów. W przypadku innych układów GPU najlepsze wyniki zapewnia używanie wyłącznie Swizzle.

Dostrajanie i optymalizowanie algorytmów GPU pod kątem poprawnego działania na każdym elemencie sprzętu może wymagać dużej wiedzy. Na szczęście czeka na nas mnóstwo utalentowanych pracowników w zakresie platform bibliotek wyższego poziomu, takich jak Mediapipe, Transformers.js, Apache TVM czy ONNX Runtime Web.

Biblioteki i platformy są odpowiednio skonfigurowane, aby radzić sobie ze złożonością zarządzania różnymi architekturami GPU oraz generowanie kodu na poziomie platformy, który będzie działał dobrze na kliencie.

Wnioski

Zespół Chrome cały czas pomaga w rozwijaniu standardów WebAssembly i WebGPU, aby ulepszać platformę internetową do obsługi zadań systemów uczących się. Inwestujemy w szybsze podstawowe elementy obliczeniowe, lepszą współpracę w zakresie standardów internetowych i zapewniamy, że zarówno duże, jak i małe modele mogą działać wydajnie na różnych urządzeniach.

Naszym celem jest maksymalizacja możliwości platformy przy jednoczesnym zachowaniu tego, co najlepsze w internecie – zasięgu, łatwości obsługi i przenośności. Nie robimy tego sami. Współpracujemy z innymi dostawcami przeglądarek w firmie W3C oraz wieloma partnerami programistycznymi.

Mamy nadzieję, że podczas pracy z WebAssembly i WebGPU będziesz pamiętać o tych kwestiach:

• Funkcja wnioskowania z wykorzystaniem AI jest teraz dostępna w internecie na różnych urządzeniach. Wykorzystuje to zalety działania na urządzeniach klienckich, w tym niższe koszty serwerów, niewielkie opóźnienia i większą prywatność.
• Choć wiele omówionych funkcji ma istotne znaczenie przede wszystkim dla autorów platformy, aplikacje mogą z nich korzystać bez większych nakładów pracy.
• Standardy internetowe są bardzo elastyczne i stale ewoluują, dlatego zawsze jesteśmy otwarci na opinie. Udostępnij swój kod dla WebAssembly i WebGPU.

Podziękowania

Chcielibyśmy podziękować zespołowi ds. grafiki internetowej firmy Intel, który w dużym stopniu przyczynił się do powstania funkcji WebGPU f16 oraz funkcji produktów z liczbą liczb całkowitych. Chcielibyśmy podziękować pozostałym członkom grup roboczych WebAssembly i WebGPU w W3C, w tym innym dostawcom przeglądarek.

Dziękujemy zespołom ds. AI i ML, zarówno w Google, jak i w społeczności open source, za bycie niezwykłymi partnerami. Oczywiście możemy też poprosić o pomoc wszystkich członków naszego zespołu, dzięki którym to wszystko jest możliwe.