더 빠른 웹 AI를 위한 WebAssembly 및 WebGPU 개선사항, 2부

이 문서는 빠른 웹 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팀은 대규모 언어 모델 추론을 위한 실험적 지원을 발표했습니다.

다음 애니메이션은 온디바이스 Chrome에서 실시간으로 실행되는 Google의 오픈소스 대규모 언어 모델 (LLM)인 Gemma를 보여줍니다.

Meta의 Segment Anything 모델에 대한 다음 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에서 실행되며 대용량 데이터에서 병렬 배열 작업을 실행하는 데 도움이 됩니다.

작년에 WebGPU의 수많은 개선 사항 중에서 Google은 웹에서 ML 및 AI 성능을 개선하기 위한 더 많은 기능을 지속적으로 추가했습니다. 최근에 우리는 16비트 부동 소수점과 패킹된 정수 닷 프로덕트라는 두 가지 새로운 기능을 출시했습니다.

16비트 부동 소수점

ML 워크로드에는 정밀도가 필요하지 않습니다. shader-f16는 WebGPU 셰이딩 언어에서 f16 유형을 사용할 수 있는 기능입니다. 이 부동 소수점 유형은 일반적인 32비트 대신 16비트를 차지합니다. f16은 범위가 더 좁고 정밀도가 낮지만 많은 ML 모델에서 이것만으로도 충분합니다.

이 기능은 몇 가지 측면에서 효율성을 높입니다.

  • 메모리 감소: f16 요소가 있는 텐서는 공간의 절반을 차지하여 메모리 사용량을 절반으로 줄입니다. GPU 계산은 메모리 대역폭에서 병목 현상이 발생하는 경우가 많으므로 메모리의 절반은 셰이더가 2배 더 빠르게 실행됨을 의미할 수 있습니다. 기술적으로는 메모리 대역폭을 절약하기 위해 f16이 필요하지 않습니다. 정밀도가 낮은 형식으로 데이터를 저장한 다음 계산을 위해 셰이더에서 전체 f32로 확장할 수 있습니다. 하지만 GPU는 데이터를 압축하고 압축해제하는 데 추가 컴퓨팅 성능을 소비합니다.

  • 데이터 변환 감소: f16은 데이터 변환을 최소화하여 컴퓨팅 사용량을 줄입니다. 정밀도가 낮은 데이터는 변환 없이 저장한 다음 직접 사용할 수 있습니다.

  • 병렬 처리 향상: 최신 GPU는 GPU 실행 단위에 더 많은 값을 동시에 적용할 수 있으므로 더 많은 병렬 계산을 수행할 수 있습니다. 예를 들어 초당 최대 5조 개의 f32 부동 소수점 연산을 지원하는 GPU는 초당 10조 개의 f16 부동 소수점 연산을 지원할 수 있습니다.

텍스트 임베딩용 WebGPU 벤치마크 스크린샷
shader-f16를 사용하는 Hugging Face의 텍스트 임베딩용 WebGPU 벤치마크 벤치마크는 Apple M1 Max 노트북에서 f32보다 3배 빠르게 벤치마크를 실행합니다.

WebLLM은 여러 대규모 언어 모델을 실행할 수 있는 프로젝트입니다. 오픈소스 머신러닝 컴파일러 프레임워크인 Apache TVM을 사용합니다.

저는 WebLLM에 Llama 3 80억 매개변수 모델을 사용하여 파리 여행을 계획해 달라고 요청했습니다. 결과에 따르면 모델의 자동 입력 단계에서 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을 명시적으로 사용 설정해야 합니다. 그런 다음 다른 부동 소수점 데이터 유형과 마찬가지로 셰이더 내에서 자유롭게 사용할 수 있습니다.

// 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는 거의 없습니다. 여기에서 패킹 정수 도트 제품이 사용됩니다. Chrome 123에서는 DP4a가 출시되었습니다.

최신 GPU에는 2개의 32비트 정수를 취해 각각을 4개의 연속으로 패킹된 8비트 정수로 해석하고 구성요소 간의 내적을 계산하라는 특별한 지침이 있습니다.

행렬 곱셈 커널은 수많은 점으로 구성되기 때문에 이는 AI와 머신러닝에 특히 유용합니다.

예를 들어 4 x 8 행렬에 8 x 1 벡터를 곱해 보겠습니다. 이를 계산하는 데는 4개의 내적을 취하여 출력 벡터의 각 값을 계산합니다. A, B, C, D입니다.

행렬 벡터 곱셈 예 다이어그램

이러한 각 출력을 계산하는 프로세스는 동일합니다. 그 중 하나를 계산하는 데 관련된 단계를 살펴보겠습니다. 계산하기 전에 먼저 8비트 정수 데이터를 f16과 같이 산술을 할 수 있는 유형으로 변환해야 합니다. 그런 다음 요소별 곱셈을 실행하고 마지막으로 모든 곱을 더합니다. 전체 행렬 벡터 곱셈의 경우, 데이터 압축해제를 위한 부동 소수점 변환 40개, 부동 소수점 수 곱셈 32개, 부동 소수점 덧셈 28개를 수행합니다.

연산이 많은 큰 행렬의 경우, 패킹된 정수 내적을 사용하면 작업량을 줄일 수 있습니다.

결과 벡터의 각 출력에 대해 기본 제공되는 WebGPU Shading Language dot4U8Packed를 사용하여 두 개의 패킹 내적 연산을 실행한 다음 결과를 함께 추가합니다. 전체적으로 행렬 벡터 곱셈에 대해 데이터 변환을 수행하지 않습니다. 저희는 패킹된 점 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.
}

다음 코드 스니펫 차이점 (차이)은 WebGPU 셰이더에서 패키징된 정수 제품을 사용하는 데 필요한 변경사항을 강조 표시합니다.

이전 - 부분적인 내적을 'sum' 변수에 누적하는 WebGPU 셰이더입니다. 루프 끝에서 `sum` 은 벡터와 입력 행렬의 한 행 사이의 전체 내적을 보유합니다.

// 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개의 부동 소수점 값을 로드하는 대신 단일 32비트 정수를 로드한다는 것입니다. 이 32비트 정수에는 8비트 정수 값 4개로 된 데이터가 저장됩니다. 그런 다음 dot4U8Packed를 호출하여 두 값의 내적을 계산합니다.

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

16비트 부동 소수점과 패킹된 정수 닷 제품은 모두 AI와 ML을 가속화하는 Chrome의 기본 기능입니다. 하드웨어가 16비트 부동 소수점을 지원하고 Chrome은 모든 기기에서 패킹된 정수 점 프로덕트를 구현하면 사용할 수 있습니다.

지금 바로 Chrome 안정화 버전에서 이러한 기능을 사용하여 더 나은 성능을 달성할 수 있습니다.

제안된 기능

앞으로 하위 그룹과 협력 행렬 곱셈이라는 두 가지 기능을 추가로 조사하고 있습니다.

하위 그룹 기능을 사용하면 SIMD 수준의 동시 로드가 통신하거나 16개 이상의 숫자에 대한 합계와 같은 집합적인 수학 연산을 수행할 수 있습니다. 이를 통해 효율적인 교차 스레드 데이터 공유가 가능합니다. 하위 그룹은 이름과 형식이 약간 다른 최신 GPU API에서 지원됩니다.

공통 세트를 WebGPU 표준화 그룹에 적용한 제안으로 정리했습니다. 또한 실험용 플래그를 사용하여 Chrome의 하위 그룹 프로토타입을 제작하고 초기 결과를 논의했습니다. 주요 문제는 이식 가능한 동작을 어떻게 보장하는가입니다.

협력 행렬 곱셈은 최근 GPU에 추가된 기능입니다. 큰 행렬 곱셈은 여러 개의 작은 행렬 곱셈으로 나눌 수 있습니다. 협력 행렬 곱셈은 단일 논리 단계에서 크기가 작은 이러한 작은 블록에 대해 곱셈을 수행합니다. 이 단계 내에서 스레드 그룹이 효율적으로 협력하여 결과를 계산합니다.

Google에서는 기본 GPU API에 대한 지원을 설문조사했으며 WebGPU 표준화 그룹에 제안서를 제공할 계획입니다. 하위 그룹과 마찬가지로 대부분의 논의는 이동성에 초점을 맞출 것으로 예상됩니다.

하위 그룹 작업의 성능을 평가하기 위해 실제 애플리케이션에서 하위 그룹에 대한 실험용 지원을 MediaPipe에 통합하고 하위 그룹 작업을 위한 Chrome 프로토타입으로 테스트했습니다.

대규모 언어 모델의 자동 입력 단계의 GPU 커널에서 하위 그룹을 사용했으므로 미리 완성 단계의 속도 향상만 보고합니다. Intel GPU에서는 하위 그룹의 성능이 기준보다 2.5배 더 빠릅니다. 그러나 이러한 개선사항은 여러 GPU에서 일관되지 않습니다.

<ph type="x-smartling-placeholder">
</ph> MediaPipe LLM 추론의 하위 그룹 speedup 스크린샷 <ph type="x-smartling-placeholder">
</ph> 차트 2. 하위 그룹을 사용하면 미리 완성이 Intel Tiger Lake GT2 GPU에서 2.5배 더 빠르게 실행되며 Chrome 및 Mediapipe에서 실험용 지원이 제공됩니다.

다음 차트는 하위 그룹을 적용하여 여러 소비자 GPU에서 행렬 곱하기 microbenchmark를 최적화한 결과를 보여줍니다. 행렬 곱셈은 대규모 언어 모델에서 비용이 더 많이 드는 작업 중 하나입니다. 데이터에 따르면 많은 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은 메모리 액세스 패턴을 최적화하여 하드웨어에 더 최적화됩니다.

메모리 재구성이 상당한 영향을 미친다는 것을 알 수 있습니다. 하위 그룹보다 더 큰 영향을 미칠 수 있습니다. GPU 6에서 재구성은 속도 향상을 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 표준을 지속적으로 발전시키고 있습니다. Google은 더 빠른 컴퓨팅 프리미티브, 웹 표준 간의 상호 운용 개선, 크고 작은 모델을 여러 기기에서 효율적으로 실행할 수 있도록 하는 데 투자하고 있습니다.

당사의 목표는 도달범위, 사용성, 이동성 등 웹의 장점을 유지하면서 플랫폼의 기능을 극대화하는 것입니다. Google만의 노력이 아닙니다. Google은 W3C의 다른 브라우저 공급업체 및 여러 개발 파트너와 협력하고 있습니다.

WebAssembly 및 WebGPU를 사용할 때 다음 사항을 기억하시기 바랍니다.

  • 이제 웹과 여러 기기에서 AI 추론을 사용할 수 있습니다. 따라서 클라이언트 기기에서 실행할 수 있어 서버 비용 절감, 짧은 지연 시간, 개인 정보 보호 향상과 같은 이점이 있습니다.
  • 논의된 많은 기능들이 주로 프레임워크 작성자와 관련이 있지만, 여러분의 애플리케이션은 많은 오버헤드 없이 이점을 누릴 수 있습니다.
  • 웹 표준은 유동적이고 변화하며 Google은 항상 의견을 기다리고 있습니다. WebAssembly 및 WebGPU를 공유합니다.

감사의 말씀

WebGPU f16과 패키징된 integer 닷 프로덕트 기능을 구현하는 데 중요한 역할을 해주신 Intel 웹 그래픽 팀에 감사의 말씀을 전하고 싶습니다. 다른 브라우저 공급업체를 포함하여 W3C의 WebAssembly 및 WebGPU 작업 그룹에 속한 다른 회원들에게도 감사의 말을 전하고 싶습니다.

훌륭한 파트너가 되어 주신 Google과 오픈소스 커뮤니티의 AI팀 및 ML팀 여러분께 감사의 말씀을 전합니다. 물론 이 모든 것을 가능하게 해 주는 모든 팀원도 있습니다.